ast/compile: fix print rewriting for arrays in := and head vars

For both partial and complete rules, when variables from destructured
arrays in assignments had been used in the rule head, the print rewriting
would fail.

Now, this situation is accounted for by starting the safe VarSet with

    r.Head.Vars()

instead of

    r.Head.Args.Vars()

The former will check each of the Args, Key and Value fields of Head: if
it's non-empty, it'll add .Vars() to the returned VarSet.

Fixes open-policy-agent#4078.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

ast/compile.go

@@ -1371,7 +1371,7 @@ func (c *Compiler) rewritePrintCalls() {
 	for _, name := range c.sorted {
 		mod := c.Modules[name]
 		WalkRules(mod, func(r *Rule) bool {
-			safe := r.Head.Args.Vars()
+			safe := r.Head.Vars()
 			safe.Update(ReservedVars)
 			WalkBodies(r, func(b Body) bool {
 				for _, err := range rewritePrintCalls(c.localvargen, c.GetArity, safe, b) {

ast/compile_test.go

@@ -3169,6 +3169,36 @@ func TestCompilerRewritePrintCalls(t *testing.T) {
 			f(__local0__) = __local2__ { true; __local2__ = {1 | __local0__[x]; __local3__ = {__local1__ | __local1__ = x}; internal.print([__local3__])} }
 			`,
 		},
+		{
+			note: "print call of var in head key",
+			module: `package test
+			f(_) = [1, 2, 3]
+			p[x] { [_, x, _] := f(true); print(x) }`,
+			exp: `package test
+			f(__local0__) = [1, 2, 3] { true }
+			p[__local2__] { data.test.f(true, __local5__); [__local1__, __local2__, __local3__] = __local5__; __local6__ = {__local4__ | __local4__ = __local2__}; internal.print([__local6__]) }
+			`,
+		},
+		{
+			note: "print call of var in head value",
+			module: `package test
+			f(_) = [1, 2, 3]
+			p = x { [_, x, _] := f(true); print(x) }`,
+			exp: `package test
+			f(__local0__) = [1, 2, 3] { true }
+			p = __local2__ { data.test.f(true, __local5__); [__local1__, __local2__, __local3__] = __local5__; __local6__ = {__local4__ | __local4__ = __local2__}; internal.print([__local6__]) }
+			`,
+		},
+		{
+			note: "print call of vars in head key and value",
+			module: `package test
+			f(_) = [1, 2, 3]
+			p[x] = y { [_, x, y] := f(true); print(x) }`,
+			exp: `package test
+			f(__local0__) = [1, 2, 3] { true }
+			p[__local2__] = __local3__ { data.test.f(true, __local5__); [__local1__, __local2__, __local3__] = __local5__; __local6__ = {__local4__ | __local4__ = __local2__}; internal.print([__local6__]) }
+			`,
+		},
 	}
 
 	for _, tc := range cases {