Fixing optimized numbers.range builtin reversed range bug (#7230)

Signed-off-by: Johan Fylling <[email protected]>
open-policy-agent · Dec 19, 2024 · a190ea3 · a190ea3
1 parent 9a7d920
commit a190ea3
Show file tree

Hide file tree

Showing 5 changed files with 72 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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...))