wasm: only shadow ruletrie with matching data refs (#3153)

Before, the planner had decided to shadow all rules whenever a with statement was targeting "data".
Now, it's looking at the ruletrie to see if whatever ref the with statement targets actually is a rule (or package ref). If so, it'll shadow, if not, it won't.

Fixes #3150.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

internal/planner/planner.go

@@ -608,11 +608,12 @@ func (p *Planner) planWith(e *ast.Expr, iter planiter) error {
 			restore[i] = [2]ir.Local{lorig, lsave}
 		}
 
-		// If any of the with statements targeted the data document we shadow
-		// the existing planned functions during expression planning. This
-		// causes the planner to re-plan any rules that may be required during
-		// planning of this expression (transitively).
-		if len(dataRefs) > 0 {
+		// If any of the `with` statements targeted the data document, overwriting
+		// parts of the ruletrie, we shadow the existing planned functions during
+		// expression planning. This causes the planner to re-plan any rules that
+		// may be required during planning of this expression (transitively).
+		shadowing := p.dataRefsShadowRuletrie(dataRefs)
+		if shadowing {
 			p.funcs.Push(map[string]string{})
 			for _, ref := range dataRefs {
 				p.rules.Push(ref)
@@ -620,7 +621,7 @@ func (p *Planner) planWith(e *ast.Expr, iter planiter) error {
 		}
 
 		err := p.planWithRec(e, paths, locals, 0, func() error {
-			if len(dataRefs) > 0 {
+			if shadowing {
 				p.funcs.Pop()
 				for i := len(dataRefs) - 1; i >= 0; i-- {
 					p.rules.Pop(dataRefs[i])
@@ -631,7 +632,7 @@ func (p *Planner) planWith(e *ast.Expr, iter planiter) error {
 
 				err := iter()
 
-				if len(dataRefs) > 0 {
+				if shadowing {
 					p.funcs.Push(map[string]string{})
 					for _, ref := range dataRefs {
 						p.rules.Push(ref)
@@ -643,7 +644,7 @@ func (p *Planner) planWith(e *ast.Expr, iter planiter) error {
 			return err
 		})
 
-		if len(dataRefs) > 0 {
+		if shadowing {
 			p.funcs.Pop()
 			for i := len(dataRefs) - 1; i >= 0; i-- {
 				p.rules.Pop(dataRefs[i])
@@ -709,6 +710,15 @@ func (p *Planner) planWithUndoRec(restore [][2]ir.Local, index int, iter planite
 	return nil
 }
 
+func (p *Planner) dataRefsShadowRuletrie(refs []ast.Ref) bool {
+	for _, ref := range refs {
+		if p.rules.Lookup(ref) != nil {
+			return true
+		}
+	}
+	return false
+}
+
 func (p *Planner) planExprTerm(e *ast.Expr, iter planiter) error {
 	return p.planTerm(e.Terms.(*ast.Term), func() error {
 		falsy := p.newLocal()

internal/planner/rules_test.go

@@ -6,6 +6,8 @@ package planner
 
 import (
 	"testing"
+
+	"github.com/open-policy-agent/opa/ast"
 )
 
 func TestFuncstack(t *testing.T) {
@@ -75,3 +77,50 @@ func TestFuncstack(t *testing.T) {
 		t.Errorf("expected fs gen to be %d, got %d", exp, act)
 	}
 }
+
+func TestDataRefsShadowRuletrie(t *testing.T) {
+	p := New()
+	rt := p.rules
+	rt.Insert(ast.MustParseRef(("data.foo.bar")))
+	rt.Insert(ast.MustParseRef(("data.foo.baz")))
+	rt.Insert(ast.MustParseRef(("data.foo.bar.quz")))
+
+	tests := []struct {
+		note string
+		refs []ast.Ref
+		exp  bool
+	}{
+		{
+			note: "no refs",
+			refs: nil,
+			exp:  false,
+		},
+		{
+			note: "data root node",
+			refs: []ast.Ref{ast.MustParseRef("data")},
+			exp:  true,
+		},
+		{
+			note: "one ref only, mismatch in first level",
+			refs: []ast.Ref{ast.MustParseRef("data.quz")},
+			exp:  false,
+		},
+		{
+			note: "two refs, matching 2nd",
+			refs: []ast.Ref{
+				ast.MustParseRef("data.quz"),
+				ast.MustParseRef("data.foo"),
+			},
+			exp: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+			act := p.dataRefsShadowRuletrie(tc.refs)
+			if tc.exp != act {
+				t.Errorf("expected %v, got %v", tc.exp, act)
+			}
+		})
+	}
+}