From a0044b1671d82cde873cbca4c2945adb43148ffa Mon Sep 17 00:00:00 2001
From: Torin Sandall <torinsandall@gmail.com>
Date: Fri, 25 Mar 2022 08:40:54 -0700
Subject: [PATCH] ast: Fix print call rewriting in else rules

The compiler was accidentally checking/rewriting print calls in the
body of else rules before the implicit args of the else rule were
processed. As a result, the compiler was generating false-positive
errors for refs to undeclared args. The root of the problem was the
usage of WalkBodies on the top-level rule (which implicitly walks the
bodies of else rules under the top-level rule). With this change, the
compiler will call WalkBodies on the head and body of each rule rather
than the entire rule itself (which includes the else chain).

Fixes #4489

Signed-off-by: Torin Sandall <torinsandall@gmail.com>
---
 ast/compile.go      |  6 ++++--
 ast/compile_test.go | 43 +++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/ast/compile.go b/ast/compile.go
index e17fa7dfba..6c738d0e10 100644
--- a/ast/compile.go
+++ b/ast/compile.go
@@ -1496,12 +1496,14 @@ func (c *Compiler) rewritePrintCalls() {
 		WalkRules(mod, func(r *Rule) bool {
 			safe := r.Head.Args.Vars()
 			safe.Update(ReservedVars)
-			WalkBodies(r, func(b Body) bool {
+			vis := func(b Body) bool {
 				for _, err := range rewritePrintCalls(c.localvargen, c.GetArity, safe, b) {
 					c.err(err)
 				}
 				return false
-			})
+			}
+			WalkBodies(r.Head, vis)
+			WalkBodies(r.Body, vis)
 			return false
 		})
 	}
diff --git a/ast/compile_test.go b/ast/compile_test.go
index 03ffac0960..f56c65a28d 100644
--- a/ast/compile_test.go
+++ b/ast/compile_test.go
@@ -4113,6 +4113,45 @@ func TestCompilerRewritePrintCalls(t *testing.T) {
 	}
 }
 
+func TestRewritePrintCallsWithElseImplicitArgs(t *testing.T) {
+
+	module := `package test
+
+	f(x, y) {
+		x = y
+	}
+
+	else = false {
+		print(x, y)
+	}`
+
+	c := NewCompiler().WithEnablePrintStatements(true)
+	opts := ParserOptions{AllFutureKeywords: true, unreleasedKeywords: true}
+	c.Compile(map[string]*Module{
+		"test.rego": MustParseModuleWithOpts(module, opts),
+	})
+
+	if c.Failed() {
+		t.Fatal(c.Errors)
+	}
+
+	exp := MustParseModuleWithOpts(`package test
+
+	f(__local0__, __local1__) = true { __local0__ = __local1__ }
+	else = false { __local6__ = {__local4__ | __local4__ = __local2__}; __local7__ = {__local5__ | __local5__ = __local3__}; internal.print([__local6__, __local7__]) }
+	`, opts)
+
+	// NOTE(tsandall): we have to patch the implicit args on the else rule
+	// because of how the parser copies the arg names across from the first
+	// rule.
+	exp.Rules[0].Else.Head.Args[0] = VarTerm("__local2__")
+	exp.Rules[0].Else.Head.Args[1] = VarTerm("__local3__")
+
+	if !exp.Equal(c.Modules["test.rego"]) {
+		t.Fatalf("Expected:\n\n%v\n\nGot:\n\n%v", exp, c.Modules["test.rego"])
+	}
+}
+
 func TestCompilerMockFunction(t *testing.T) {
 	c := NewCompiler()
 	c.Modules["test"] = MustParseModule(`
@@ -4299,7 +4338,7 @@ func TestCompilerCheckUnusedAssignedVar(t *testing.T) {
 		{
 			note: "rule with nested closure",
 			module: `package test
-				p { 
+				p {
 					x := 1
 					a := 1
 					{ y | y := [ z | z:=[1,2,3][a]; z > 1 ][_] }
@@ -4312,7 +4351,7 @@ func TestCompilerCheckUnusedAssignedVar(t *testing.T) {
 		{
 			note: "rule with nested closure and unused inner var",
 			module: `package test
-				p { 
+				p {
 					x := 1
 					{ y | y := [ z | z:=[1,2,3][x]; z > 1; a := 2 ][_] }
 				}