Skip to content

Commit

Permalink
topdown: add numbers.range_step built-in function (#6187)
Browse files Browse the repository at this point in the history
Fixes: #6186
Signed-off-by: Sebastian Spaink <sebastianspaink@gmail.com>
  • Loading branch information
sspaink authored Aug 30, 2023
1 parent 519eea7 commit c28eaf7
Show file tree
Hide file tree
Showing 5 changed files with 212 additions and 16 deletions.
18 changes: 18 additions & 0 deletions ast/builtins.go
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,7 @@ var DefaultBuiltins = [...]*Builtin{

// Numbers
NumbersRange,
NumbersRangeStep,
RandIntn,

// Encoding
Expand Down Expand Up @@ -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
*/
Expand Down
28 changes: 28 additions & 0 deletions builtin_metadata.json
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,7 @@
"minus",
"mul",
"numbers.range",
"numbers.range_step",
"plus",
"rand.intn",
"rem",
Expand Down Expand Up @@ -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": [
{
Expand Down
23 changes: 23 additions & 0 deletions capabilities.json
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand Down
87 changes: 87 additions & 0 deletions test/cases/testdata/numbersrangestep/test-numbersrangestep.yaml
Original file line number Diff line number Diff line change
@@ -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
72 changes: 56 additions & 16 deletions topdown/numbers.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
}

0 comments on commit c28eaf7

Please sign in to comment.