From 2a0bd9f8af84696413341e724e79d416163cd753 Mon Sep 17 00:00:00 2001
From: Anders Eknert <anders@styra.com>
Date: Tue, 6 Aug 2024 11:59:45 +0200
Subject: [PATCH] perf: Walk less (#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                     | 15 ++--------
 bundle/regal/ast/ast_test.rego                | 21 -------------
 bundle/regal/ast/search.rego                  | 30 +++++++++++++++----
 .../rules/bugs/unused_output_variable.rego    |  2 +-
 .../regal/rules/bugs/var_shadows_builtin.rego |  2 +-
 .../regal/rules/custom/naming_convention.rego |  2 +-
 .../idiomatic/use_some_for_output_vars.rego   | 16 ++--------
 .../rules/idiomatic/use_strings_count.rego    |  2 +-
 .../regal/rules/style/default_over_else.rego  |  9 ++----
 .../regal/rules/style/prefer_snake_case.rego  |  2 +-
 .../regal/rules/style/unnecessary_some.rego   |  7 ++---
 .../rules/style/use_assignment_operator.rego  |  8 ++---
 .../rules/testing/metasyntactic_variable.rego |  2 +-
 mess.rego                                     |  7 -----
 14 files changed, 43 insertions(+), 82 deletions(-)
 delete mode 100644 mess.rego

diff --git a/bundle/regal/ast/ast.rego b/bundle/regal/ast/ast.rego
index bda864ce..89ed73b2 100644
--- a/bundle/regal/ast/ast.rego
+++ b/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
diff --git a/bundle/regal/ast/ast_test.rego b/bundle/regal/ast/ast_test.rego
index 734ec521..a330605b 100644
--- a/bundle/regal/ast/ast_test.rego
+++ b/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
 
diff --git a/bundle/regal/ast/search.rego b/bundle/regal/ast/search.rego
index 843b5fca..6199a781 100644
--- a/bundle/regal/ast/search.rego
+++ b/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,6 +193,26 @@ 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
@@ -202,7 +220,7 @@ vars[rule_index][context] contains var if {
 #   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)
 }
 
diff --git a/bundle/regal/rules/bugs/unused_output_variable.rego b/bundle/regal/rules/bugs/unused_output_variable.rego
index dfa0588c..1c21cc3d 100644
--- a/bundle/regal/rules/bugs/unused_output_variable.rego
+++ b/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, "$")
 }
diff --git a/bundle/regal/rules/bugs/var_shadows_builtin.rego b/bundle/regal/rules/bugs/var_shadows_builtin.rego
index 3637018c..bb547665 100644
--- a/bundle/regal/rules/bugs/var_shadows_builtin.rego
+++ b/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
 
diff --git a/bundle/regal/rules/custom/naming_convention.rego b/bundle/regal/rules/custom/naming_convention.rego
index 7730bef0..daba4522 100644
--- a/bundle/regal/rules/custom/naming_convention.rego
+++ b/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)
 
diff --git a/bundle/regal/rules/idiomatic/use_some_for_output_vars.rego b/bundle/regal/rules/idiomatic/use_some_for_output_vars.rego
index c6e32db4..e2af5cb0 100644
--- a/bundle/regal/rules/idiomatic/use_some_for_output_vars.rego
+++ b/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))
 }
diff --git a/bundle/regal/rules/idiomatic/use_strings_count.rego b/bundle/regal/rules/idiomatic/use_strings_count.rego
index 7a46ea6e..496c5e2e 100644
--- a/bundle/regal/rules/idiomatic/use_strings_count.rego
+++ b/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"
diff --git a/bundle/regal/rules/style/default_over_else.rego b/bundle/regal/rules/style/default_over_else.rego
index 1550217b..f2de6148 100644
--- a/bundle/regal/rules/style/default_over_else.rego
+++ b/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)
diff --git a/bundle/regal/rules/style/prefer_snake_case.rego b/bundle/regal/rules/style/prefer_snake_case.rego
index 502ccd33..362c029f 100644
--- a/bundle/regal/rules/style/prefer_snake_case.rego
+++ b/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))
diff --git a/bundle/regal/rules/style/unnecessary_some.rego b/bundle/regal/rules/style/unnecessary_some.rego
index 82345c33..55bf5c15 100644
--- a/bundle/regal/rules/style/unnecessary_some.rego
+++ b/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"
diff --git a/bundle/regal/rules/style/use_assignment_operator.rego b/bundle/regal/rules/style/use_assignment_operator.rego
index aa1af3fc..d1b5eae9 100644
--- a/bundle/regal/rules/style/use_assignment_operator.rego
+++ b/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)}}))
 }
diff --git a/bundle/regal/rules/testing/metasyntactic_variable.rego b/bundle/regal/rules/testing/metasyntactic_variable.rego
index 2005e52b..0afe483a 100644
--- a/bundle/regal/rules/testing/metasyntactic_variable.rego
+++ b/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
 
diff --git a/mess.rego b/mess.rego
deleted file mode 100644
index 9979dd7e..00000000
--- a/mess.rego
+++ /dev/null
@@ -1,7 +0,0 @@
-package p
-
-import rego.v1
-
-# comment
-
-allow := true