planner: Correct crash in virtual doc iteration

There was an issue with iterating over virtual docs when planning
an expression that had `with` keywords. The iterator would potentially
be called multiple times, so the "with" values needed to be pushed
and popped accordingly. This essentially copies the pattern used in
the topdown eval flow.

Fixes: open-policy-agent#2601
Signed-off-by: Patrick East <east.patrick@gmail.com>

internal/planner/planner.go

@@ -537,22 +537,42 @@ func (p *Planner) planWith(e *ast.Expr, iter planiter) error {
 			}
 		}
 
-		return p.planWithRec(e, paths, locals, 0, func() error {
-
+		err := p.planWithRec(e, paths, locals, 0, func() error {
 			if len(dataRefs) > 0 {
 				p.funcs.Pop()
 				for i := len(dataRefs) - 1; i >= 0; i-- {
 					p.rules.Pop(dataRefs[i])
 				}
 			}
 
-			return p.planWithUndoRec(restore, 0, iter)
+			err := p.planWithUndoRec(restore, 0, func() error {
+
+				err := iter()
+
+				if len(dataRefs) > 0 {
+					p.funcs.Push(map[string]string{})
+					for _, ref := range dataRefs {
+						p.rules.Push(ref)
+					}
+				}
+				return err
+			})
+
+			return err
 		})
+
+		if len(dataRefs) > 0 {
+			p.funcs.Pop()
+			for i := len(dataRefs) - 1; i >= 0; i-- {
+				p.rules.Pop(dataRefs[i])
+			}
+		}
+		return err
+
 	})
 }
 
 func (p *Planner) planWithRec(e *ast.Expr, targets [][]int, values []ir.Local, index int, iter planiter) error {
-
 	if index >= len(e.With) {
 		return p.planExpr(e.NoWith(), iter)
 	}

internal/planner/planner_test.go

@@ -195,6 +195,17 @@ func TestPlannerHelloWorld(t *testing.T) {
 				q = 2`,
 			},
 		},
+		{
+			note:    "with keyword - virtual doc iteration",
+			queries: []string{`x = data[i][j] with data.bar as 1; y = "a"`},
+			modules: []string{
+				`package foo
+				
+				p = 0
+				q = 1
+				r = 2`,
+			},
+		},
 	}
 
 	for _, tc := range tests {

test/wasm/assets/016_with.yaml

@@ -229,4 +229,58 @@ cases:
   - note: with conflict
     query: |
       input = x with input.foo as 1 with input.foo.bar as 2
-    want_error: with target conflict
+    want_error: with target conflict
+  - note: with virtual doc iteration
+    query: |
+      x := data[i][j] with data.bar.p as 3 with data.bar.q as 4; y = data.bar.p; z = data.bar.q
+    modules:
+      - |
+        package foo
+        p = 0
+        q = 1
+        r = 2
+      - |
+        package bar
+        p = -1
+    data: {
+      'bar': {
+        'q': -2,
+      }
+    }
+    want_result: [
+      {
+        "i": "foo",
+        "j": "p",
+        "x": 0,
+        "y": -1,
+        "z": -2,
+      },
+      {
+        "i": "foo",
+        "j": "q",
+        "x": 1,
+        "y": -1,
+        "z": -2,
+      },
+      {
+        "i": "foo",
+        "j": "r",
+        "x": 2,
+        "y": -1,
+        "z": -2,
+      },
+      {
+        "i": "bar",
+        "j": "p",
+        "x": 3,
+        "y": -1,
+        "z": -2,
+      },
+      {
+        "i": "bar",
+        "j": "q",
+        "x": 4,
+        "y": -1,
+        "z": -2,
+      },
+    ]