ast: restore location on unsafe-var errors from rewritten head vars

[arpitjain099]

What

opa check on a policy with an unsafe variable that appears in a rule head and is rewritten (for example, declared with some) reports the error without a location:

$ opa check p.rego
1 error occurred: rego_unsafe_var_error: var x is unsafe

As noted in #8719, this regressed in v1.14.0 (v1.13 reported p.rego:3: rego_unsafe_var_error: var x is unsafe).

Why

In checkSafetyRuleHeads, the error is built with vars[v].Location, where vars comes from r.Head.Vars() and is keyed by the head variables. The location is read after v has been reassigned to its original name via RewrittenVars:

for v := range unsafe {
    if w, ok := c.RewrittenVars[v]; ok {
        v = w
    }
    ...
    NewError(UnsafeVarErr, vars[v].Location, ...)

When the head var is rewritten (e.g. by some x), vars is keyed by the rewritten name, so vars[v] after the reassignment misses and returns a zero entry with a nil location.

Fix

Capture the location from the original key before reassigning v. With the fix:

$ opa check p.rego
1 error occurred: p.rego:3: rego_unsafe_var_error: var x is unsafe

Added a regression case to TestCompilerCheckSafetyVarLoc (it nil-panics on the location lookup without the fix). The full v1/ast suite and gofmt pass.

Fixes #8719

[johanfylling]

Thank you! 🙂

[johanfylling]

We could consider the root cause of the issue to be where the rewrite is happening. Could there be a slight risk of a regression here if the location isn't always attached to the rewritten var?

Maybe we could make this handle a wider variety of upstream mishaps by simply using the location from v or w; whichever is non-nil.