diff --git a/ast/builtins.go b/ast/builtins.go index 6c4cac5c28..0dfac22979 100644 --- a/ast/builtins.go +++ b/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 */ diff --git a/builtin_metadata.json b/builtin_metadata.json index 28c2359d17..71d89d6fbf 100644 --- a/builtin_metadata.json +++ b/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": [ { diff --git a/capabilities.json b/capabilities.json index 1b0b80322b..50b9a1969c 100644 --- a/capabilities.json +++ b/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": { diff --git a/test/cases/testdata/numbersrangestep/test-numbersrangestep.yaml b/test/cases/testdata/numbersrangestep/test-numbersrangestep.yaml new file mode 100644 index 0000000000..6388cc3d10 --- /dev/null +++ b/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 diff --git a/topdown/numbers.go b/topdown/numbers.go index b0953b0fdb..27f3156b8a 100644 --- a/topdown/numbers.go +++ b/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) }