Merge pull request #60 from thiremani/feat-array-comparison-mask

feat(compiler): element-wise array comparison masks + per-slot condition propagation

Author: thiremani

README.md

@@ -143,11 +143,11 @@ b = Square(2.2)          # float specialization
 c = Square(arr)          # squares each array element
 a, b, c
 
-# Range and filter (non-accumulated)
+# Range, mask, and filter (non-accumulated)
 d = Square(1:3)          # range: final iteration result
 e = Square(arr[1:3])     # array-range: final iteration result
-f = Square(arr > 3)      # filtered array
-g = Square(arr[1:3] > 3) # filtered array-range
+f = Square(arr > 3)      # element-wise mask: each element kept where > 3, else 0
+g = Square(arr[1:3] > 3) # array-range filter: final iteration result
 d, e, f, g
 
 # Accumulation forms
@@ -162,11 +162,11 @@ Output:
 
 ```text
 25 4.84 [1 4 9 25]
-4 9 [25] 0
+4 9 [0 0 0 25] 0
 [1 4] [4 9] [0 4] [0 9 25]
 ```
 
-No generics syntax, no type parameters. Write the template once — call it with a scalar, an array, or a filtered view. `arr > 3` filters the array to elements greater than 3 and passes that subset through.
+No generics syntax, no type parameters. Write the template once — call it with a scalar, an array, or a masked array. `arr > 3` masks the array element-wise — each element kept where it's greater than 3, else 0 (same length) — and passes that through.
 
 Generated code is equivalent to handwritten specialized code. There is no runtime overhead.
 
