topdown: Fix namespacing to use caller bindings

This change fixes an issue in partial eval where variables in negated
expressions were getting plugged with the wrong set of caller
bindings. The caller bindings identify the context in which plugging
is being performed so that variables can be namespaced and conflicts
can be avoided.

For example, given the query `p[x]` and the rule `p[y] { y = input;
not y = 1; ... }` the first expression in the rule body would
eventually be plugged in the context of the calling query whereas the
second expression would be plugged using a special set of bindings
that ensured namespacing was performed. The problem was that the
partial eval negation implementation was incorrectly choosing the
special binding set.

This change updates partial eval to use the bindings from the top of
the evaluator stack in all cases where variables are plugged with
namespacing enabled. This lets us remove the special "sentinel"
bindings that were used to ensure namespacing was performed and
ensures that variables that appear in the original query are not
namespaced (which is the desired behaviour.)

Fixes #1814

Signed-off-by: Torin Sandall <torinsandall@gmail.com>

topdown/bindings.go

@@ -6,7 +6,6 @@ package topdown
 
 import (
 	"fmt"
-	"math"
 	"strings"
 
 	"github.com/open-policy-agent/opa/ast"
@@ -31,10 +30,6 @@ func (u *undo) Undo() {
 	u.next.Undo()
 }
 
-// sentinel is a binding list that will never be used. The binding list
-// identifier increments from zero.
-var sentinel = newBindings(math.MaxUint64, nil)
-
 type bindings struct {
 	id     uint64
 	values bindingsArrayHashmap

topdown/eval.go

@@ -33,6 +33,7 @@ type eval struct {
 	queryID         uint64
 	queryIDFact     *queryIDFactory
 	parent          *eval
+	caller          *eval
 	cancel          Cancel
 	query           ast.Body
 	index           int
@@ -386,27 +387,13 @@ func (e *eval) evalNotPartial(iter evalIterator) error {
 	negation := expr.Complement().NoWith()
 	child := e.closure(ast.NewBody(negation))
 
-	var caller *bindings
-
-	if e.parent == nil {
-		// The top-level query is being evaluated. Do not namespace variables from this
-		// query.
-		caller = e.bindings
-	} else {
-		// The top-level query is NOT being evaluated. All variables in the queries
-		// that are emitted by partial evaluation of this query MUST be namespaced. A
-		// sentinel is used because the plug operations do not namespace if the caller
-		// is nil.
-		caller = sentinel
-	}
-
 	// Unknowns is the set of variables that are marked as unknown. The variables
 	// are namespaced with the query ID that they originate in. This ensures that
 	// variables across two or more queries are identified uniquely.
 	//
 	// NOTE(tsandall): this is greedy in the sense that we only need variable
 	// dependencies of the negation.
-	unknowns := e.saveSet.Vars(caller)
+	unknowns := e.saveSet.Vars(e.caller.bindings)
 
 	// Run partial evaluation, plugging the result and applying copy propagation to
 	// each result. Since the result may require support, push a new query onto the
@@ -417,7 +404,7 @@ func (e *eval) evalNotPartial(iter evalIterator) error {
 
 	child.eval(func(*eval) error {
 		query := e.saveStack.Peek()
-		plugged := query.Plug(caller)
+		plugged := query.Plug(e.caller.bindings)
 		result := applyCopyPropagation(p, e.instr, plugged)
 		savedQueries = append(savedQueries, result)
 		return nil
@@ -834,7 +821,7 @@ func (e *eval) biunifyComprehensionPartial(a, b *ast.Term, b1, b2 *bindings, swa
 	// needed.) Eventually we may want to make the logic a bit smarter.
 	var extras []*ast.Expr
 
-	err := b1.Iter(sentinel, func(k, v *ast.Term) error {
+	err := b1.Iter(e.caller.bindings, func(k, v *ast.Term) error {
 		extras = append(extras, ast.Equality.Expr(k, v))
 		return nil
 	})
@@ -862,7 +849,7 @@ func (e *eval) biunifyComprehensionPartial(a, b *ast.Term, b1, b2 *bindings, swa
 		body.Append(e)
 	}
 
-	b1.Namespace(a, sentinel)
+	b1.Namespace(a, e.caller.bindings)
 
 	// The other term might need to be plugged so include the bindings. The
 	// bindings for the comprehension term are saved (for compatibility) but
@@ -1740,16 +1727,16 @@ func (e evalVirtualPartial) partialEvalSupportRule(iter unifyIterator, rule *ast
 		child.traceExit(rule)
 
 		current := e.e.saveStack.PopQuery()
-		plugged := current.Plug(child.bindings)
+		plugged := current.Plug(e.e.caller.bindings)
 
 		var key, value *ast.Term
 
 		if rule.Head.Key != nil {
-			key = child.bindings.PlugNamespaced(rule.Head.Key, child.bindings)
+			key = child.bindings.PlugNamespaced(rule.Head.Key, e.e.caller.bindings)
 		}
 
 		if rule.Head.Value != nil {
-			value = child.bindings.PlugNamespaced(rule.Head.Value, child.bindings)
+			value = child.bindings.PlugNamespaced(rule.Head.Value, e.e.caller.bindings)
 		}
 
 		head := ast.NewHead(rule.Head.Name, key, value)
@@ -1985,9 +1972,9 @@ func (e evalVirtualComplete) partialEvalSupportRule(iter unifyIterator, rule *as
 		child.traceExit(rule)
 
 		current := e.e.saveStack.PopQuery()
-		plugged := current.Plug(child.bindings)
+		plugged := current.Plug(e.e.caller.bindings)
 
-		head := ast.NewHead(rule.Head.Name, nil, child.bindings.PlugNamespaced(rule.Head.Value, child.bindings))
+		head := ast.NewHead(rule.Head.Name, nil, child.bindings.PlugNamespaced(rule.Head.Value, e.e.caller.bindings))
 		p := copypropagation.New(head.Vars()).WithEnsureNonEmptyBody(true)
 
 		e.e.saveSupport.Insert(path, &ast.Rule{

topdown/query.go

@@ -181,6 +181,7 @@ func (q *Query) PartialRun(ctx context.Context) (partials []ast.Body, support []
 		genvarprefix:    q.genvarprefix,
 		runtime:         q.runtime,
 	}
+	e.caller = e
 	q.startTimer(metrics.RegoPartialEval)
 	defer q.stopTimer(metrics.RegoPartialEval)
 
@@ -265,6 +266,7 @@ func (q *Query) Iter(ctx context.Context, iter func(QueryResult) error) error {
 		genvarprefix: q.genvarprefix,
 		runtime:      q.runtime,
 	}
+	e.caller = e
 	q.startTimer(metrics.RegoQueryEval)
 	err := e.Run(func(e *eval) error {
 		qr := QueryResult{}

topdown/topdown_partial_test.go

@@ -472,12 +472,12 @@ func TestTopDownPartialEval(t *testing.T) {
 		{
 			note:        "comprehensions: save",
 			query:       `x = [true | true]; y = {true | true}; z = {a: true | a = "foo"}`,
-			wantQueries: []string{`x = [true | true]; y = {true | true}; z = {a0: true | a0 = "foo"}`},
+			wantQueries: []string{`x = [true | true]; y = {true | true}; z = {a: true | a = "foo"}`},
 		},
 		{
 			note:        "comprehensions: closure",
 			query:       `i = 1; xs = [x | x = data.foo[i]]`,
-			wantQueries: []string{`xs = [x0 | x0 = data.foo[i0]; i0 = 1]; i = 1`},
+			wantQueries: []string{`xs = [x | x = data.foo[1]; 1 = 1]; i = 1`},
 		},
 		{
 			note:  "save: sub path",
@@ -718,7 +718,7 @@ func TestTopDownPartialEval(t *testing.T) {
 			wantSupport: []string{
 				`package partial.test
 				p = 1 { input.x = 1 }
-				p = x { input.x = x }
+				p = x1 { input.x = x1 }
 				default p = 0
 				`,
 			},
@@ -1394,6 +1394,22 @@ func TestTopDownPartialEval(t *testing.T) {
 				`,
 			},
 		},
+		{
+			note:  "negation: inlining namespaced variables",
+			query: "data.test.p[x]",
+			modules: []string{
+				`package test
+
+				p[y] {
+					y = input
+					not y = 1
+				}
+				`,
+			},
+			wantQueries: []string{
+				`x = input; not x = 1; x`,
+			},
+		},
 		{
 			note:  "disable inlining: complete doc",
 			query: "data.test.p = true",
@@ -1489,6 +1505,50 @@ func TestTopDownPartialEval(t *testing.T) {
 			},
 			disableInlining: []string{`data.test.q`},
 		},
+		{
+			note:            "disable inlining: partial rule namespaced variables (negation)",
+			query:           "data.test.p[x]",
+			disableInlining: []string{"data.test.p"},
+			modules: []string{
+				`package test
+
+				p[y] {
+					y = input
+					not y = 1
+				}
+				`,
+			},
+			wantQueries: []string{
+				`data.partial.test.p[x]`,
+			},
+			wantSupport: []string{
+				`package partial.test
+
+				p[y1] { y1 = input; not y1 = 1 }`,
+			},
+		},
+		{
+			note:            "disable inlining: complete rule namespaced variables (negation)",
+			query:           "data.test.p = x",
+			disableInlining: []string{"data.test.p"},
+			modules: []string{
+				`package test
+
+				p = y {
+					y = input
+					not y = 1
+				}
+				`,
+			},
+			wantQueries: []string{
+				`data.partial.test.p = x`,
+			},
+			wantSupport: []string{
+				`package partial.test
+
+				p = y1 { y1 = input; not y1 = 1 }`,
+			},
+		},
 		{
 			note:  "comprehensions: ref heads (with namespacing)",
 			query: "data.test.p = true; input.x = x",
@@ -1504,7 +1564,7 @@ func TestTopDownPartialEval(t *testing.T) {
 		{
 			note:        "comprehensions: ref heads (with live vars)",
 			query:       "x = [0]; y = {true | x[0]}",
-			wantQueries: []string{`y = {true | x0[0]; x0 = [0]}; x = [0]`},
+			wantQueries: []string{`y = {true | x[0]; x = [0]}; x = [0]`},
 		},
 	}