Print undeclared var names

v1/ast/compile.go

@@ -2682,11 +2682,12 @@ func (c *Compiler) rewritePrintCalls() {
 				}
 
 				bodyVis := func(b Body) bool {
-					modrec, errs := rewritePrintCalls(c.localvargen, c.GetArity, vis.vars, b)
+					modrec, errs := rewritePrintCalls(c.localvargen, c.GetArity, vis.vars, c.RewrittenVars, b)
 					if modrec {
 						modified = true
 					}
-					if !c.err(errs...) {
+					if len(errs) > 0 {
+						c.err(errs...)
 						return true
 					}
 					return false
@@ -2732,44 +2733,44 @@ func checkVoidCalls(env *TypeEnv, x any) Errors {
 // The expression would be rewritten to:
 //
 //	print({__local0__ | __local0__ = "the value of x is:"}, {__local1__ | __local1__ = input.x})
-func rewritePrintCalls(gen *localVarGenerator, getArity func(Ref) int, globals VarSet, body Body) (bool, Errors) {
+func rewritePrintCalls(gen *localVarGenerator, getArity func(Ref) int, globals VarSet, rewritten map[Var]Var, body Body) (bool, Errors) {
 
 	var errs Errors
 	var modified bool
 
-	// Visit comprehension bodies recursively to ensure print statements inside
-	// those bodies only close over variables that are safe.
+	// Visit nested bodies recursively to ensure print statements inside those
+	// bodies only close over variables that are safe.
 	for i := range body {
-		if ContainsClosures(body[i]) {
+		if containsNestedBody(body[i]) {
 			safe := outputVarsForBody(body[:i], getArity, globals, nil)
 			safe.Update(globals)
 			WalkClosures(body[i], func(x any) bool {
 				var modrec bool
 				var errsrec Errors
 				switch x := x.(type) {
 				case *SetComprehension:
-					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, x.Body)
+					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, rewritten, x.Body)
 				case *ArrayComprehension:
-					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, x.Body)
+					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, rewritten, x.Body)
 				case *ObjectComprehension:
-					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, x.Body)
+					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, rewritten, x.Body)
 				case *Every:
 					safe.Update(x.KeyValueVars())
-					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, x.Body)
+					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, rewritten, x.Body)
 				case *Not:
-					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, x.Body)
+					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, rewritten, x.Body)
 				case *LogicalAnd:
 					var modR bool
 					var errsR Errors
-					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, x.Lhs)
-					modR, errsR = rewritePrintCalls(gen, getArity, safe, x.Rhs)
+					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, rewritten, x.Lhs)
+					modR, errsR = rewritePrintCalls(gen, getArity, safe, rewritten, x.Rhs)
 					modrec = modrec || modR
 					errsrec = append(errsrec, errsR...)
 				case *LogicalOr:
 					var modR bool
 					var errsR Errors
-					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, x.Lhs)
-					modR, errsR = rewritePrintCalls(gen, getArity, safe, x.Rhs)
+					modrec, errsrec = rewritePrintCalls(gen, getArity, safe, rewritten, x.Lhs)
+					modR, errsR = rewritePrintCalls(gen, getArity, safe, rewritten, x.Rhs)
 					modrec = modrec || modR
 					errsrec = append(errsrec, errsR...)
 				}
