ast/compile: also replace declared vars in with's target

It might be uncommon to do this, but it's not wrong to expect this to work.

Fixes open-policy-agent#6979.

Signed-off-by: Stephan Renatus <[email protected]>

ast/compile.go

@@ -5332,6 +5332,7 @@ func rewriteDeclaredVarsInExpr(g *localVarGenerator, stack *localDeclaredVars, e
 		case *Term:
 			stop, errs = rewriteDeclaredVarsInTerm(g, stack, x, errs, strict)
 		case *With:
+			errs = rewriteDeclaredVarsInTermRecursive(g, stack, x.Target, errs, strict)
 			errs = rewriteDeclaredVarsInTermRecursive(g, stack, x.Value, errs, strict)
 			stop = true
 		}

ast/compile_test.go

@@ -5673,6 +5673,23 @@ func TestRewriteDeclaredVars(t *testing.T) {
 				p { __local1__ = data.test.y; data.test.q[[__local1__, __local0__]] }
 			`,
 		},
+		{
+			note: "with: rewrite target",
+			module: `
+				package test
+				p {
+					x := "foo"
+					true with input[x] as 1
+				}
+			`,
+			exp: `
+				package test
+				p {
+					__local0__ = "foo";
+					true with input[__local0__] as 1
+				}
+			`,
+		},
 		{
 			note: "single-value rule with ref head",
 			module: `
@@ -6038,6 +6055,26 @@ func TestRewriteDeclaredVars(t *testing.T) {
 				}
 			`,
 		},
+		{
+			note: "rewrite every: with modifier on body, using every's key+value",
+			module: `
+				package test
+				# import future.keywords.in
+				# import future.keywords.every
+				p {
+					every x, y in input { true with data.test.q[x][y] as 100 }
+				}
+			`,
+			exp: `
+				package test
+				p {
+				    __local2__ = input
+					every __local0__, __local1__ in __local2__ {
+						true with data.test.q[__local0__][__local1__] as 100
+					}
+				}
+			`,
+		},
 		{
 			note: "rewrite closures",
 			module: `