Return correct location of unsafe var in object (#8371)

Signed-off-by: Sebastian Spaink <sebastianspaink@gmail.com>

v1/ast/compile.go

@@ -1545,7 +1545,7 @@ func (c *Compiler) checkSafetyRuleHeads() {
 							v = w
 						}
 						if !v.IsGenerated() {
-							if !c.err(NewError(UnsafeVarErr, r.Loc(), "var %v is unsafe", v)) {
+							if !c.err(NewError(UnsafeVarErr, vars[v].Location, "var %v is unsafe", v)) {
 								return true
 							}
 						}
@@ -4431,9 +4431,9 @@ type unsafeVars map[*Expr]VarSet
 
 func (vs unsafeVars) Add(e *Expr, v Var) {
 	if u, ok := vs[e]; ok {
-		u[v] = struct{}{}
+		u[v] = struct{ *Location }{}
 	} else {
-		vs[e] = VarSet{v: struct{}{}}
+		vs[e] = VarSet{v: struct{ *Location }{}}
 	}
 }
 

v1/ast/compile_test.go

@@ -1819,27 +1819,78 @@ func TestCompilerCheckSafetyBodyErrors(t *testing.T) {
 
 func TestCompilerCheckSafetyVarLoc(t *testing.T) {
 
-	_, err := CompileModules(map[string]string{"test.rego": `package test
+	tests := []struct {
+		module       string
+		expectedErrs []struct {
+			message  string
+			location int
+		}
+	}{
+		{
+			module: `package test
 import rego.v1
 
 p if {
 	not x
 	x > y
-}`})
+}`,
+			expectedErrs: []struct {
+				message  string
+				location int
+			}{
+				{
+					"var x is unsafe",
+					5,
+				},
+				{
+					"var y is unsafe",
+					6,
+				},
+			},
+		},
+		{
+			module: `package play
 
-	if err == nil {
-		t.Fatal("expected error")
+obj := {
+    "foo": "bar",
+    "baz": qux,
+}
+`,
+			expectedErrs: []struct {
+				message  string
+				location int
+			}{
+				{
+					"var qux is unsafe",
+					5,
+				},
+			},
+		},
 	}
 
-	errs := err.(Errors)
+	for _, tc := range tests {
+		t.Run(tc.module, func(t *testing.T) {
+			_, err := CompileModules(map[string]string{"test.rego": tc.module})
 
-	if !strings.Contains(errs[0].Message, "var x is unsafe") || errs[0].Location.Row != 5 {
-		t.Fatal("expected error on row 5 but got:", err)
+			if err == nil {
+				t.Fatal("expected error")
+			}
+
+			var errs Errors
+			errors.As(err, &errs)
+
+			if len(errs) != len(tc.expectedErrs) {
+				t.Fatalf("expected %d errors, got %d", len(tc.expectedErrs), len(errs))
+			}
+
+			for i := range errs {
+				if !strings.Contains(errs[i].Message, tc.expectedErrs[i].message) || errs[i].Location.Row != tc.expectedErrs[i].location {
+					t.Fatalf("expected error on row %d but got: %s", tc.expectedErrs[i].location, errs[i])
+				}
+			}
+		})
 	}
 
-	if !strings.Contains(errs[1].Message, "var y is unsafe") || errs[1].Location.Row != 6 {
-		t.Fatal("expected y is unsafe on row 6 but got:", err)
-	}
 }
 
 func TestCompilerCheckSafetyFunctionAndContainsKeyword(t *testing.T) {

v1/ast/varset.go

@@ -12,7 +12,7 @@ import (
 )
 
 // VarSet represents a set of variables.
-type VarSet map[Var]struct{}
+type VarSet map[Var]struct{ *Location }
 
 // NewVarSet returns a new VarSet containing the specified variables.
 func NewVarSet(vs ...Var) VarSet {
@@ -30,7 +30,16 @@ func NewVarSetOfSize(size int) VarSet {
 
 // Add updates the set to include the variable "v".
 func (s VarSet) Add(v Var) {
-	s[v] = struct{}{}
+	if _, ok := s[v]; !ok {
+		s[v] = struct{ *Location }{}
+	}
+}
+
+func (s VarSet) AddLocation(v Var, l *Location) {
+	if entry, ok := s[v]; ok {
+		entry.Location = l
+		s[v] = entry
+	}
 }
 
 // Contains returns true if the set contains the variable "v".
@@ -54,6 +63,7 @@ func (s VarSet) Diff(vs VarSet) VarSet {
 	for v := range s {
 		if !vs.Contains(v) {
 			r.Add(v)
+			r.AddLocation(v, s[v].Location)
 		}
 	}
 	return r

v1/ast/visit.go

@@ -878,6 +878,7 @@ func (vis *VarVisitor) visit(v any) bool {
 	}
 	if v, ok := v.(Var); ok {
 		vis.Add(v)
+		return true
 	}
 	return false
 }
@@ -941,6 +942,9 @@ func (vis *VarVisitor) Walk(x any) {
 		vis.Walk(x.Value.Value)
 	case *Term:
 		vis.Walk(x.Value)
+		if vVar, ok := x.Value.(Var); ok {
+			vis.vars.AddLocation(vVar, x.Location)
+		}
 	case Ref:
 		for i := range x {
 			vis.Walk(x[i].Value)
@@ -1007,6 +1011,9 @@ func (vis *VarVisitor) WalkRef(ref Ref) {
 	}
 	for _, term := range ref {
 		vis.Walk(term.Value)
+		if vVar, ok := term.Value.(Var); ok {
+			vis.vars.AddLocation(vVar, term.Location)
+		}
 	}
 }