@@ -195,7 +195,7 @@ x = [1 2 3 4 5]
 y = [1.1 2.2 3.3]
 ```
 
-Arrays are safe by construction — out-of-bounds access is not possible. Comparisons like `arr > 2` produce filtered views that work anywhere an array does.
+Arrays are safe by construction — out-of-bounds access is not possible. Comparisons like `arr > 2` produce element-wise masks (each element kept where it holds, else 0) that work anywhere an array does.
 
 ---
 

compiler/array.go

@@ -607,6 +607,16 @@ func (c *Compiler) forEachArrayPair(
 	})
 }
 
+// arraySym wraps an array vector and element type as a Symbol, bitcast to the
+// opaque i8* array handle used for array values.
+func (c *Compiler) arraySym(array llvm.Value, elemType Type) *Symbol {
+	i8p := llvm.PointerType(c.Context.Int8Type(), 0)
+	return &Symbol{
+		Type: Array{Headers: nil, ColTypes: []Type{elemType}, Length: 0},
+		Val:  c.builder.CreateBitCast(array, i8p, "arr_i8p"),
+	}
+}
+
 func (c *Compiler) compileArrayArrayInfix(op string, left *Symbol, right *Symbol, resElem Type) *Symbol {
 	leftArrType := left.Type.(Array)
 	rightArrType := right.Type.(Array)
@@ -639,10 +649,7 @@ func (c *Compiler) compileArrayArrayInfix(op string, left *Symbol, right *Symbol
 	})
 
 	// Return result array
-	i8p := llvm.PointerType(c.Context.Int8Type(), 0)
-	resSym := &Symbol{Type: Array{Headers: nil, ColTypes: []Type{resElem}, Length: 0}}
-	resSym.Val = c.builder.CreateBitCast(resVec, i8p, "arr_i8p")
-	return resSym
+	return c.arraySym(resVec, resElem)
 }
 
 func (c *Compiler) compileArrayConcat(left *Symbol, right *Symbol, leftElem Type, rightElem Type, resElem Type) *Symbol {
@@ -666,10 +673,7 @@ func (c *Compiler) compileArrayConcat(left *Symbol, right *Symbol, leftElem Type
 	c.CopyArrayInto(resVec, right, rightElem, resElem, leftLen, true)
 
 	// Return concatenated array
-	i8p := llvm.PointerType(c.Context.Int8Type(), 0)
-	resSym := &Symbol{Type: Array{Headers: nil, ColTypes: []Type{resElem}, Length: 0}}
-	resSym.Val = c.builder.CreateBitCast(resVec, i8p, "arr_i8p")
-	return resSym
+	return c.arraySym(resVec, resElem)
 }
 
 func (c *Compiler) compileArrayScalarInfix(op string, arr *Symbol, scalar *Symbol, resElem Type, arrayOnLeft bool) *Symbol {
@@ -695,83 +699,74 @@ func (c *Compiler) compileArrayScalarInfix(op string, arr *Symbol, scalar *Symbo
 		c.ArraySetOwnForType(resElem, resVec, iter, resultVal)
 	})
 
-	i8p := llvm.PointerType(c.Context.Int8Type(), 0)
-	resSym := &Symbol{Type: Array{Headers: nil, ColTypes: []Type{resElem}, Length: 0}}
-	resSym.Val = c.builder.CreateBitCast(resVec, i8p, "arr_i8p")
-	return resSym
+	return c.arraySym(resVec, resElem)
 }
 
-// compileArrayFilter dispatches value-position comparison filtering when at least
-// one operand is an array. It keeps LHS values where the comparison holds.
-// For scalar-op-array, the scalar LHS is repeated for each true element.
-func (c *Compiler) compileArrayFilter(op string, left *Symbol, right *Symbol, expected Type) *Symbol {
+// compileArrayMask dispatches value-position comparison masking when at least one
+// operand is an array. Each position yields the LHS element where the comparison
+// holds, otherwise 0 — preserving length and position (a 0 marks a failed compare,
+// never a dropped element). This is Pluto's scalar "comparison yields LHS-or-0"
+// applied per element. Array-array zips to min length; array-scalar uses the
+// array's length and broadcasts the scalar.
+func (c *Compiler) compileArrayMask(op string, left *Symbol, right *Symbol, expected Type) *Symbol {
 	l := c.derefIfPointer(left, "")
 	r := c.derefIfPointer(right, "")
-
-	resArr := expected.(Array)
-	acc := c.NewArrayAccumulator(resArr)
+	resElem := expected.(Array).ColTypes[0]
 
 	if l.Type.Kind() == ArrayKind && r.Type.Kind() == ArrayKind {
-		return c.compileArrayArrayFilter(op, l, r, acc)
+		return c.compileArrayArrayMask(op, l, r, resElem)
 	}
 	if l.Type.Kind() == ArrayKind {
-		return c.compileArrayScalarFilter(op, l, r, acc, true)
+		return c.compileArrayScalarMask(op, l, r, resElem, true)
 	}
 	if r.Type.Kind() == ArrayKind {
-		return c.compileArrayScalarFilter(op, r, l, acc, false)
+		return c.compileArrayScalarMask(op, r, l, resElem, false)
 	}
-	panic(fmt.Sprintf("compileArrayFilter expects at least one array operand, got %s and %s", l.Type, r.Type))
+	panic(fmt.Sprintf("compileArrayMask expects at least one array operand, got %s and %s", l.Type, r.Type))
 }
 
-// filterPush conditionally appends sym to acc based on cond.
-// Emits a branch: if cond is true, push sym; otherwise skip.
-func (c *Compiler) filterPush(acc *ArrayAccumulator, sym *Symbol, cond llvm.Value) {
-	copyBlock, nextBlock := c.createIfCont(cond, "filter_copy", "filter_next")
+// maskStore writes lhs into resVec[iter] when cond holds; on failure it leaves the
+// slot at its preallocated zero (0 for numbers, "" for strings), realizing
+// `cond ? lhs : 0` per cell. The set copies the (borrowed) element, so the source
+// array keeps ownership and the result owns an independent copy.
+func (c *Compiler) maskStore(resElem Type, resVec llvm.Value, iter llvm.Value, lhs *Symbol, cond llvm.Value) {
+	setBlock, contBlock := c.createIfCont(cond, "mask_set", "mask_cont")
 
-	c.builder.SetInsertPointAtEnd(copyBlock)
-	c.PushVal(acc, sym)
-	c.builder.CreateBr(nextBlock)
+	c.builder.SetInsertPointAtEnd(setBlock)
+	c.ArraySetForType(resElem, resVec, iter, lhs.Val)
+	c.builder.CreateBr(contBlock)
 
-	c.builder.SetInsertPointAtEnd(nextBlock)
+	c.builder.SetInsertPointAtEnd(contBlock)
 }
 
-func (c *Compiler) compileArrayArrayFilter(op string, left *Symbol, right *Symbol, acc *ArrayAccumulator) *Symbol {
-	c.forEachArrayPair(left, right, c.arrayPairMinLen(left, right), func(_ llvm.Value, leftSym *Symbol, rightSym *Symbol) {
-		cmpResult := defaultOps[opKey{
-			Operator:  op,
-			LeftType:  opType(leftSym.Type.Key()),
-			RightType: opType(rightSym.Type.Key()),
-		}](c, leftSym, rightSym, true)
-
-		c.filterPush(acc, leftSym, cmpResult.Val)
+func (c *Compiler) compileArrayArrayMask(op string, left *Symbol, right *Symbol, resElem Type) *Symbol {
+	loopLen := c.arrayPairMinLen(left, right)
+	resVec := c.CreateArrayForType(resElem, loopLen)
+	c.forEachArrayPair(left, right, loopLen, func(iter llvm.Value, leftSym *Symbol, rightSym *Symbol) {
+		lhs, cond := c.compareScalars(op, leftSym, rightSym)
+		c.maskStore(resElem, resVec, iter, lhs, cond)
 	})
 
-	return c.ArrayAccResult(acc)
+	return c.arraySym(resVec, resElem)
 }
 
-func (c *Compiler) compileArrayScalarFilter(op string, arr *Symbol, scalar *Symbol, acc *ArrayAccumulator, arrayOnLeft bool) *Symbol {
+func (c *Compiler) compileArrayScalarMask(op string, arr *Symbol, scalar *Symbol, resElem Type, arrayOnLeft bool) *Symbol {
 	arrElem := arr.Type.(Array).ColTypes[0]
-	c.forEachArrayPair(arr, scalar, c.ArrayLen(arr, arrElem), func(_ llvm.Value, elemSym *Symbol, scalarSym *Symbol) {
-		leftSym := elemSym
-		rightSym := scalarSym
-		pushSym := elemSym
+	loopLen := c.ArrayLen(arr, arrElem)
+	resVec := c.CreateArrayForType(resElem, loopLen)
+	c.forEachArrayPair(arr, scalar, loopLen, func(iter llvm.Value, elemSym *Symbol, scalarSym *Symbol) {
+		// Preserve operand order for non-commutative comparisons. The masked value
+		// is always the comparison's LHS (the array element, or the scalar when it
+		// is on the left), which compareScalars returns alongside the result.
+		leftSym, rightSym := elemSym, scalarSym
 		if !arrayOnLeft {
-			// Preserve original operand order for non-commutative comparisons,
-			// and keep scalar LHS values on true comparisons.
-			leftSym = scalarSym
-			rightSym = elemSym
-			pushSym = scalarSym
+			leftSym, rightSym = scalarSym, elemSym
 		}
-		cmpResult := defaultOps[opKey{
-			Operator:  op,
-			LeftType:  opType(leftSym.Type.Key()),
-			RightType: opType(rightSym.Type.Key()),
-		}](c, leftSym, rightSym, true)
-
-		c.filterPush(acc, pushSym, cmpResult.Val)
+		lhs, cond := c.compareScalars(op, leftSym, rightSym)
+		c.maskStore(resElem, resVec, iter, lhs, cond)
 	})
 
-	return c.ArrayAccResult(acc)
+	return c.arraySym(resVec, resElem)
 }
 
 func (c *Compiler) compileArrayUnaryPrefix(op string, arr *Symbol, result Array) *Symbol {
@@ -799,10 +794,7 @@ func (c *Compiler) compileArrayUnaryPrefix(op string, arr *Symbol, result Array)
 		c.ArraySetOwnForType(resElem, resVec, idx, resultVal)
 	})
 
-	i8p := llvm.PointerType(c.Context.Int8Type(), 0)
-	resSym := &Symbol{Type: result}
-	resSym.Val = c.builder.CreateBitCast(resVec, i8p, "arr_i8p")
-	return resSym
+	return c.arraySym(resVec, resElem)
 }
 
 // Array string conversion function
@@ -916,8 +908,9 @@ func (c *Compiler) compileArrayRangeBasic(expr *ast.ArrayRangeExpression) []*Sym
 		}}
 	}
 	// Scalar element access does not retain the source array pointer.
-	// Release temporary array sources on all scalar return paths.
-	defer c.freeTemporary(expr.Array, []*Symbol{arraySym})
+	// Release temporary array sources on all scalar return paths; a consumed
+	// frame mask ((arr > 1)[i]) is recorded so the sweep skips it.
+	defer c.freeConsumedTemporary(expr.Array, []*Symbol{arraySym})
 
 	// Scalar index - element access
 	arrElemType := arrType.ColTypes[0]