Fix false positivie in prefer-equals-comparison (#1828)

We'd previously fail to account for comprehension term vars in this rule.
I'm pretty sure the logic here needs to extend into other locations in Regal,
but this at least fixes the immediate issue as reported.

Fixes #1826

Signed-off-by: Anders Eknert <anders.eknert@apple.com>

Author: anderseknert

bundle/regal/rules/idiomatic/prefer-equals-comparison/prefer_equals_comparison.rego

@@ -4,6 +4,7 @@ package regal.rules.idiomatic["prefer-equals-comparison"]
 
 import data.regal.ast
 import data.regal.result
+import data.regal.util
 
 report contains violation if {
 	some rule_index, expr
@@ -26,5 +27,33 @@ _unassignable(term, _) if {
 
 _unassignable(term, rule_index) if {
 	term.type == "var"
-	not ast.is_output_var(input.rules[to_number(rule_index)], term)
+	ri := to_number(rule_index)
+	not ast.is_output_var(input.rules[ri], term)
+	not _is_declared_comp_term(term, rule_index)
+}
+
+_is_declared_comp_term(term, rule_index) if {
+	some comp_term in _comprehension_term_vars_in_scope(rule_index, term.location)
+	comp_term.type == "var"
+	comp_term.value == term.value
+}
+
+_comprehension_term_vars_in_scope(rule_index, location) := {node |
+	loc := util.to_location_object(location)
+
+	some comp in ast.found.comprehensions[rule_index]
+	util.contains_location(util.to_location_object(comp.location), loc)
+
+	fields := {
+		"arraycomprehension": ["term"],
+		"objectcomprehension": ["key", "value"],
+		"setcomprehension": ["term"],
+	}[comp.type]
+
+	some field in fields
+	term := comp.value[field]
+
+	walk(term, [_, node])
+
+	node.type == "var"
 }

bundle/regal/rules/idiomatic/prefer-equals-comparison/prefer_equals_comparison_test.rego

@@ -117,3 +117,39 @@ test_success_not_impossible_assignment_output_var_equals_static_ref if {
 
 	r == set()
 }
+
+test_success_not_impossible_assignment_set_comprehension_term if {
+	r := rule.report with input as ast.policy(`r := {x | x = y[_].z}`)
+
+	r == set()
+}
+
+test_success_not_impossible_assignment_set_comprehension_term_complex if {
+	r := rule.report with input as ast.policy(`r := {{x} | x = y[_].z}`)
+
+	r == set()
+}
+
+test_success_not_impossible_assignment_array_comprehension_term if {
+	r := rule.report with input as ast.policy(`r := [x | x = y[_].z]`)
+
+	r == set()
+}
+
+test_success_not_impossible_assignment_array_comprehension_term_complex if {
+	r := rule.report with input as ast.policy(`r := [[x] | x = y[_].z]`)
+
+	r == set()
+}
+
+test_success_not_impossible_assignment_object_comprehension_term if {
+	r := rule.report with input as ast.policy(`r := {k: v | k = y[v].k}`)
+
+	r == set()
+}
+
+test_success_not_impossible_assignment_object_comprehension_term_complex if {
+	r := rule.report with input as ast.policy(`r := {{k}: {v} | k = y[v].k}`)
+
+	r == set()
+}

bundle/regal/rules/imports/unresolved-reference/unresolved_reference.rego

@@ -143,5 +143,5 @@ _to_location_object(loc, text, file) := {"location": {
 	from_col := substring(text, col - 1, -1)
 	ref_text := substring(from_col, 0, indexof(from_col, " "))
 
-	end_col := to_number(vals[1]) + count(ref_text)
+	end_col := col + count(ref_text)
 }

bundle/regal/util/util.rego

@@ -228,3 +228,23 @@ parse_bool("FALSE") := false
 # METADATA
 # description: creates a string where s is repeated n times
 repeat(s, n) := replace(sprintf("%-*s", [n, " "]), " ", s)
+
+# METADATA
+# description: true if location 'sub' is fully contained within location 'sup'
+contains_location(sup, sub) if {
+	sup.row < sub.row
+	sup.end.row > sub.end.row
+} else if {
+	sup.row == sub.row
+	sup.col <= sub.col
+	sup.end.row > sub.end.row
+} else if {
+	sup.row < sub.row
+	sup.end.row == sub.end.row
+	sup.end.col >= sub.end.col
+} else if {
+	sup.row == sub.row
+	sup.col <= sub.col
+	sup.end.row == sub.end.row
+	sup.end.col >= sub.end.col
+}