diff --git a/v1/test/cases/testdata/v0/numbersrange/test-numbersrange-0258.yaml b/v1/test/cases/testdata/v0/numbersrange/test-numbersrange-0258.yaml index b2fe4f23cb..1e8dd378ba 100644 --- a/v1/test/cases/testdata/v0/numbersrange/test-numbersrange-0258.yaml +++ b/v1/test/cases/testdata/v0/numbersrange/test-numbersrange-0258.yaml @@ -19,3 +19,20 @@ cases: - -1 - -2 - -3 + - note: numbersrange/descending (cheap optimization) + query: data.generated.p = x + modules: + - | + package generated + + p := __local0__ { + numbers.range(5, 2, __local1__) + __local0__ = __local1__ + } + data: {} + want_result: + - x: + - 5 + - 4 + - 3 + - 2 diff --git a/v1/test/cases/testdata/v0/numbersrangestep/test-numbersrangestep.yaml b/v1/test/cases/testdata/v0/numbersrangestep/test-numbersrangestep.yaml index a203acc3ce..dc05e621a5 100644 --- a/v1/test/cases/testdata/v0/numbersrangestep/test-numbersrangestep.yaml +++ b/v1/test/cases/testdata/v0/numbersrangestep/test-numbersrangestep.yaml @@ -34,6 +34,21 @@ cases: - -6 - -8 - -10 + - note: numbersrangestep/descending (cheap optimization) + query: data.test.p = x + modules: + - | + package test + + p = num { + num := numbers.range_step(10, 3, 2) + } + want_result: + - x: + - 10 + - 8 + - 6 + - 4 - note: numbersrangestep/negative query: data.test.p = x modules: diff --git a/v1/test/cases/testdata/v1/numbersrange/test-numbersrange-0258.yaml b/v1/test/cases/testdata/v1/numbersrange/test-numbersrange-0258.yaml index 65f480ff25..15f1d2d2c2 100644 --- a/v1/test/cases/testdata/v1/numbersrange/test-numbersrange-0258.yaml +++ b/v1/test/cases/testdata/v1/numbersrange/test-numbersrange-0258.yaml @@ -19,3 +19,20 @@ cases: - -1 - -2 - -3 + - note: numbersrange/descending (cheap optimization) + query: data.generated.p = x + modules: + - | + package generated + + p := __local0__ if { + numbers.range(5, 2, __local1__) + __local0__ = __local1__ + } + data: {} + want_result: + - x: + - 5 + - 4 + - 3 + - 2 diff --git a/v1/test/cases/testdata/v1/numbersrangestep/test-numbersrangestep.yaml b/v1/test/cases/testdata/v1/numbersrangestep/test-numbersrangestep.yaml index 92d922f08a..4851274d0a 100644 --- a/v1/test/cases/testdata/v1/numbersrangestep/test-numbersrangestep.yaml +++ b/v1/test/cases/testdata/v1/numbersrangestep/test-numbersrangestep.yaml @@ -34,6 +34,21 @@ cases: - -6 - -8 - -10 + - note: numbersrangestep/descending (cheap optimization) + query: data.test.p = x + modules: + - | + package test + + p = num if { + num := numbers.range_step(10, 3, 2) + } + want_result: + - x: + - 10 + - 8 + - 6 + - 4 - note: numbersrangestep/negative query: data.test.p = x modules: diff --git a/v1/topdown/numbers.go b/v1/topdown/numbers.go index 4f25879efe..05a4f96f0e 100644 --- a/v1/topdown/numbers.go +++ b/v1/topdown/numbers.go @@ -107,8 +107,14 @@ func generateCheapRange(operands []*ast.Term, iter func(*ast.Term) error) error terms := make([]*ast.Term, 0, y+1) - for i := x; i <= y; i += step { - terms = append(terms, ast.InternedIntNumberTerm(i)) + if x <= y { + for i := x; i <= y; i += step { + terms = append(terms, ast.InternedIntNumberTerm(i)) + } + } else { + for i := x; i >= y; i -= step { + terms = append(terms, ast.InternedIntNumberTerm(i)) + } } return iter(ast.ArrayTerm(terms...))