perf: improve performance of new unresolved-reference rule (#1530)

We've had a few heavy rules in the past, but nothing like this
one. That's not really due to the implementation of the rule, but
simply the fact that it aggregates every ref found everywhere, which
will be a lot no matter what.

Using the standard "regal linting itself" benchmark on main before the
rule was merged and after shows a 10%+ increase in memory allocated as
well as wall clock time:

```
1111071791 ns/op	3001557280 B/op	58184205 allocs/op
1289814875 ns/op	3363166528 B/op	64788506 allocs/op - unresolved-reference merged
```

There are some design issues to address with aggregate rules that likely
would mean bigger wins across the board, but the focus of this PR is solely
to optimize the code for the rule itself.

While this remains an expensive rule, it's now closer to our other more costly rules.
In total, this saves about 2.7 million allocations of the 6.7 million allocs increased
with its addition. The steps in which this was done are roughly as follows:

```
1249527041 ns/op	3337275400 B/op	64261924 allocs/op - exclude input refs
1247609750 ns/op	3325728208 B/op	64061112 allocs/op - get rid of _import_aliases
1264200584 ns/op	3324900160 B/op	64037543 allocs/op - simplify _all_full_path_refs
1182358875 ns/op	3322581056 B/op	64013278 allocs/op - inline parts of _is_resolved_ref
1189268416 ns/op	3317109520 B/op	63942894 allocs/op - skip if name in ast.rule_and_function_names
1182895291 ns/op	3305238432 B/op	63739044 allocs/op - only aggregate data used in report
1179427375 ns/op	3273346464 B/op	63188538 allocs/op - custom result.location alternative
1161155834 ns/op	3210160768 B/op	62267845 allocs/op - defer to_location_object to report
1130275167 ns/op	3204268280 B/op	62150928 allocs/op - a few small tweaks
```

I have a few more ideas for improvements, but as they'd touch on more than just this rule,
this is probably a good place to stop for now.

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

Author: anderseknert

bundle/regal/ast/ast.rego

@@ -145,15 +145,15 @@ function_arg_names(rule) := [arg.value | some arg in rule.head.args]
 
 # METADATA
 # description: all the rule and function names in the input AST
-rule_and_function_names contains ref_to_string(rule.head.ref) if some rule in input.rules
+rule_and_function_names contains ref_static_to_string(rule.head.ref) if some rule in input.rules
 
 # METADATA
 # description: all identifiers in the input AST (rule and function names, plus imported names)
 identifiers := rule_and_function_names | imported_identifiers
 
 # METADATA
 # description: all rule names in the input AST (excluding functions)
-rule_names contains ref_to_string(rule.head.ref) if some rule in rules
+rule_names contains ref_static_to_string(rule.head.ref) if some rule in rules
 
 # METADATA
 # description: |
@@ -273,10 +273,7 @@ _format_part(part) := sprintf(".%s", [part.value]) if {
 #   non-static (i.e. variable) value, if any:
 #   foo.bar -> foo.bar
 #   foo.bar[baz] -> foo.bar
-ref_static_to_string(ref) := str if {
-	rs := ref_to_string(ref)
-	str := _trim_from_var(rs, regex.find_n(`\[[^"]`, rs, 1))
-}
+ref_static_to_string(ref) := _trim_from_var(rs, regex.find_n(`\[[^"]`, rs, 1)) if rs := ref_to_string(ref)
 
 _trim_from_var(ref_str, vars) := ref_str if {
 	count(vars) == 0

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

@@ -9,112 +9,119 @@ import data.regal.util
 
 # METADATA
 # description: collects exported and full of used refs from each module
-aggregate contains result.aggregate(rego.metadata.chain(), {
-	"exported_rules": object.keys(ast.rule_head_locations),
-	"expanded_refs": _all_full_path_refs,
-	"prefix_tree": _prefix_tree,
-})
-
-_import_aliases[alias_key] := string_value if {
-	some alias_key, value in ast.resolved_imports
-	string_value := concat(".", value)
+aggregate contains entry if {
+	exported := object.keys(ast.rule_head_locations)
+
+	entry := result.aggregate(rego.metadata.chain(), {
+		"exported_rules": exported,
+		"expanded_refs": _all_full_path_refs,
+		"prefix_tree": {prefix_path |
+			some rule_name in exported
+			rule_path := split(rule_name, ".")
+			some i in numbers.range(1, count(rule_path))
+			prefix_path := array.slice(rule_path, 0, i)
+		},
+	})
 }
 
 # an import is shadowed if it shares name with a rule
 _shadowed_imports contains rule_name if {
 	some rule_name in ast.rule_names
-	_import_aliases[rule_name]
+	ast.resolved_imports[rule_name]
 }
 
 # an import is shadowed if it shares name with a variable (or function argument)
 _shadowed_imports contains var_name if {
 	var_name := ast.found.vars[_][_][_].value
-	_import_aliases[var_name]
+	ast.resolved_imports[var_name]
 }
 
 _refs contains ref if {
-	term := ast.found.refs[_][_].value
-	name := ast.ref_static_to_string(term)
+	terms := ast.found.refs[_][_].value
+	terms[0].value != "input"
 
-	not name in ast.builtin_names
+	name := ast.ref_static_to_string(terms)
 
-	path := split(name, ".")
-	not path[0] in _shadowed_imports
+	not name in ast.builtin_names
+	not name in ast.rule_and_function_names
+	not terms[0].value in _shadowed_imports
 
-	ref := object.union(result.location(term), {
+	# util.to_location_row inlined for some extra performance
+	row := to_number(regex.replace(terms[0].location, `^(\d+):.*`, "$1"))
+	ref := {
 		"name": name,
-		"path": path,
-	})
+		"text": input.regal.file.lines[row - 1],
+		"location": terms[0].location,
+	}
 }
 
-_all_full_path_refs[ref.name] contains ref if {
+_all_full_path_refs[ref.name] contains [ref.location, ref.text] if {
 	some ref in _refs
-	ref.path[0] == "data"
+	startswith(ref.name, "data.")
 }
 
-_all_full_path_refs[expanded_ref] contains ref if {
+_all_full_path_refs[expanded_ref] contains [ref.location, ref.text] if {
 	some ref in _refs
-	full_source_prefix := _import_aliases[ref.path[0]]
-
-	full_path_array := array.concat([full_source_prefix], util.rest(ref.path))
-
-	expanded_ref := concat(".", full_path_array)
+	path := split(ref.name, ".")
+	expanded_ref := concat(".", array.concat(ast.resolved_imports[path[0]], util.rest(path)))
 }
 
 # METADATA
 # schemas:
 #   - input: schema.regal.aggregate
 aggregate_report contains violation if {
-	all_exports_in_bundle := {export |
-		some entry in input.aggregate
-		some export in entry.aggregate_data.exported_rules
-	}
-
-	prefix_tree := {prefix |
-		some entry in input.aggregate
-		some prefix in entry.aggregate_data.prefix_tree
-	}
+	all_exports := {export | export := input.aggregate[_].aggregate_data.exported_rules[_]}
+	prefix_tree := {prefix | prefix := input.aggregate[_].aggregate_data.prefix_tree[_]}
 
 	some entry in input.aggregate
-	some ref_full_name, ref_locations in entry.aggregate_data.expanded_refs
+	some name, refs in entry.aggregate_data.expanded_refs
 
-	# Ignore everything after the first "[" in the ref name. E.g. foo.bar[0].baz becomes foo.bar
-	simplified_ref_name := split(ref_full_name, "[")[0]
+	# ignore everything after the first "[" in the ref name. E.g. foo.bar[0].baz becomes foo.bar
+	ref_name := split(name, "[")[0]
+	ref_path := split(ref_name, ".")
 
-	not _is_resolved_ref(simplified_ref_name, prefix_tree, all_exports_in_bundle)
+	# a reference is considered resolved with respect to a rule if
+	# 1: it is the prefix of a rule
+	# 2: it indexes into a rule - we do not consider the possible data
+	# 3: the reference is ignored in the config
+	not ref_path in prefix_tree
+	not _is_resolved_ref(ref_path, all_exports)
+	not _is_excepted(ref_name)
 
-	some ref in ref_locations
-	violation := result.fail(rego.metadata.chain(), result.location(ref))
+	some ref in refs
+
+	violation := result.fail(rego.metadata.chain(), _to_location_object(ref[0], ref[1], entry.aggregate_source.file))
 }
 
-# METADATA
-# description: a reference is valid with respect to a rule if
-# custom:
-#   1: it is the prefix of a rule
-#   2: it indexes into a rule - we do not consider the possible data
-#   3: the reference is ignored in the config
-_is_resolved_ref(ref_full_name, _, _) if {
+_is_excepted(ref_full_name) if {
 	some exception in config.rules.imports["unresolved-reference"]["except-paths"]
 	glob.match(exception, [], ref_full_name)
 }
 
-_is_resolved_ref(ref_full_name, prefix_tree, _) if {
-	ref_full_path := split(ref_full_name, ".")
-	ref_full_path in prefix_tree
-}
-
-_is_resolved_ref(ref_full_name, _, all_exports_in_bundle) if {
-	ref_full_path := split(ref_full_name, ".")
-
-	some i in numbers.range(0, count(ref_full_path))
+_is_resolved_ref(ref_full_path, all_exports) if {
+	some i in numbers.range(1, count(ref_full_path))
 	path_prefix := concat(".", array.slice(ref_full_path, 0, i))
 
-	path_prefix in all_exports_in_bundle
+	path_prefix in all_exports
 }
 
-_prefix_tree contains prefix_path if {
-	some rule_name, _ in ast.rule_head_locations
-	rule_path := split(rule_name, ".")
-	some i in numbers.range(0, count(rule_path))
-	prefix_path := array.slice(rule_path, 0, i)
+# like util.to_location_object, but with text and file passed in
+# as we don't have access to the usual input.regal.file attributes
+# in the context of reporting aggregated data
+_to_location_object(loc, text, file) := {"location": {
+	"file": file,
+	"row": row,
+	"col": col,
+	"text": text,
+	"end": {
+		"row": end_row,
+		"col": end_col,
+	},
+}} if {
+	[r, c, er, ec] := split(loc, ":")
+
+	row := to_number(r)
+	col := to_number(c)
+	end_row := to_number(er)
+	end_col := to_number(ec)
 }

bundle/regal/util/util.rego

@@ -67,6 +67,10 @@ to_location_object(loc) := {
 
 to_location_object(loc) := loc if is_object(loc)
 
+# METADATA
+# description: efficiently extracts row number from location string
+to_location_row(loc) := to_number(regex.replace(loc, `^(\d+):.*`, "$1"))
+
 _location_to_text(row, col, end_row, end_col) := substring(
 	input.regal.file.lines[row - 1],
 	col - 1,

pkg/linter/linter_test.go

@@ -728,7 +728,7 @@ import data.unresolved`,
 	}
 }
 
-// 1111071791 ns/op	3001557280 B/op	58184205 allocs/op
+// 1130275167 ns/op	3204268280 B/op	62150928 allocs/op - a few small tweaks
 // ...
 func BenchmarkRegalLintingItself(b *testing.B) {
 	linter := NewLinter().
@@ -755,7 +755,7 @@ func BenchmarkRegalLintingItself(b *testing.B) {
 	}
 }
 
-// 1011992167 ns/op	2949576920 B/op	57079296 allocs/op
+// 1082386000 ns/op	3147603112 B/op	61016216 allocs/op
 // ...
 func BenchmarkRegalLintingItselfPrepareOnce(b *testing.B) {
 	ctx := context.Background()