topdown: add numbers.range_step built-in function (#6187)

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

Author: sspaink

ast/builtins.go

@@ -136,6 +136,7 @@ var DefaultBuiltins = [...]*Builtin{
 
 	// Numbers
 	NumbersRange,
+	NumbersRangeStep,
 	RandIntn,
 
 	// Encoding
@@ -1347,6 +1348,23 @@ 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.
+	If the provided "step" is less then 1, an error will be thrown.
+	If "b" is not in the range of the provided "step", "b" won't be included in the result.
+	`,
+	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` in `step` increments"),
+	),
+}
+
 /**
  * Units
  */

builtin_metadata.json

@@ -110,6 +110,7 @@
       "minus",
       "mul",
       "numbers.range",
+      "numbers.range_step",
       "plus",
       "rand.intn",
       "rem",
@@ -11558,6 +11559,33 @@
     },
     "wasm": true
   },
+  "numbers.range_step": {
+    "args": [
+      {
+        "name": "a",
+        "type": "number"
+      },
+      {
+        "name": "b",
+        "type": "number"
+      },
+      {
+        "name": "step",
+        "type": "number"
+      }
+    ],
+    "available": [
+      "edge"
+    ],
+    "description": "Returns an array of numbers in the given (inclusive) range incremented by a positive step.\n\tIf \"a==b\", then \"range == [a]\"; if \"a \u003e b\", then \"range\" is in descending order.\n\tIf the provided \"step\" is less then 1, an error will be thrown.\n\tIf \"b\" is not in the range of the provided \"step\", \"b\" won't be included in the result.\n\t",
+    "introduced": "edge",
+    "result": {
+      "description": "the range between `a` and `b` in `step` increments",
+      "name": "range",
+      "type": "array[number]"
+    },
+    "wasm": false
+  },
   "object.filter": {
     "args": [
       {

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": {

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

@@ -0,0 +1,87 @@
+cases:
+- note: numbersrangestep/ascending
+  query: data.test.p = x
+  modules:
+  - |
+    package test
+
+    p = num {
+      num := numbers.range_step(0, 10, 2)
+    }
+  want_result:
+  - x:
+    - 0
+    - 2
+    - 4
+    - 6
+    - 8
+    - 10
+- note: numbersrangestep/descending
+  query: data.test.p = x
+  modules:
+  - |
+    package test
+
+    p = num {
+      num := numbers.range_step(0, -10, 2)
+    }
+  want_result:
+  - x:
+    - 0
+    - -2
+    - -4
+    - -6
+    - -8
+    - -10
+- note: numbersrangestep/negative
+  query: data.test.p = x
+  modules:
+  - |
+    package test
+
+    p = num {
+      num := numbers.range_step(0, 10, -2)
+    }
+  want_error: 'numbers.range_step: step must be a positive number above zero'
+  want_error_code: eval_builtin_error
+  strict_error: true
+- note: numbersrangestep/memoryexample
+  query: data.test.p = x
+  modules:
+  - |
+    package test
+
+    p = num {
+      num := numbers.range_step(1024, 4096, 1024)
+    }
+  want_result:
+  - x:
+    - 1024
+    - 2048
+    - 3072
+    - 4096
+- note: numbersrangestep/equal
+  query: data.test.p = x
+  modules:
+  - |
+    package test
+
+    p = num {
+      num := numbers.range_step(2, 2, 2)
+    }
+  want_result:
+  - x:
+    - 2
+- note: numbersrangestep/notinrange
+  query: data.test.p = x
+  modules:
+  - |
+    package test
+
+    p = num {
+      num := numbers.range_step(2, 5, 2)
+    }
+  want_result:
+  - x:
+    - 2
+    - 4

topdown/numbers.go

@@ -28,32 +28,71 @@ func builtinNumbersRange(bctx BuiltinContext, operands []*ast.Term, iter func(*a
 		return err
 	}
 
-	result := ast.NewArray()
+	ast, err := generateRange(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 := generateRange(bctx, x, y, step, "numbers.range_step")
+	if err != nil {
+		return err
+	}
+
+	return iter(ast)
+}
+
+func generateRange(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
 		}
+		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 +134,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)
 }