diff --git a/internal/planner/planner.go b/internal/planner/planner.go index 1d2ffec2ff..c9eb0f3ad6 100644 --- a/internal/planner/planner.go +++ b/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() diff --git a/internal/planner/rules_test.go b/internal/planner/rules_test.go index 072afce980..1b6c355331 100644 --- a/internal/planner/rules_test.go +++ b/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) + } + }) + } +}