diff --git a/ast/compile.go b/ast/compile.go
index f446606d16..983af7374e 100644
--- a/ast/compile.go
+++ b/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) {
diff --git a/ast/compile_test.go b/ast/compile_test.go
index dce2baf216..223f7eac7b 100644
--- a/ast/compile_test.go
+++ 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 {