@@ -2821,6 +2822,9 @@ func rewritePrintCalls(gen *localVarGenerator, getArity func(Ref) int, globals V
 			if vars.DiffCount(safe) > 0 {
 				unsafe := vars.Diff(safe)
 				for _, v := range unsafe.Sorted() {
+					if w, ok := rewritten[v]; ok {
+						v = w
+					}
 					errs = append(errs, NewError(CompileErr, args[j].Loc(), "var %v is undeclared", v))
 				}
 			}
@@ -2847,6 +2851,18 @@ func rewritePrintCalls(gen *localVarGenerator, getArity func(Ref) int, globals V
 	return modified, nil
 }
 
+// containsNestedBody returns true if x contains any node that carries a nested
+// body which rewritePrintCalls needs to descend into. This is a superset of
+// ContainsClosures, which ignores not/and/or expressions.
+func containsNestedBody(x any) bool {
+	found := false
+	WalkClosures(x, func(any) bool {
+		found = true
+		return found
+	})
+	return found
+}
+
 func erasePrintCalls(node any) bool {
 	var modified bool
 	NewGenericVisitor(func(x any) bool {
@@ -3867,7 +3883,7 @@ func (qc *queryCompiler) rewritePrintCalls(_ *QueryContext, body Body) (Body, er
 		return cpy, nil
 	}
 	gen := newLocalVarGenerator("q", body)
-	if _, errs := rewritePrintCalls(gen, qc.compiler.GetArity, ReservedVars, body); len(errs) > 0 {
+	if _, errs := rewritePrintCalls(gen, qc.compiler.GetArity, ReservedVars, qc.RewrittenVars(), body); len(errs) > 0 {
 		return nil, errs
 	}
 	return body, nil

v1/ast/compile_cases_test.go

@@ -52,7 +52,7 @@ func runCompileCase(t *testing.T, tc compilecases.TestCase) {
 		modules[name] = parsed
 	}
 
-	c := NewCompiler().WithStrict(tc.Strict)
+	c := NewCompiler().WithStrict(tc.Strict).WithEnablePrintStatements(tc.PrintStatements)
 	c.Compile(modules)
 
 	if !c.Failed() {

v1/ast/compile_test.go

@@ -7821,6 +7821,7 @@ func TestCompilerRewritePrintCallsErrors(t *testing.T) {
 		note    string
 		module  string
 		exp     error
+		exps    []string // when set, asserts the full set of error messages
 		errCode string
 	}{
 		{
@@ -7845,6 +7846,7 @@ func TestCompilerRewritePrintCallsErrors(t *testing.T) {
 			p if { {1 | print(x)} = {1 | print(7)} }
 			`,
 			exp:     errors.New("var x is undeclared"),
+			exps:    []string{"var x is undeclared"},
 			errCode: CompileErr,
 		},
 		{
@@ -7855,6 +7857,40 @@ func TestCompilerRewritePrintCallsErrors(t *testing.T) {
 			exp:     errors.New("print(42) used as value"),
 			errCode: TypeErr,
 		},
+		{
+			note: "some declaration inside comprehension body",
+			module: `package test
+
+			p if {
+				xs := [1 | some x; print(x)]
+				xs == xs
+			}`,
+			exp:     errors.New("var x is undeclared"),
+			exps:    []string{"var x is undeclared"},
+			errCode: CompileErr,
+		},
+		{
+			note: "declared var shadowed inside comprehension",
+			module: `package test
+
+			p if {
+				x := 1
+				ys := [x | some x; print(x)]
+				ys == ys
+			}`,
+			exp:     errors.New("var x is undeclared"),
+			exps:    []string{"var x is undeclared"},
+			errCode: CompileErr,
+		},
+		{
+			note: "unsafe var alongside function argument",
+			module: `package test
+
+			f(x) if { print(x, y) }`,
+			exp:     errors.New("var y is undeclared"),
+			exps:    []string{"var y is undeclared"},
+			errCode: CompileErr,
+		},
 	}
 
 	for _, tc := range cases {
@@ -7869,6 +7905,15 @@ func TestCompilerRewritePrintCallsErrors(t *testing.T) {
 			if c.Errors[0].Code != tc.errCode || c.Errors[0].Message != tc.exp.Error() {
 				t.Fatal("unexpected error:", c.Errors)
 			}
+			if tc.exps != nil {
+				got := make([]string, len(c.Errors))
+				for i, err := range c.Errors {
+					got[i] = err.Message
+				}
+				if !slices.Equal(got, tc.exps) {
+					t.Fatalf("expected errors %v but got %v", tc.exps, got)
+				}
+			}
 		})
 	}
 }
@@ -7880,6 +7925,38 @@ func TestCompilterRewritePrintCallsNestedComprehensionLocalsSafe(t *testing.T) {
 		note   string
 		module string
 	}{
+		{
+			note: "comprehension in every domain",
+			module: `package test
+			p if { every z in [c | c := 1; print(c)] { z == z } }`,
+		},
+		{
+			note: "comprehension in every body",
+			module: `package test
+			p if { every z in [1] { y := [1 | print(z)]; y == y } }`,
+		},
+		{
+			note: "comprehension in rule head",
+			module: `package test
+			p := {1 | print("h")}`,
+		},
+		{
+			note: "comprehension in with value",
+			module: `package test
+			q := 1
+			p if { q with input as [1 | print(input)] }`,
+		},
+		{
+			note: "print after not expression",
+			module: `package test
+			q := 1
+			r if { not q; print(q) }`,
+		},
+		{
+			note: "outer and every-key vars visible in nested comprehension",
+			module: `package test
+			p if { a := 1; every z in [1] { b := [1 | print(a, z)]; b == b } }`,
+		},
 		{
 			note: "print variable from nested comprehension without error",
 			module: `package test
@@ -7898,6 +7975,12 @@ func TestCompilterRewritePrintCallsNestedComprehensionLocalsSafe(t *testing.T) {
 			c := NewCompiler().WithEnablePrintStatements(true)
 			c.Compile(map[string]*Module{"test.rego": module(tc.module)})
 			assertNotFailed(t, c)
+
+			// Every print call must have been rewritten into internal.print,
+			// including those in nested bodies.
+			if str := c.Modules["test.rego"].String(); strings.Contains(str, " print(") {
+				t.Fatalf("expected all print calls to be rewritten, got:\n\n%v", str)
+			}
 		})
 	}
 }

v1/test/compilecases/cases.go

@@ -39,6 +39,7 @@ type TestCase struct {
 	RegoVersion          string   `json:"rego_version,omitempty"           yaml:"rego_version,omitempty"`          // rego version to parse the modules as: v0, v1 (default), or v0-compat-v1
 	Strict               bool     `json:"strict,omitempty"                 yaml:"strict,omitempty"`                // enable the compiler's strict mode
 	ExperimentalKeywords bool     `json:"experimental_keywords,omitempty"  yaml:"experimental_keywords,omitempty"` // opt-in to experimental future keywords
+	PrintStatements      bool     `json:"print_statements,omitempty"       yaml:"print_statements,omitempty"`      // keep print() calls instead of erasing them, as required to reach diagnostics about their operands
 	WantErrors           []Error  `json:"want_errors"                      yaml:"want_errors"`                     // diagnostics the compilation must produce
 	Exhaustive           bool     `json:"exhaustive,omitempty"             yaml:"exhaustive,omitempty"`            // require want_errors to be the complete set, not a subset
 }

v1/test/compilecases/testdata/v1/print/test-print-undeclared-var.yaml

@@ -0,0 +1,65 @@
+---
+# print() operands are captured in comprehensions so that an undefined operand
+# does not stop the print from being reached. That rewrite runs after local vars
+# have been renamed to generated ones, so a diagnostic about an operand has to
+# map the name back to the one the author wrote. See issue #5624.
+#
+# These cases need print_statements: with print erased, there is no operand left
+# to diagnose and the modules compile.
+cases:
+  - note: print/undeclared-var-declared-by-some
+    print_statements: true
+    modules:
+      - |
+        package test
+
+        p if {
+        	some x
+        	print(x)
+        }
+    want_errors:
+      - code: rego_compile_error
+        row: 5
+        col: 8
+        message: var x is undeclared
+    exhaustive: true
+  - note: print/undeclared-var-declared-by-some-reports-each-operand
+    print_statements: true
+    modules:
+      - |
+        package test
+
+        p if {
+        	some x, y
+        	print(x, y)
+        }
+    want_errors:
+      - code: rego_compile_error
+        row: 5
+        col: 8
+        message: var x is undeclared
+      - code: rego_compile_error
+        row: 5
+        col: 11
+        message: var y is undeclared
+    exhaustive: true
+  - note: print/undeclared-var-in-every-body-does-not-implicate-key-var
+    # `z` is bound by the every, so only `q` is undeclared. Reporting is per
+    # print operand, and the body is walked once: a var that is safe inside the
+    # nested body must not be reported, and `q` must not be reported twice.
+    print_statements: true
+    modules:
+      - |
+        package test
+
+        p if {
+        	every z in [1] {
+        		print(q, z)
+        	}
+        }
+    want_errors:
+      - code: rego_compile_error
+        row: 5
+        col: 9
+        message: var q is undeclared
+    exhaustive: true