diff --git a/internal/planner/planner.go b/internal/planner/planner.go index a68765feb7..3105e773f6 100644 --- a/internal/planner/planner.go +++ b/internal/planner/planner.go @@ -1360,6 +1360,10 @@ func (p *Planner) planObjectComprehension(oc *ast.ObjectComprehension, iter plan func (p *Planner) planComprehension(body ast.Body, closureIter planiter, target ir.Local, iter planiter) error { + // Variables that have been introduced in this comprehension have + // no effect on other parts of the policy, so they'll be dropped + // below. + p.vars.Push(map[ast.Var]ir.Local{}) prev := p.curr p.curr = &ir.Block{} ploc := p.loc @@ -1373,6 +1377,7 @@ func (p *Planner) planComprehension(body ast.Body, closureIter planiter, target block := p.curr p.curr = prev p.loc = ploc + p.vars.Pop() p.appendStmt(&ir.BlockStmt{ Blocks: []*ir.Block{ diff --git a/internal/planner/varstack_test.go b/internal/planner/varstack_test.go new file mode 100644 index 0000000000..e342065645 --- /dev/null +++ b/internal/planner/varstack_test.go @@ -0,0 +1,44 @@ +package planner + +import ( + "testing" + + "github.com/open-policy-agent/opa/ast" + "github.com/open-policy-agent/opa/internal/ir" +) + +func TestVarStackPushPop(t *testing.T) { + p := New() + vs := newVarstack() + vs.Push(map[ast.Var]ir.Local{}) + loc0, loc1 := p.newLocal(), p.newLocal() + vs.Put(ast.Var("x"), loc0) + vs.Push(map[ast.Var]ir.Local{}) + loc, ok := vs.Get(ast.Var("x")) + if exp, act := true, ok; exp != act { + t.Errorf("Get(x): expected %v, got %v", exp, act) + } + if exp, act := loc0, loc; exp != act { + t.Errorf("Get(x) expected %v, got %v", exp, act) + } + vs.Put(ast.Var("y"), loc1) + pop := vs.Pop() + if exp, act := 1, len(pop); exp != act { + t.Errorf("Pop(): expected len %v, got %v", exp, act) + } + + // x still there + loc, ok = vs.Get(ast.Var("x")) + if exp, act := true, ok; exp != act { + t.Errorf("Get(x): expected %v, got %v", exp, act) + } + if exp, act := loc0, loc; exp != act { + t.Errorf("Get(x) expected %v, got %v", exp, act) + } + + // y not there + _, ok = vs.Get(ast.Var("y")) + if exp, act := false, ok; exp != act { + t.Errorf("Get(y): expected %v, got %v", exp, act) + } +} diff --git a/test/cases/testdata/comprehensions/test-comprehensions-and-vars.yaml b/test/cases/testdata/comprehensions/test-comprehensions-and-vars.yaml new file mode 100644 index 0000000000..477a5fe202 --- /dev/null +++ b/test/cases/testdata/comprehensions/test-comprehensions-and-vars.yaml @@ -0,0 +1,16 @@ +cases: +- data: + modules: + - | + package test + xs := {"a", "b", "c"} + p = x { + y := { x | xs[x] } + z := { x | xs[x] } + count(y) == count(z) + x := count(y) + } + note: comprehensions/var bindings have no effect outside + query: data.test.p = x + want_result: + - x: 3 \ No newline at end of file