topdown: exit early for functions

Turns out the work for this is shared with the other document kinds in the indexer result. All we have to do here is to actually exit early if possible. Signed-off-by: Stephan Renatus <[email protected]>
open-policy-agent · Oct 19, 2021 · d1eab2f · d1eab2f
1 parent c1a2b04
commit d1eab2f
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 6 deletions.
diff --git a/ast/index_test.go b/ast/index_test.go
@@ -963,11 +963,89 @@ r = data.a {
 	input.y = 1
 }`),
 			input: `{"x": 1, "y": 1}`,
+		},
+		{
+			note:       "function: no early exit: single rule",
+			expectedEE: false,
+			module: MustParseModule(`package test
+r(x) {
+	input.x = x
+}
+r = 3 {
+	input.y = 2
+}`),
+			input: `{"x": 1}`,
 			expectedRS: []string{
-				`r = data.a { input.x = 1 }`,
-				`r = data.a { input.y = 1 }`,
+				`r(x) { input.x = x }`,
 			},
 		},
+		{
+			note:       "function: no early exit: different constant value",
+			expectedEE: false,
+			module: MustParseModule(`package test
+r(x) {
+	input.x = x
+}
+r(y) = 2 {
+	input.x = 1
+	input.y = y
+}`),
+			input: `{"x": 1, "y": 2}`,
+			expectedRS: []string{
+				`r(x) { input.x = x }`,
+				`r(y) = 2 { input.x = 1; input.y = y }`,
+			},
+		},
+		{
+			note:       "function: same constant value",
+			expectedEE: true,
+			module: MustParseModule(`package test
+r(x) {
+	input.x = x
+}
+r(y) {
+	input.y = y
+}`),
+			input: `{"x": 1, "y": 1}`,
+		},
+		{
+			note:       "function: no early exit: one with with non-constant value",
+			expectedEE: false,
+			module: MustParseModule(`package test
+r(x) {
+	input.x = x
+}
+r(y) = x {
+	input.y = y
+	x = "foo"
+}`),
+			input: `{"x": 1, "y": 1}`,
+		},
+		{
+			note:       "function: same ref value (input)",
+			expectedEE: true,
+			module: MustParseModule(`package test
+r(x) = input.a {
+	input.x = x
+}
+r(y) = input.a {
+	input.y = y
+}`),
+			input: `{"x": 1, "y": 1}`,
+		},
+		{
+			note:       "function: same ref value (data)",
+			expectedEE: true,
+			module: MustParseModule(`package test
+r(x) = data.a {
+	input.x = x
+}
+r(y) = data.a {
+	input.y = y
+}`),
+			input: `{"x": 1, "y": 1}`,
+		},
+
 		// NOTE(sr): The remaining cases record the limitations of the current implementation:
 		// Any matching rules whose values contain non-constant values are not compared, and
 		// cancel early exit.

diff --git a/topdown/eval.go b/topdown/eval.go
@@ -1623,14 +1623,14 @@ func (e evalFunc) eval(iter unifyIterator) error {
 
 	var prev *ast.Term
 
-	for i := range e.ir.Rules {
-		next, err := e.evalOneRule(iter, e.ir.Rules[i], cacheKey, prev)
+	for _, rule := range e.ir.Rules {
+		next, err := e.evalOneRule(iter, rule, cacheKey, prev)
 		if err != nil {
 			return err
 		}
 		if next == nil {
-			for _, rule := range e.ir.Else[e.ir.Rules[i]] {
-				next, err = e.evalOneRule(iter, rule, cacheKey, prev)
+			for _, erule := range e.ir.Else[rule] {
+				next, err = e.evalOneRule(iter, erule, cacheKey, prev)
 				if err != nil {
 					return err
 				}
@@ -1641,6 +1641,11 @@ func (e evalFunc) eval(iter unifyIterator) error {
 		}
 		if next != nil {
 			prev = next
+
+			if e.ir.EarlyExit {
+				e.e.traceExitEarly(rule)
+				break
+			}
 		}
 	}