perf: Walk less (StyraInc#965)

- Use `ast.found.refs` instead of `walk` in `use-some-for-output-vars`
- Cheaper `walk` over `else` in `default-over-else` and `use-assignment-operator`
- Remove some unnecessary code in `external-reference`

The plan was also to consolidate `ast.found.symbols` with `ast.found.refs`, but
that had disastrous consequences leading to something I think is a bug in OPA.
I have demonstrated that to @charlieegan3, but will need to consult the OPA team
once they're back from vacation.

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

bundle/regal/ast/ast.rego

@@ -147,17 +147,6 @@ is_ref(value) if value.type == "ref"
 
 is_ref(value) if value[0].type == "ref"
 
-refs[rule_index] contains value if {
-	some i, rule in _rules
-
-	# converting to string until https://github.com/open-policy-agent/opa/issues/6736 is fixed
-	rule_index := sprintf("%d", [i])
-
-	walk(rule, [_, value])
-
-	is_ref(value)
-}
-
 # METADATA
 # description: |
 #   a map containing all function calls (built-in and custom) in the input AST
@@ -168,7 +157,7 @@ function_calls[rule_index] contains call if {
 	# converting to string until https://github.com/open-policy-agent/opa/issues/6736 is fixed
 	rule_index := sprintf("%d", [i])
 
-	some ref in refs[rule_index]
+	some ref in found.refs[rule_index]
 
 	name := ref_to_string(ref[0].value)
 	args := [arg |
@@ -191,7 +180,7 @@ _exclude_arg(_, _, arg) if arg.type == "call"
 # ignore here, as it's covered elsewhere
 _exclude_arg("assign", 0, _)
 
-all_rules_refs contains refs[_][_]
+all_rules_refs contains found.refs[_][_]
 
 # METADATA
 # description: set containing all references found in the input AST

bundle/regal/ast/ast_test.rego

@@ -248,27 +248,6 @@ test_find_names_in_scope if {
 	in_scope == {"bar", "global", "comp", "allow", "a", "b", "c", "d", "e"}
 }
 
-test_find_some_decl_vars if {
-	policy := `
-	package p
-
-	import rego.v1
-
-	allow if {
-		foo := 1
-		some x
-		input[x]
-		some y, z
-		input[y][z] == x
-	}`
-
-	module := regal.parse_module("p.rego", policy)
-
-	some_vars := ast.find_some_decl_vars(module.rules[0]) with input as module
-
-	var_names(some_vars) == {"x", "y", "z"}
-}
-
 test_find_some_decl_names_in_scope if {
 	policy := `package p
 

bundle/regal/ast/search.rego

@@ -153,13 +153,11 @@ _rule_index(rule) := sprintf("%d", [i]) if {
 	r == rule
 }
 
-find_some_decl_vars(rule) := [var | some var in vars[_rule_index(rule)]["some"]] # regal ignore:external-reference
-
 # METADATA
 # description: |
 #   traverses all nodes under provided node (using `walk`), and returns an array with
 #   all variables declared via assignment (:=), `some`, `every` and in comprehensions
-#   DEPRECATED: uses ast.vars instead
+#   DEPRECATED: uses ast.found.vars instead
 find_vars(node) := [var |
 	walk(node, [path, value])
 
@@ -183,7 +181,7 @@ _rules := data.workspace.parsed[input.regal.file.uri].rules if not input.rules
 #   - some
 #   - somein
 #   - ref
-vars[rule_index][context] contains var if {
+found.vars[rule_index][context] contains var if {
 	some i, rule in _rules
 
 	# converting to string until https://github.com/open-policy-agent/opa/issues/6736 is fixed
@@ -195,14 +193,34 @@ vars[rule_index][context] contains var if {
 	some var in vars
 }
 
+found.refs[rule_index] contains value if {
+	some i, rule in _rules
+
+	# converting to string until https://github.com/open-policy-agent/opa/issues/6736 is fixed
+	rule_index := sprintf("%d", [i])
+
+	walk(rule, [_, value])
+
+	is_ref(value)
+}
+
+found.symbols[rule_index] contains value.symbols if {
+	some i, rule in _rules
+
+	# converting to string until https://github.com/open-policy-agent/opa/issues/6736 is fixed
+	rule_index := sprintf("%d", [i])
+
+	walk(rule, [_, value])
+}
+
 # METADATA
 # description: |
 #   finds all vars declared in `rule` *before* the `location` provided
 #   note: this isn't 100% accurate, as it doesn't take into account `=`
 #   assignments / unification, but it's likely good enough since other rules
 #   recommend against those
 find_vars_in_local_scope(rule, location) := [var |
-	var := vars[_rule_index(rule)][_][_] # regal ignore:external-reference
+	var := found.vars[_rule_index(rule)][_][_] # regal ignore:external-reference
 
 	not startswith(var.value, "$")
 	_before_location(rule, var, location)
@@ -265,7 +283,7 @@ find_names_in_scope(rule, location) := names if {
 #   find all variables declared via `some` declarations (and *not* `some .. in`)
 #   in the scope of the given location
 find_some_decl_names_in_scope(rule, location) := {some_var.value |
-	some some_var in find_some_decl_vars(rule)
+	some some_var in found.vars[_rule_index(rule)]["some"] # regal ignore:external-reference
 	_before_location(rule, some_var, location)
 }
 

bundle/regal/rules/bugs/unused_output_variable.rego

@@ -42,7 +42,7 @@ report contains violation if {
 
 _ref_vars[rule_index][var.value] contains var if {
 	some rule_index
-	var := ast.vars[rule_index].ref[_]
+	var := ast.found.vars[rule_index].ref[_]
 
 	not startswith(var.value, "$")
 }

bundle/regal/rules/bugs/var_shadows_builtin.rego

@@ -8,7 +8,7 @@ import data.regal.ast
 import data.regal.result
 
 report contains violation if {
-	var := ast.vars[_][_][_]
+	var := ast.found.vars[_][_][_]
 
 	var.value in ast.builtin_namespaces
 

bundle/regal/rules/custom/naming_convention.rego

@@ -81,7 +81,7 @@ report contains violation if {
 
 	target in {"var", "variable"}
 
-	var := ast.vars[_][_][_]
+	var := ast.found.vars[_][_][_]
 
 	not regex.match(convention.pattern, var.value)
 

bundle/regal/rules/idiomatic/use_some_for_output_vars.rego

@@ -8,24 +8,14 @@ import data.regal.ast
 import data.regal.result
 
 report contains violation if {
-	# can't use ast.all_refs here as we need to
-	# refer to the `rule` below
-	some rule in input.rules
-
-	walk(rule, [_, value])
-
-	value.type == "ref"
-	ref := value.value
-
-	some i, elem in ref
+	some rule_index, i
+	elem := ast.found.refs[rule_index][_].value[i]
 
 	# first item can't be a loop or key ref
 	i != 0
 	elem.type == "var"
 	not startswith(elem.value, "$")
-
-	scope := ast.find_names_in_scope(rule, elem.location)
-	not elem.value in scope
+	not elem.value in ast.find_names_in_scope(input.rules[to_number(rule_index)], elem.location)
 
 	violation := result.fail(rego.metadata.chain(), result.location(elem))
 }

bundle/regal/rules/idiomatic/use_strings_count.rego

@@ -19,7 +19,7 @@ notices contains result.notice(rego.metadata.chain()) if not capabilities.has_ob
 report contains violation if {
 	some rule in input.rules
 
-	ref := ast.refs[_][_]
+	ref := ast.found.refs[_][_]
 
 	ref[0].value[0].type == "var"
 	ref[0].value[0].value == "count"

bundle/regal/rules/style/default_over_else.rego

@@ -16,13 +16,10 @@ report contains violation if {
 	# walking is expensive but necessary here, since there could be
 	# any number of `else` clauses nested below. no need to traverse
 	# the rule if there isn't a single `else` present though!
-	rule["else"]
+	walk(rule["else"], [_, value])
 
-	walk(rule, [_, value])
-
-	# quoting is needed as `else` is a keyword
-	else_body := value["else"].body
-	else_head := value["else"].head
+	else_body := value.body
+	else_head := value.head
 
 	# we don't know for sure, but if all that's in the body is an empty
 	# `true`, it's likely an implicit body (i.e. one not printed)

bundle/regal/rules/style/prefer_snake_case.rego

@@ -17,7 +17,7 @@ report contains violation if {
 }
 
 report contains violation if {
-	var := ast.vars[_][_][_]
+	var := ast.found.vars[_][_][_]
 	not util.is_snake_case(var.value)
 
 	violation := result.fail(rego.metadata.chain(), result.ranged_location_from_text(var))

bundle/regal/rules/style/unnecessary_some.rego

@@ -11,11 +11,8 @@ report contains violation if {
 	# No need to traverse rules here if we're not importing `in`
 	ast.imports_keyword(input.imports, "in")
 
-	some rule in input.rules
-
-	walk(rule, [_, value])
-
-	symbols := value.symbols
+	some rule_index
+	symbols := ast.found.symbols[rule_index][_]
 
 	symbols[0].type == "call"
 	symbols[0].value[0].type == "ref"

bundle/regal/rules/style/use_assignment_operator.rego

@@ -41,21 +41,19 @@ report contains violation if {
 	# walking is expensive but necessary here, since there could be
 	# any number of `else` clauses nested below. no need to traverse
 	# the rule if there isn't a single `else` present though!
-	rule["else"]
 
 	# NOTE: the same logic is used in default-over-else
 	# we should consider having a helper function to return
 	# all else clauses, for a given rule, as potentially that
 	# would be cached on the second invocation of the function
-	walk(rule, [_, value])
-	value["else"]
+	walk(rule["else"], [_, value])
 
 	# extract the text from location to see if '=' is used for
 	# assignment
-	text := base64.decode(value["else"].head.location.text)
+	text := base64.decode(value.head.location.text)
 	regex.match(`^else\s*=`, text)
 
-	loc := result.location(value["else"].head)
+	loc := result.location(value.head)
 
 	violation := result.fail(rego.metadata.chain(), object.union(loc, {"location": {"col": eq_col(loc)}}))
 }

bundle/regal/rules/testing/metasyntactic_variable.rego

@@ -43,7 +43,7 @@ report contains violation if {
 
 report contains violation if {
 	some i
-	var := ast.vars[i][_][_]
+	var := ast.found.vars[i][_][_]
 
 	lower(var.value) in metasyntactic