From afe33a5d7cc4f67b9687a2e372d428ac90ff5a60 Mon Sep 17 00:00:00 2001
From: Stephan Renatus <stephan.renatus@gmail.com>
Date: Tue, 30 Mar 2021 13:29:02 +0200
Subject: [PATCH] planner: forget variables from comprehensions (#3327)

Before, variables introduced in a comprehension, like the `x` in

    { x | xs[x] }

would have been recorded in `p.vars`, and have an effect outside of
the scope of this comprehension, leading to wrong plans.

Now, we'll push a new frame for learning about variables we've not
known before, and reset to that state after planning the
comprehensions, thereby forgetting them.

Fixes #3325.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
---
 internal/planner/planner.go                   |  5 +++
 internal/planner/varstack_test.go             | 44 +++++++++++++++++++
 .../test-comprehensions-and-vars.yaml         | 16 +++++++
 3 files changed, 65 insertions(+)
 create mode 100644 internal/planner/varstack_test.go
 create mode 100644 test/cases/testdata/comprehensions/test-comprehensions-and-vars.yaml

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