feat: add numbers.range_step built-in function

Signed-off-by: Sebastian Spaink <[email protected]>

ast/builtins.go

@@ -136,6 +136,7 @@ var DefaultBuiltins = [...]*Builtin{
 
 	// Numbers
 	NumbersRange,
+	NumbersRangeStep,
 	RandIntn,
 
 	// Encoding
@@ -1347,6 +1348,19 @@ var NumbersRange = &Builtin{
 	),
 }
 
+var NumbersRangeStep = &Builtin{
+	Name:        "numbers.range_step",
+	Description: "Returns an array of numbers in the given (inclusive) range incremented by a positive step. If `a==b`, then `range == [a]`; if `a > b`, then `range` is in descending order.",
+	Decl: types.NewFunction(
+		types.Args(
+			types.Named("a", types.N),
+			types.Named("b", types.N),
+			types.Named("step", types.N),
+		),
+		types.Named("range", types.NewArray(nil, types.N)).Description("the range between `a` and `b` incremented by a positive step."),
+	),
+}
+
 /**
  * Units
  */

capabilities.json

@@ -2895,6 +2895,29 @@
         "type": "function"
       }
     },
+    {
+      "name": "numbers.range_step",
+      "decl": {
+        "args": [
+          {
+            "type": "number"
+          },
+          {
+            "type": "number"
+          },
+          {
+            "type": "number"
+          }
+        ],
+        "result": {
+          "dynamic": {
+            "type": "number"
+          },
+          "type": "array"
+        },
+        "type": "function"
+      }
+    },
     {
       "name": "object.filter",
       "decl": {
@@ -4662,4 +4685,4 @@
   "features": [
     "rule_head_ref_string_prefixes"
   ]
-}
+}

test/cases/testdata/numbersrangestep/test-numbersrangestep-1.yaml

@@ -0,0 +1,20 @@
+cases:
+- data: {}
+  modules:
+  - |
+    package generated
+
+    p = __local0__ {
+      numbers.range_step(0, 10, 2, __local1__)
+      __local0__ = __local1__
+    }
+  note: numbersrangestep/ascending
+  query: data.generated.p = x
+  want_result:
+  - x:
+    - 0
+    - 2
+    - 4
+    - 6
+    - 8
+    - 10

test/cases/testdata/numbersrangestep/test-numbersrangestep-2.yaml

@@ -0,0 +1,20 @@
+cases:
+- data: {}
+  modules:
+  - |
+    package generated
+
+    p = __local0__ {
+      numbers.range_step(0, -10, 2, __local1__)
+      __local0__ = __local1__
+    }
+  note: numbersrangestep/descending
+  query: data.generated.p = x
+  want_result:
+  - x:
+    - 0
+    - -2
+    - -4
+    - -6
+    - -8
+    - -10

test/cases/testdata/numbersrangestep/test-numbersrangestep-3.yaml

@@ -0,0 +1,15 @@
+cases:
+- data: {}
+  modules:
+  - |
+    package generated
+
+    p = __local0__ {
+      numbers.range_step(0, 10, -2, __local1__)
+      __local0__ = __local1__
+    }
+  note: numbersrangestep/negative
+  query: data.generated.p = x
+  want_error: 'numbers.range_step: step must be a positive number above zero'
+  want_error_code: eval_builtin_error
+  strict_error: true

test/cases/testdata/numbersrangestep/test-numbersrangestep-4.yaml

@@ -0,0 +1,18 @@
+cases:
+- data: {}
+  modules:
+  - |
+    package generated
+
+    p = __local0__ {
+      numbers.range_step(1024, 4096, 1024, __local1__)
+      __local0__ = __local1__
+    }
+  note: numbersrangestep/memoryexample
+  query: data.generated.p = x
+  want_result:
+  - x:
+    - 1024
+    - 2048
+    - 3072
+    - 4096

topdown/numbers.go

@@ -28,32 +28,72 @@ func builtinNumbersRange(bctx BuiltinContext, operands []*ast.Term, iter func(*a
 		return err
 	}
 
-	result := ast.NewArray()
+	ast, err := builtinRange(bctx, x, y, one, "numbers.range")
+	if err != nil {
+		return err
+	}
+
+	return iter(ast)
+}
+
+func builtinNumbersRangeStep(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
+
+	x, err := builtins.BigIntOperand(operands[0].Value, 1)
+	if err != nil {
+		return err
+	}
+
+	y, err := builtins.BigIntOperand(operands[1].Value, 2)
+	if err != nil {
+		return err
+	}
+
+	step, err := builtins.BigIntOperand(operands[2].Value, 3)
+	if err != nil {
+		return err
+	}
+
+	if step.Cmp(big.NewInt(0)) <= 0 {
+		return fmt.Errorf("numbers.range_step: step must be a positive number above zero")
+	}
+
+	ast, err := builtinRange(bctx, x, y, step, "numbers.range_step")
+	if err != nil {
+		return err
+	}
+
+	return iter(ast)
+}
+
+func builtinRange(bctx BuiltinContext, x *big.Int, y *big.Int, step *big.Int, funcName string) (*ast.Term, error) {
+
 	cmp := x.Cmp(y)
+
+	comp := func(i *big.Int, y *big.Int) bool { return i.Cmp(y) <= 0 }
+	iter := func(i *big.Int) *big.Int { return i.Add(i, step) }
+
+	if cmp > 0 {
+		comp = func(i *big.Int, y *big.Int) bool { return i.Cmp(y) >= 0 }
+		iter = func(i *big.Int) *big.Int { return i.Sub(i, step) }
+	}
+
+	result := ast.NewArray()
 	haltErr := Halt{
 		Err: &Error{
 			Code:    CancelErr,
-			Message: "numbers.range: timed out before generating all numbers in range",
+			Message: fmt.Sprintf("%s: timed out before generating all numbers in range", funcName),
 		},
 	}
 
-	if cmp <= 0 {
-		for i := new(big.Int).Set(x); i.Cmp(y) <= 0; i = i.Add(i, one) {
-			if bctx.Cancel != nil && bctx.Cancel.Cancelled() {
-				return haltErr
-			}
-			result = result.Append(ast.NewTerm(builtins.IntToNumber(i)))
-		}
-	} else {
-		for i := new(big.Int).Set(x); i.Cmp(y) >= 0; i = i.Sub(i, one) {
-			if bctx.Cancel != nil && bctx.Cancel.Cancelled() {
-				return haltErr
-			}
-			result = result.Append(ast.NewTerm(builtins.IntToNumber(i)))
+	for i := new(big.Int).Set(x); comp(i, y); i = iter(i) {
+		if bctx.Cancel != nil && bctx.Cancel.Cancelled() {
+			return nil, haltErr
 		}
+		fmt.Println(i)
+		result = result.Append(ast.NewTerm(builtins.IntToNumber(i)))
 	}
 
-	return iter(ast.NewTerm(result))
+	return ast.NewTerm(result), nil
 }
 
 func builtinRandIntn(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
@@ -95,5 +135,6 @@ func builtinRandIntn(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.T
 
 func init() {
 	RegisterBuiltinFunc(ast.NumbersRange.Name, builtinNumbersRange)
+	RegisterBuiltinFunc(ast.NumbersRangeStep.Name, builtinNumbersRangeStep)
 	RegisterBuiltinFunc(ast.RandIntn.Name, builtinRandIntn)
 }