From 35eafec33c40e85b922ce6edb4e0424dd2b0320f Mon Sep 17 00:00:00 2001
From: Anders Eknert <anders@styra.com>
Date: Fri, 25 Apr 2025 01:00:22 +0200
Subject: [PATCH] Scan nested expressions

A few of the first Regal rules only ever checked the top level of
a body, and would not traverse into e.g. comprehensions or every
bodies. This has been annoying for long, and while doing this
properly is certainly more expensive (+2 million allocs) I think
we can afford that at this point in order to prioritize correctness.
May very well be that many will have Regal find "new" issues in the
next version due to this.

Also:
- We now also separate ast.found.references and ast.found.calls. This
  too comes at a cost, but makes it much more obvious what our code is
  doing, and we can now afforc prioritizing that too.
- Improve `narrow-argument` rule to also take array indices into account
  for narrowing (i.e. only `arr[1].value[0]` referenced)
- Fix new golangci-lint rule that says public functions should be above
  private ones.

Fixes #82

Signed-off-by: Anders Eknert <anders@styra.com>
---
 .regal/config.yaml                            |   1 +
 bundle/regal/ast/ast.rego                     |  27 +--
 bundle/regal/ast/ast_test.rego                | 156 -----------------
 bundle/regal/ast/search.rego                  |  66 +++----
 bundle/regal/ast/search_test.rego             | 163 ++++++++++++++++--
 bundle/regal/lsp/completion/ref_names.rego    |   6 +
 .../constant_condition.rego                   |  13 +-
 .../constant_condition_test.rego              |  80 ++++++++-
 .../deprecated_builtin.rego                   |   3 +-
 .../if_object_literal_test.rego               |  20 ++-
 .../bugs/impossible-not/impossible_not.rego   |   2 +-
 .../impossible-not/impossible_not_test.rego   |  25 +++
 .../not_equals_in_loop.rego                   |   5 +-
 .../redundant_existence_check.rego            |  10 +-
 .../sprintf_arguments_mismatch.rego           |   4 +-
 .../sprintf_arguments_mismatch_test.rego      |  42 ++++-
 .../unassigned_return_value.rego              |   2 +-
 .../unassigned_return_value_test.rego         |  31 ++++
 .../unused_output_variable.rego               |  10 +-
 .../forbidden_function_call.rego              |   2 +-
 .../narrow-argument/narrow_argument.rego      |  19 +-
 .../narrow-argument/narrow_argument_test.rego |  37 ++++
 .../non_raw_regex_pattern.rego                |   2 +-
 .../prefer_set_or_object_rule.rego            |   2 +-
 .../use-strings-count/use_strings_count.rego  |   2 +-
 .../prefer_package_imports.rego               |   5 +-
 .../unresolved-import/unresolved_import.rego  |  15 +-
 .../defer-assignment/defer_assignment.rego    |   5 +-
 .../with_outside_test_context.rego            |   4 +-
 .../comprehension_term_assignment.rego        |   4 +-
 .../double-negative/double_negative.rego      |   5 +-
 .../pointless_reassignment.rego               |   2 +-
 .../rules/style/rule-length/rule_length.rego  |   4 +-
 .../unconditional_assignment.rego             |  13 +-
 .../unnecessary-some/unnecessary_some.rego    |  24 ++-
 .../style/yoda-condition/yoda_condition.rego  |   2 +-
 .../dubious_print_sprintf.rego                |   2 +-
 .../print_or_trace_call.rego                  |   2 +-
 internal/embeds/schemas/regal-ast.json        |   3 -
 internal/lsp/config/watcher.go                |  54 +++---
 internal/lsp/server.go                        | 158 ++++++++---------
 pkg/fixer/fixer.go                            |  16 +-
 42 files changed, 615 insertions(+), 433 deletions(-)

diff --git a/.regal/config.yaml b/.regal/config.yaml
index 45bb4081..406999be 100644
--- a/.regal/config.yaml
+++ b/.regal/config.yaml
@@ -17,6 +17,7 @@ rules:
       level: error
       exclude-args:
         - cfg
+        - metadata
         - rule
   style:
     line-length:
diff --git a/bundle/regal/ast/ast.rego b/bundle/regal/ast/ast.rego
index 75abd0b7..da0b1aea 100644
--- a/bundle/regal/ast/ast.rego
+++ b/bundle/regal/ast/ast.rego
@@ -203,7 +203,7 @@ rule_index_strings := [sprintf("%d", [i]) | some i, _ in _rules]
 #   keyed by rule index
 function_calls[rule_index] contains call if {
 	some rule_index in rule_index_strings
-	some ref in found.refs[rule_index]
+	some ref in found.calls[rule_index]
 
 	name := ref_to_string(ref[0].value)
 	args := [arg |
@@ -344,21 +344,6 @@ all_functions := object.union(config.capabilities.builtins, function_decls(input
 #   scope of the input AST
 all_function_names := object.keys(all_functions)
 
-# METADATA
-# description: set containing all negated expressions in input AST
-negated_expressions[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 := rule_index_strings[i]
-
-	some node in ["head", "body", "else"]
-
-	walk(rule[node], [_, value])
-
-	value.negated == true
-}
-
 # METADATA
 # description: |
 #   true if rule head contains no identifier, but is a chained rule body immediately following the previous one:
@@ -398,15 +383,15 @@ var_in_call(calls, rule_index, name) if has_named_var(calls[rule_index][_].args[
 
 # METADATA
 # description: answers wether provided expression is an assignment (using `:=`)
-is_assignment(expr) if {
-	expr.terms[0].type == "ref"
-	expr.terms[0].value[0].type == "var"
-	expr.terms[0].value[0].value == "assign"
+is_assignment(term) if {
+	term.type == "ref"
+	term.value[0].type == "var"
+	term.value[0].value == "assign"
 }
 
 # METADATA
 # description: returns the terms in an assignment (`:=`) expression, or undefined if not assignment
-assignment_terms(expr) := [expr.terms[1], expr.terms[2]] if is_assignment(expr)
+assignment_terms(terms) := [terms[1], terms[2]] if is_assignment(terms[0])
 
 # METADATA
 # description: |
diff --git a/bundle/regal/ast/ast_test.rego b/bundle/regal/ast/ast_test.rego
index 974a5f9c..ba71636e 100644
--- a/bundle/regal/ast/ast_test.rego
+++ b/bundle/regal/ast/ast_test.rego
@@ -4,89 +4,6 @@ import data.regal.ast
 import data.regal.capabilities
 import data.regal.config
 
-test_find_vars if {
-	policy := `
-	package p
-
-	global := "foo"
-
-	allow if {
-		a := global
-		b := [c | c := input[d]]
-
-		every e in input {
-			e == "foo"
-		}
-
-		every f, g in input.bar {
-			f == g
-		}
-
-		some h, i
-		input.bar[h][i]
-		some j in input
-		some k, l in input
-
-		[m, n, o] := [1, 2, 3]
-
-		[p, [q, _]] := [1, [2, 1]]
-
-		some _, [r, s] in [["foo", "bar"], [1, 2]]
-
-		{"x": t} := {"x": 1}
-
-		some [u] in [[1]]
-	}
-	`
-
-	vars := ast.find_vars(regal.parse_module("p.rego", policy).rules) with config.capabilities as capabilities.provided
-	names := {var.value |
-		some var in vars
-		var.type == "var"
-	}
-
-	names == {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u"}
-}
-
-test_find_vars_comprehension_lhs if {
-	policy := `
-	package p
-
-	allow if {
-		a := [b | input[b]]
-		c := {d | input[d]}
-		e := {f: g | g := input[f]}
-	}
-	`
-
-	vars := ast.find_vars(regal.parse_module("p.rego", policy).rules) with config.capabilities as capabilities.provided
-	names := {var.value |
-		some var in vars
-		var.type == "var"
-	}
-
-	names == {"a", "b", "c", "d", "e", "f", "g"}
-}
-
-test_find_vars_function_ret_return_args if {
-	policy := `
-	package p
-
-	allow if {
-		walk(input, [path, value])
-	}
-	`
-
-	module := regal.parse_module("p.rego", policy)
-	vars := ast.find_vars(module.rules) with config.capabilities as capabilities.provided with input.rules as []
-	names := {var.value |
-		some var in vars
-		var.type == "var"
-	}
-
-	names == {"path", "value"}
-}
-
 # https://github.com/StyraInc/regal/issues/168
 test_function_decls_multiple_same_name if {
 	policy := `package p
@@ -132,62 +49,6 @@ allow := true
 	]
 }
 
-test_find_vars_in_local_scope if {
-	policy := `
-	package p
-
-	global := "foo"
-
-	allow if {
-		a := global
-		b := [c | c := input[d]]
-
-		every e in input {
-			f == "foo"
-			g := "bar"
-			h == "foo"
-		}
-	}`
-
-	module := regal.parse_module("p.rego", policy)
-
-	allow_rule := module.rules[1]
-
-	var_locations := {
-		"a": {"col": 3, "row": 9},
-		"b": {"col": 3, "row": 10},
-		"c": {"col": 13, "row": 10},
-		"d": {"col": 9, "row": 12},
-		"e": {"col": 4, "row": 14},
-	}
-
-	var_names(ast.find_vars_in_local_scope(allow_rule, var_locations.a)) with input as module == set()
-	var_names(ast.find_vars_in_local_scope(allow_rule, var_locations.b)) with input as module == {"a"}
-	var_names(ast.find_vars_in_local_scope(allow_rule, var_locations.c)) with input as module == {"a", "b", "c"}
-	var_names(ast.find_vars_in_local_scope(allow_rule, var_locations.d)) with input as module == {"a", "b", "c", "d"}
-	var_names(ast.find_vars_in_local_scope(allow_rule, var_locations.e)) with input as module == {"a", "b", "c", "d", "e"}
-}
-
-test_find_vars_in_local_scope_complex_comprehension_term if {
-	policy := `
-	package p
-
-	allow if {
-		a := [{"b": b} | c := input[b]]
-	}`
-
-	module := regal.parse_module("p.rego", policy)
-
-	allow_rule := module.rules[0]
-
-	ast.find_vars_in_local_scope(allow_rule, {"col": 10, "row": 10}) with input as module == [
-		{"location": {"col": 3, "row": 7, "text": "YQ=="}, "type": "var", "value": "a"},
-		{"location": {"col": 15, "row": 7, "text": "Yg=="}, "type": "var", "value": "b"},
-		{"location": {"col": 20, "row": 7, "text": "Yw=="}, "type": "var", "value": "c"},
-		{"location": {"col": 31, "row": 7, "text": "Yg=="}, "type": "var", "value": "b"},
-	]
-}
-
 test_find_names_in_scope if {
 	policy := `
 	package p
@@ -218,23 +79,6 @@ test_find_names_in_scope if {
 	in_scope == {"bar", "global", "comp", "allow", "a", "b", "c", "d", "e"}
 }
 
-test_find_some_decl_names_in_scope if {
-	policy := `package p
-
-	allow if {
-		foo := 1
-		some x
-		input[x]
-		some y, z
-		input[y][z] == x
-	}`
-
-	module := regal.parse_module("p.rego", policy)
-
-	{"x"} == ast.find_some_decl_names_in_scope(module.rules[0], {"col": 1, "row": 6}) with input as module
-	{"x", "y", "z"} == ast.find_some_decl_names_in_scope(module.rules[0], {"col": 1, "row": 8}) with input as module
-}
-
 var_names(vars) := {var.value | some var in vars}
 
 test_provided_capabilities_never_undefined if capabilities.provided == {} with data.internal as {}
diff --git a/bundle/regal/ast/search.rego b/bundle/regal/ast/search.rego
index fcc773ca..a37965a0 100644
--- a/bundle/regal/ast/search.rego
+++ b/bundle/regal/ast/search.rego
@@ -10,19 +10,13 @@ _find_nested_vars(obj) := [value |
 
 # simple assignment, i.e. `x := 100` returns `x`
 # always returns a single var, but wrapped in an
-# array for consistency
-_find_assign_vars(value) := var if {
-	value[1].type == "var"
-	var := [value[1]]
-}
-
+# array for consistency or
 # 'destructuring' array assignment, i.e.
-# [a, b, c] := [1, 2, 3]
-# or
-# {a: b} := {"foo": "bar"}
-_find_assign_vars(value) := vars if {
-	value[1].type in {"array", "object"}
-	vars := _find_nested_vars(value[1])
+# [a, b, c] := [1, 2, 3] or {a: b} := {"foo": "bar"}
+_find_assign_vars(value) := [value] if {
+	value.type == "var"
+} else := _find_nested_vars(value) if {
+	value.type in {"array", "object"}
 }
 
 # var declared via `some`, i.e. `some x` or `some x, y`
@@ -32,16 +26,10 @@ _find_some_decl_vars(value) := [v |
 ]
 
 # single var declared via `some in`, i.e. `some x in y`
-_find_some_in_decl_vars(value) := vars if {
-	arr := value[0].value
-	count(arr) == 3
-
-	vars := _find_nested_vars(arr[1])
-}
+_find_some_in_decl_vars(arr) := _find_nested_vars(arr[1]) if count(arr) == 3
 
 # two vars declared via `some in`, i.e. `some x, y in z`
-_find_some_in_decl_vars(value) := vars if {
-	arr := value[0].value
+_find_some_in_decl_vars(arr) := vars if {
 	count(arr) == 4
 
 	vars := [v |
@@ -128,14 +116,14 @@ _find_vars(value, last) := {"term": find_term_vars(function_ret_args(fn_name, va
 # left-hand side is equally dubious, but we'll treat `x = 1` as `x := 1` for
 # the purpose of this function until we have a more robust way of dealing with
 # unification
-_find_vars(value, last) := {"assign": _find_assign_vars(value)} if {
+_find_vars(value, last) := {"assign": _find_assign_vars(value[1])} if {
 	last == "terms"
 	value[0].type == "ref"
 	value[0].value[0].type == "var"
 	value[0].value[0].value in {"assign", "eq"}
 }
 
-_find_vars(value, last) := {"somein": _find_some_in_decl_vars(value)} if {
+_find_vars(value, last) := {"somein": _find_some_in_decl_vars(value[0].value)} if {
 	last == "symbols"
 	value[0].type == "call"
 }
@@ -208,10 +196,13 @@ _rules := data.workspace.parsed[input.regal.file.uri].rules if not input.rules
 #   - somein
 #   - ref
 found.vars[rule_index][context] contains var if {
-	some i, rule_index in rule_index_strings
+	some i, rule in _rules
+
+	rule_index := rule_index_strings[i]
+
 	some node in ["head", "body", "else"]
 
-	walk(_rules[i][node], [path, value])
+	walk(rule[node], [path, value])
 
 	last := regal.last(path)
 	last in {"terms", "symbols", "args"}
@@ -233,7 +224,6 @@ found.vars[rule_index].ref contains var if {
 
 # METADATA
 # description: all refs found in module
-# scope: document
 found.refs[rule_index] contains value if {
 	some i, rule in _rules
 
@@ -245,7 +235,9 @@ found.refs[rule_index] contains value if {
 	value.type == "ref"
 }
 
-found.refs[rule_index] contains value if {
+# METADATA
+# description: all calls found in module
+found.calls[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
@@ -280,6 +272,21 @@ found.comprehensions[rule_index] contains value if {
 	value.type in {"arraycomprehension", "objectcomprehension", "setcomprehension"}
 }
 
+# METADATA
+# description: set containing all negated expressions in input AST
+found.expressions[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 := rule_index_strings[i]
+
+	some node in ["head", "body", "else"]
+
+	walk(rule[node], [_, value])
+
+	value.terms
+}
+
 # METADATA
 # description: |
 #   finds all vars declared in `rule` *before* the `location` provided
@@ -363,10 +370,3 @@ find_some_decl_names_in_scope(rule, location) := {some_var.value |
 	some some_var in found.vars[_rule_index(rule)]["some"]
 	_before_location(rule.head, some_var, location)
 }
-
-# METADATA
-# description: all expressions in module
-exprs[rule_index][expr_index] := expr if {
-	some rule_index, rule in input.rules
-	some expr_index, expr in rule.body
-}
diff --git a/bundle/regal/ast/search_test.rego b/bundle/regal/ast/search_test.rego
index f64a65ad..4e8a306e 100644
--- a/bundle/regal/ast/search_test.rego
+++ b/bundle/regal/ast/search_test.rego
@@ -1,28 +1,161 @@
 package regal.ast_test
 
 import data.regal.ast
+import data.regal.capabilities
+import data.regal.config
 
-test_exprs if {
-	inp := regal.parse_module("foo.rego", `package example
+test_find_vars if {
+	policy := `
+	package p
 
-import rego.v1
+	global := "foo"
 
-allow if input.x == 1
+	allow if {
+		a := global
+		b := [c | c := input[d]]
 
-allow if {
-	input.y == 2
-	input.z == 3
+		every e in input {
+			e == "foo"
+		}
+
+		every f, g in input.bar {
+			f == g
+		}
+
+		some h, i
+		input.bar[h][i]
+		some j in input
+		some k, l in input
+
+		[m, n, o] := [1, 2, 3]
+
+		[p, [q, _]] := [1, [2, 1]]
+
+		some _, [r, s] in [["foo", "bar"], [1, 2]]
+
+		{"x": t} := {"x": 1}
+
+		some [u] in [[1]]
+	}
+	`
+
+	vars := ast.find_vars(regal.parse_module("p.rego", policy).rules) with config.capabilities as capabilities.provided
+	names := {var.value |
+		some var in vars
+		var.type == "var"
+	}
+
+	names == {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u"}
+}
+
+test_find_vars_comprehension_lhs if {
+	policy := `
+	package p
+
+	allow if {
+		a := [b | input[b]]
+		c := {d | input[d]}
+		e := {f: g | g := input[f]}
+	}
+	`
+
+	vars := ast.find_vars(regal.parse_module("p.rego", policy).rules) with config.capabilities as capabilities.provided
+	names := {var.value |
+		some var in vars
+		var.type == "var"
+	}
+
+	names == {"a", "b", "c", "d", "e", "f", "g"}
+}
+
+test_find_vars_function_ret_return_args if {
+	policy := `
+	package p
+
+	allow if {
+		walk(input, [path, value])
+	}
+	`
+
+	module := regal.parse_module("p.rego", policy)
+	vars := ast.find_vars(module.rules) with config.capabilities as capabilities.provided with input.rules as []
+	names := {var.value |
+		some var in vars
+		var.type == "var"
+	}
+
+	names == {"path", "value"}
 }
-`)
 
-	result := ast.exprs with input as inp
+test_find_some_decl_names_in_scope if {
+	policy := `package p
+
+	allow if {
+		foo := 1
+		some x
+		input[x]
+		some y, z
+		input[y][z] == x
+	}`
+
+	module := regal.parse_module("p.rego", policy)
+
+	{"x"} == ast.find_some_decl_names_in_scope(module.rules[0], {"col": 1, "row": 6}) with input as module
+	{"x", "y", "z"} == ast.find_some_decl_names_in_scope(module.rules[0], {"col": 1, "row": 8}) with input as module
+}
+
+test_find_vars_in_local_scope if {
+	policy := `
+	package p
+
+	global := "foo"
+
+	allow if {
+		a := global
+		b := [c | c := input[d]]
+
+		every e in input {
+			f == "foo"
+			g := "bar"
+			h == "foo"
+		}
+	}`
+
+	module := regal.parse_module("p.rego", policy)
+
+	allow_rule := module.rules[1]
+
+	var_locations := {
+		"a": {"col": 3, "row": 9},
+		"b": {"col": 3, "row": 10},
+		"c": {"col": 13, "row": 10},
+		"d": {"col": 9, "row": 12},
+		"e": {"col": 4, "row": 14},
+	}
+
+	var_names(ast.find_vars_in_local_scope(allow_rule, var_locations.a)) with input as module == set()
+	var_names(ast.find_vars_in_local_scope(allow_rule, var_locations.b)) with input as module == {"a"}
+	var_names(ast.find_vars_in_local_scope(allow_rule, var_locations.c)) with input as module == {"a", "b", "c"}
+	var_names(ast.find_vars_in_local_scope(allow_rule, var_locations.d)) with input as module == {"a", "b", "c", "d"}
+	var_names(ast.find_vars_in_local_scope(allow_rule, var_locations.e)) with input as module == {"a", "b", "c", "d", "e"}
+}
+
+test_find_vars_in_local_scope_complex_comprehension_term if {
+	policy := `
+	package p
+
+	allow if {
+		a := [{"b": b} | c := input[b]]
+	}`
 
-	count(result) == 2 # rules
+	module := regal.parse_module("p.rego", policy)
 
-	count(result[0]) == 1
-	result[0][0].terms[0].value[0].value == "equal"
+	allow_rule := module.rules[0]
 
-	count(result[1]) == 2
-	result[1][0].terms[0].value[0].value == "equal"
-	result[1][1].terms[0].value[0].value == "equal"
+	ast.find_vars_in_local_scope(allow_rule, {"col": 10, "row": 10}) with input as module == [
+		{"location": {"col": 3, "row": 7, "text": "YQ=="}, "type": "var", "value": "a"},
+		{"location": {"col": 15, "row": 7, "text": "Yg=="}, "type": "var", "value": "b"},
+		{"location": {"col": 20, "row": 7, "text": "Yw=="}, "type": "var", "value": "c"},
+		{"location": {"col": 31, "row": 7, "text": "Yg=="}, "type": "var", "value": "b"},
+	]
 }
diff --git a/bundle/regal/lsp/completion/ref_names.rego b/bundle/regal/lsp/completion/ref_names.rego
index d0754698..bdbbd548 100644
--- a/bundle/regal/lsp/completion/ref_names.rego
+++ b/bundle/regal/lsp/completion/ref_names.rego
@@ -7,6 +7,12 @@ import data.regal.ast
 #   returns a list of ref names that are used in the module
 #   built-in functions are not included as they are provided by another completions provider
 # scope: document
+ref_names contains name if {
+	name := ast.ref_static_to_string(ast.found.calls[_][_].value)
+
+	not name in ast.builtin_names
+}
+
 ref_names contains name if {
 	name := ast.ref_static_to_string(ast.found.refs[_][_].value)
 
diff --git a/bundle/regal/rules/bugs/constant-condition/constant_condition.rego b/bundle/regal/rules/bugs/constant-condition/constant_condition.rego
index 42bba7ba..5e667f28 100644
--- a/bundle/regal/rules/bugs/constant-condition/constant_condition.rego
+++ b/bundle/regal/rules/bugs/constant-condition/constant_condition.rego
@@ -5,20 +5,13 @@ package regal.rules.bugs["constant-condition"]
 import data.regal.ast
 import data.regal.result
 
-# NOTE: The constant condition checks currently don't do nesting!
-# Additionally, there are several other conditions that could be considered
-# constant, or if not, redundant... so this rule should be expanded in time
-
 # METADATA
 # description: single scalar value, like a lone `true` inside a rule body
 # scope: rule
 report contains violation if {
-	terms := input.rules[_].body[_].terms
+	terms := ast.found.expressions[_][_].terms
 
-	# We could probably include arrays and objects too, as a single compound value
-	# is not very useful, but it's not as clear cut as scalars, as you could have
-	# something like {"a": foo(input.x) == "bar"} which is not a constant condition,
-	# however meaningless it may be. Maybe consider for another rule?
+	# We could include composite types too, but less comomon and more expensive to check
 	terms.type in ast.scalar_types
 
 	violation := result.fail(rego.metadata.chain(), result.location(terms))
@@ -30,7 +23,7 @@ report contains violation if {
 report contains violation if {
 	operators := {"equal", "gt", "gte", "lt", "lte", "neq"}
 
-	expr := input.rules[_].body[_]
+	expr := ast.found.expressions[_][_]
 
 	expr.terms[0].value[0].type == "var"
 	expr.terms[0].value[0].value in operators
diff --git a/bundle/regal/rules/bugs/constant-condition/constant_condition_test.rego b/bundle/regal/rules/bugs/constant-condition/constant_condition_test.rego
index f09e8a1c..52c4ba94 100644
--- a/bundle/regal/rules/bugs/constant-condition/constant_condition_test.rego
+++ b/bundle/regal/rules/bugs/constant-condition/constant_condition_test.rego
@@ -12,7 +12,45 @@ test_fail_simple_constant_condition if {
 	r == {{
 		"category": "bugs",
 		"description": "Constant condition",
-		"location": {"col": 2, "file": "policy.rego", "row": 4, "text": "\t1", "end": {"row": 4, "col": 3}},
+		"location": {
+			"col": 2,
+			"file": "policy.rego",
+			"row": 4,
+			"text": "\t1",
+			"end": {
+				"row": 4,
+				"col": 3,
+			},
+		},
+		"related_resources": [{
+			"description": "documentation",
+			"ref": config.docs.resolve_url("$baseUrl/$category/constant-condition", "bugs"),
+		}],
+		"title": "constant-condition",
+		"level": "error",
+	}}
+}
+
+test_fail_simple_constant_condition_nested if {
+	r := rule.report with input as ast.policy(`allow if {
+		every x in [1, 2] {
+			1
+			x == 2
+		}
+	}`)
+
+	r == {{
+		"category": "bugs",
+		"description": "Constant condition",
+		"location": {
+			"col": 4,
+			"end": {
+				"col": 5,
+				"row": 5,
+			},
+			"file": "policy.rego",
+			"row": 5, "text": "\t\t\t1",
+		},
 		"related_resources": [{
 			"description": "documentation",
 			"ref": config.docs.resolve_url("$baseUrl/$category/constant-condition", "bugs"),
@@ -55,7 +93,45 @@ test_fail_operator_constant_condition if {
 	r == {{
 		"category": "bugs",
 		"description": "Constant condition",
-		"location": {"col": 2, "file": "policy.rego", "row": 4, "text": "\t1 == 1", "end": {"col": 8, "row": 4}},
+		"location": {
+			"col": 2,
+			"file": "policy.rego",
+			"row": 4,
+			"text": "\t1 == 1",
+			"end": {
+				"col": 8,
+				"row": 4,
+			},
+		},
+		"related_resources": [{
+			"description": "documentation",
+			"ref": config.docs.resolve_url("$baseUrl/$category/constant-condition", "bugs"),
+		}],
+		"title": "constant-condition",
+		"level": "error",
+	}}
+}
+
+test_fail_operator_constant_condition_nested if {
+	r := rule.report with input as ast.policy(`nested := [1 |
+		c := [2 |
+			1 == 1
+		]
+	]`)
+
+	r == {{
+		"category": "bugs",
+		"description": "Constant condition",
+		"location": {
+			"col": 4,
+			"end": {
+				"col": 10,
+				"row": 5,
+			},
+			"file": "policy.rego",
+			"row": 5,
+			"text": "\t\t\t1 == 1",
+		},
 		"related_resources": [{
 			"description": "documentation",
 			"ref": config.docs.resolve_url("$baseUrl/$category/constant-condition", "bugs"),
diff --git a/bundle/regal/rules/bugs/deprecated-builtin/deprecated_builtin.rego b/bundle/regal/rules/bugs/deprecated-builtin/deprecated_builtin.rego
index 415b0e42..0461f47e 100644
--- a/bundle/regal/rules/bugs/deprecated-builtin/deprecated_builtin.rego
+++ b/bundle/regal/rules/bugs/deprecated-builtin/deprecated_builtin.rego
@@ -29,8 +29,7 @@ report contains violation if {
 	# bail out early if no the deprecated built-ins are in capabilities
 	util.intersects(object.keys(config.capabilities.builtins), deprecated_builtins)
 
-	ref := ast.found.refs[_][_]
-	call := ref[0]
+	call := ast.found.calls[_][_][0]
 
 	ast.ref_to_string(call.value) in deprecated_builtins
 
diff --git a/bundle/regal/rules/bugs/if-object-literal/if_object_literal_test.rego b/bundle/regal/rules/bugs/if-object-literal/if_object_literal_test.rego
index 57aeb663..012f2881 100644
--- a/bundle/regal/rules/bugs/if-object-literal/if_object_literal_test.rego
+++ b/bundle/regal/rules/bugs/if-object-literal/if_object_literal_test.rego
@@ -9,11 +9,27 @@ test_fail[name] if {
 	some name, [policy, location] in {
 		"empty_object": [
 			`rule if {}`,
-			{"col": 9, "row": 5, "text": "rule if {}", "end": {"col": 11, "row": 5}},
+			{
+				"col": 9,
+				"row": 5,
+				"text": "rule if {}",
+				"end": {
+					"col": 11,
+					"row": 5,
+				},
+			},
 		],
 		"non_empty_object": [
 			`rule if {"x": input.x}`,
-			{"col": 9, "row": 5, "text": `rule if {"x": input.x}`, "end": {"col": 23, "row": 5}},
+			{
+				"col": 9,
+				"row": 5,
+				"text": `rule if {"x": input.x}`,
+				"end": {
+					"col": 23,
+					"row": 5,
+				},
+			},
 		],
 	}
 
diff --git a/bundle/regal/rules/bugs/impossible-not/impossible_not.rego b/bundle/regal/rules/bugs/impossible-not/impossible_not.rego
index f8e9ef26..0eb6a31f 100644
--- a/bundle/regal/rules/bugs/impossible-not/impossible_not.rego
+++ b/bundle/regal/rules/bugs/impossible-not/impossible_not.rego
@@ -25,7 +25,7 @@ _multivalue_rules contains path if {
 
 _negated_refs contains negated_ref if {
 	some rule_index, value
-	ast.negated_expressions[rule_index][value]
+	ast.found.expressions[rule_index][value].negated
 
 	# if terms is an array, it's a function call, and most likely not "impossible"
 	is_object(value.terms)
diff --git a/bundle/regal/rules/bugs/impossible-not/impossible_not_test.rego b/bundle/regal/rules/bugs/impossible-not/impossible_not_test.rego
index 04620420..8ceb78df 100644
--- a/bundle/regal/rules/bugs/impossible-not/impossible_not_test.rego
+++ b/bundle/regal/rules/bugs/impossible-not/impossible_not_test.rego
@@ -34,6 +34,31 @@ test_fail_multivalue_not_reference_same_package if {
 	})
 }
 
+test_fail_multivalue_not_reference_same_package_nested_expression if {
+	agg1 := rule.aggregate with input as regal.parse_module("p1.rego", `package foo
+
+	partial contains "foo"`)
+
+	agg2 := rule.aggregate with input as regal.parse_module("p2.rego", `package foo
+
+	test_foo if {
+		comprehension := [1 | not partial]
+	}`)
+
+	r := rule.aggregate_report with input as {"aggregate": (agg1 | agg2)}
+
+	r == expected_with_location({
+		"col": 29,
+		"end": {
+			"col": 36,
+			"row": 4,
+		},
+		"file": "p2.rego",
+		"row": 4,
+		"text": "not partial",
+	})
+}
+
 test_fail_multivalue_not_reference_different_package_using_direct_reference if {
 	agg1 := rule.aggregate with input as regal.parse_module("p1.rego", `package foo
 
diff --git a/bundle/regal/rules/bugs/not-equals-in-loop/not_equals_in_loop.rego b/bundle/regal/rules/bugs/not-equals-in-loop/not_equals_in_loop.rego
index 26ab665e..9dea6385 100644
--- a/bundle/regal/rules/bugs/not-equals-in-loop/not_equals_in_loop.rego
+++ b/bundle/regal/rules/bugs/not-equals-in-loop/not_equals_in_loop.rego
@@ -6,7 +6,10 @@ import data.regal.ast
 import data.regal.result
 
 report contains violation if {
-	terms := ast.exprs[_][_].terms
+	some rule_index, i
+	ast.found.expressions[rule_index][i].terms[0].type == "ref"
+
+	terms := ast.found.expressions[rule_index][i].terms
 
 	terms[0].type == "ref"
 	terms[0].value[0].type == "var"
diff --git a/bundle/regal/rules/bugs/redundant-existence-check/redundant_existence_check.rego b/bundle/regal/rules/bugs/redundant-existence-check/redundant_existence_check.rego
index 61a09532..52c31a03 100644
--- a/bundle/regal/rules/bugs/redundant-existence-check/redundant_existence_check.rego
+++ b/bundle/regal/rules/bugs/redundant-existence-check/redundant_existence_check.rego
@@ -9,7 +9,7 @@ import data.regal.result
 # description: check rule bodies for redundant existence checks
 report contains violation if {
 	some rule_index, rule in input.rules
-	some expr_index, expr in ast.exprs[rule_index]
+	some expr_index, expr in _exprs[rule_index]
 
 	expr.terms.type == "ref"
 
@@ -52,10 +52,16 @@ report contains violation if {
 
 	rule.head.value.type == "ref"
 
-	some expr in ast.exprs[rule_index]
+	some expr in _exprs[rule_index]
 
 	expr.terms.type == "ref"
 	ast.ref_value_equal(expr.terms.value, rule.head.value.value)
 
 	violation := result.fail(rego.metadata.chain(), result.ranged_from_ref(expr.terms.value))
 }
+
+# all top-level expressions in module
+_exprs[rule_index][expr_index] := expr if {
+	some rule_index, rule in input.rules
+	some expr_index, expr in rule.body
+}
diff --git a/bundle/regal/rules/bugs/sprintf-arguments-mismatch/sprintf_arguments_mismatch.rego b/bundle/regal/rules/bugs/sprintf-arguments-mismatch/sprintf_arguments_mismatch.rego
index 0f50baec..08f513de 100644
--- a/bundle/regal/rules/bugs/sprintf-arguments-mismatch/sprintf_arguments_mismatch.rego
+++ b/bundle/regal/rules/bugs/sprintf-arguments-mismatch/sprintf_arguments_mismatch.rego
@@ -63,11 +63,11 @@ _first_arg_value(rule_index, term) := found if {
 	trow := util.to_location_object(term.location).row
 
 	found := [rhs |
-		some expr in ast.exprs[to_number(rule_index)]
+		some expr in ast.found.expressions[rule_index]
 
 		util.to_location_object(expr.location).row < trow
 
-		[lhs, rhs] := ast.assignment_terms(expr)
+		[lhs, rhs] := ast.assignment_terms(expr.terms)
 		lhs.type == "var"
 		lhs.value == term.value
 		rhs.type == "string"
diff --git a/bundle/regal/rules/bugs/sprintf-arguments-mismatch/sprintf_arguments_mismatch_test.rego b/bundle/regal/rules/bugs/sprintf-arguments-mismatch/sprintf_arguments_mismatch_test.rego
index b064693d..3a5db38f 100644
--- a/bundle/regal/rules/bugs/sprintf-arguments-mismatch/sprintf_arguments_mismatch_test.rego
+++ b/bundle/regal/rules/bugs/sprintf-arguments-mismatch/sprintf_arguments_mismatch_test.rego
@@ -14,7 +14,10 @@ test_fail_too_many_values_in_array if {
 		"location": {
 			"row": 5,
 			"col": 14,
-			"end": {"col": 26, "row": 5},
+			"end": {
+				"col": 26,
+				"row": 5,
+			},
 			"file": "policy.rego",
 			"text": "x := sprintf(\"%s\", [1, 2])",
 		},
@@ -26,6 +29,33 @@ test_fail_too_many_values_in_array if {
 	}}
 }
 
+test_fail_too_many_values_in_array_nested if {
+	r := rule.report with input as ast.with_rego_v1(`x := [1 |
+		y := [s | s := sprintf("%s", [1, 2])]
+	]`)
+
+	r == {{
+		"category": "bugs",
+		"description": "Mismatch in `sprintf` arguments count",
+		"level": "error",
+		"location": {
+			"col": 26,
+			"end": {
+				"col": 38,
+				"row": 6,
+			},
+			"file": "policy.rego",
+			"row": 6,
+			"text": "\t\ty := [s | s := sprintf(\"%s\", [1, 2])]",
+		},
+		"related_resources": [{
+			"description": "documentation",
+			"ref": config.docs.resolve_url("$baseUrl/$category/sprintf-arguments-mismatch", "bugs"),
+		}],
+		"title": "sprintf-arguments-mismatch",
+	}}
+}
+
 test_fail_too_few_values_in_array if {
 	r := rule.report with input as ast.with_rego_v1(`x := sprintf("%s%v", [1])`)
 	r == {{
@@ -35,7 +65,10 @@ test_fail_too_few_values_in_array if {
 		"location": {
 			"row": 5,
 			"col": 14,
-			"end": {"col": 25, "row": 5},
+			"end": {
+				"col": 25,
+				"row": 5,
+			},
 			"file": "policy.rego",
 			"text": `x := sprintf("%s%v", [1])`,
 		},
@@ -87,7 +120,10 @@ test_fail_first_arg_is_variable_with_nonmatching_pattern if {
 		"level": "error",
 		"location": {
 			"col": 11,
-			"end": {"col": 21, "row": 7},
+			"end": {
+				"col": 21,
+				"row": 7,
+			},
 			"file": "policy.rego",
 			"row": 7,
 			"text": "\t\tsprintf(s, [\"foo\"])",
diff --git a/bundle/regal/rules/bugs/unassigned-return-value/unassigned_return_value.rego b/bundle/regal/rules/bugs/unassigned-return-value/unassigned_return_value.rego
index 493f31bf..682e42e6 100644
--- a/bundle/regal/rules/bugs/unassigned-return-value/unassigned_return_value.rego
+++ b/bundle/regal/rules/bugs/unassigned-return-value/unassigned_return_value.rego
@@ -7,7 +7,7 @@ import data.regal.config
 import data.regal.result
 
 report contains violation if {
-	terms := ast.exprs[_][_].terms
+	terms := ast.found.expressions[_][_].terms
 
 	terms[0].type == "ref"
 	terms[0].value[0].type == "var"
diff --git a/bundle/regal/rules/bugs/unassigned-return-value/unassigned_return_value_test.rego b/bundle/regal/rules/bugs/unassigned-return-value/unassigned_return_value_test.rego
index 2d35055f..c3ec2873 100644
--- a/bundle/regal/rules/bugs/unassigned-return-value/unassigned_return_value_test.rego
+++ b/bundle/regal/rules/bugs/unassigned-return-value/unassigned_return_value_test.rego
@@ -33,6 +33,37 @@ test_fail_unused_return_value if {
 	}}
 }
 
+test_fail_unused_return_value_nested if {
+	r := rule.report with input as ast.with_rego_v1(`allow if {
+		comprehension := [x |
+			indexof("s", "s")
+			x := 1
+		]
+	}`)
+		with config.capabilities as capabilities.provided
+
+	r == {{
+		"category": "bugs",
+		"description": "Non-boolean return value unassigned",
+		"level": "error",
+		"location": {
+			"col": 4,
+			"end": {
+				"col": 11,
+				"row": 7,
+			},
+			"file": "policy.rego",
+			"row": 7,
+			"text": "\t\t\tindexof(\"s\", \"s\")",
+		},
+		"related_resources": [{
+			"description": "documentation",
+			"ref": config.docs.resolve_url("$baseUrl/$category/unassigned-return-value", "bugs"),
+		}],
+		"title": "unassigned-return-value",
+	}}
+}
+
 test_success_unused_boolean_return_value if {
 	r := rule.report with input as ast.policy(`allow if { startswith("s", "s") }`)
 		with config.capabilities as capabilities.provided
diff --git a/bundle/regal/rules/bugs/unused-output-variable/unused_output_variable.rego b/bundle/regal/rules/bugs/unused-output-variable/unused_output_variable.rego
index 9a35b7da..5d1a7f24 100644
--- a/bundle/regal/rules/bugs/unused-output-variable/unused_output_variable.rego
+++ b/bundle/regal/rules/bugs/unused-output-variable/unused_output_variable.rego
@@ -21,12 +21,13 @@ import data.regal.result
 #   ```
 report contains violation if {
 	some rule_index
-	var_refs := _ref_vars[rule_index][_]
+	var_refs := ast.found.vars[rule_index].ref
 
 	count(var_refs) == 1
 
 	some var in var_refs
 
+	not ast.is_wildcard(var)
 	not ast.var_in_head(input.rules[to_number(rule_index)].head, var.value)
 	not ast.var_in_call(ast.function_calls, rule_index, var.value)
 	not _ref_base_vars[rule_index][var.value]
@@ -39,13 +40,6 @@ report contains violation if {
 	violation := result.fail(rego.metadata.chain(), result.location(var))
 }
 
-_ref_vars[rule_index][var.value] contains var if {
-	some rule_index
-	var := ast.found.vars[rule_index].ref[_]
-
-	not ast.is_wildcard(var)
-}
-
 # "a" in "a[foo]", and not foo
 _ref_base_vars[rule_index][term.value] contains term if {
 	some rule_index
diff --git a/bundle/regal/rules/custom/forbidden-function-call/forbidden_function_call.rego b/bundle/regal/rules/custom/forbidden-function-call/forbidden_function_call.rego
index bd89642a..3a6786b6 100644
--- a/bundle/regal/rules/custom/forbidden-function-call/forbidden_function_call.rego
+++ b/bundle/regal/rules/custom/forbidden-function-call/forbidden_function_call.rego
@@ -13,7 +13,7 @@ report contains violation if {
 	# avoid traversal if no forbidden function is called
 	util.intersects(forbidden, ast.builtin_functions_called)
 
-	ref := ast.found.refs[_][_]
+	ref := ast.found.calls[_][_]
 	name := ast.ref_to_string(ref[0].value)
 	name in forbidden
 
diff --git a/bundle/regal/rules/custom/narrow-argument/narrow_argument.rego b/bundle/regal/rules/custom/narrow-argument/narrow_argument.rego
index 4d643864..fd48ddf9 100644
--- a/bundle/regal/rules/custom/narrow-argument/narrow_argument.rego
+++ b/bundle/regal/rules/custom/narrow-argument/narrow_argument.rego
@@ -17,7 +17,6 @@ report contains violation if {
 	not _arg_used_in_call(indices, arg)
 
 	location := result.location(_first_named_arg_location(indices, arg))
-
 	violation := result.fail(rego.metadata.chain(), object.union(
 		location,
 		{"description": _message(count(refs), arg, narrowed)},
@@ -42,8 +41,9 @@ _narrow(refs) := narrowed if {
 	arr := util.any_set_item(refs)
 
 	count(arr) > 1
+	not _nested(arr)
 
-	narrowed := concat(".", arr)
+	narrowed := ast.ref_to_string(_to_terms(arr))
 }
 
 _narrow(refs) := narrowed if {
@@ -52,8 +52,9 @@ _narrow(refs) := narrowed if {
 	prefix := util.longest_prefix(refs)
 
 	count(prefix) > 1
+	not _nested(prefix)
 
-	narrowed := concat(".", prefix)
+	narrowed := ast.ref_to_string(_to_terms(prefix))
 }
 
 _first_named_arg_location(indices, name) := [arg.location |
@@ -126,6 +127,14 @@ _first_var_pos(ref) := pos if {
 	][0]
 } else := count(ref) + 1
 
-_exclude_arg(cfg, name) if {
-	name in cfg["exclude-args"]
+_exclude_arg(cfg, name) if name in cfg["exclude-args"]
+
+_to_terms(arr) := [_to_term(item) | some item in arr]
+
+_to_term(value) := {"type": "number", "value": value} if is_number(value)
+_to_term(value) := {"type": "string", "value": value} if is_string(value)
+
+_nested(arr) if {
+	some item in arr
+	not type_name(item) in ast.scalar_types
 }
diff --git a/bundle/regal/rules/custom/narrow-argument/narrow_argument_test.rego b/bundle/regal/rules/custom/narrow-argument/narrow_argument_test.rego
index 31deb4be..53d8d8a1 100644
--- a/bundle/regal/rules/custom/narrow-argument/narrow_argument_test.rego
+++ b/bundle/regal/rules/custom/narrow-argument/narrow_argument_test.rego
@@ -63,11 +63,48 @@ test_fail_can_be_narrowed_prefixed_ref if {
 	}}
 }
 
+test_fail_can_be_narrowed_prefixed_array_ref if {
+	r := rule.report with input as ast.policy(`
+		fun(arr) if arr[0].y.number == 1
+		fun(arr) if arr[0].y.string == "1"
+	`)
+		with config.rules as {"custom": {"narrow-argument": {"level": "error"}}}
+
+	r == {{
+		"category": "custom",
+		"description": "Argument arr always referenced by a common prefix, value passed can be narrowed to arr[0].y",
+		"level": "error",
+		"location": {
+			"col": 7,
+			"end": {
+				"col": 10,
+				"row": 4,
+			},
+			"file": "policy.rego",
+			"row": 4,
+			"text": "\t\tfun(arr) if arr[0].y.number == 1",
+		},
+		"related_resources": [{
+			"description": "documentation",
+			"ref": config.docs.resolve_url("$baseUrl/$category/narrow-argument", "custom"),
+		}],
+		"title": "narrow-argument",
+	}}
+}
+
 test_success_can_not_be_narrowed_arg_is_least_common_denominator if {
 	r := rule.report with input as ast.policy(`
 		fun(obj) if obj.typ == "string"
 		fun(obj) if obj.val == "string"
 	`)
+		with config.rules as {"custom": {"narrow-argument": {"level": "error"}}}
+
+	r == set()
+}
+
+test_success_nested_or_variable_path_not_narrowed if {
+	r := rule.report with input as ast.policy(`foo(lines) := lines[bar - 1]`)
+		with config.rules as {"custom": {"narrow-argument": {"level": "error"}}}
 
 	r == set()
 }
diff --git a/bundle/regal/rules/idiomatic/non-raw-regex-pattern/non_raw_regex_pattern.rego b/bundle/regal/rules/idiomatic/non-raw-regex-pattern/non_raw_regex_pattern.rego
index 25ab0e59..dce8fa8a 100644
--- a/bundle/regal/rules/idiomatic/non-raw-regex-pattern/non_raw_regex_pattern.rego
+++ b/bundle/regal/rules/idiomatic/non-raw-regex-pattern/non_raw_regex_pattern.rego
@@ -10,7 +10,7 @@ report contains violation if {
 	# skip traversing refs if no builtin regex function calls are registered
 	util.intersects(_re_pattern_function_names, ast.builtin_functions_called)
 
-	value := ast.found.refs[_][_]
+	value := ast.found.calls[_][_]
 
 	value[0].value[0].type == "var"
 	value[0].value[0].value == "regex"
diff --git a/bundle/regal/rules/idiomatic/prefer-set-or-object-rule/prefer_set_or_object_rule.rego b/bundle/regal/rules/idiomatic/prefer-set-or-object-rule/prefer_set_or_object_rule.rego
index a41ada28..ed41b225 100644
--- a/bundle/regal/rules/idiomatic/prefer-set-or-object-rule/prefer_set_or_object_rule.rego
+++ b/bundle/regal/rules/idiomatic/prefer-set-or-object-rule/prefer_set_or_object_rule.rego
@@ -54,7 +54,7 @@ _is_array_conversion(value) if {
 
 	count(body) == 1
 
-	[lhs, rhs] := ast.assignment_terms(body[0])
+	[lhs, rhs] := ast.assignment_terms(body[0].terms)
 
 	# Assignment to comprehension variable
 	lhs.type == "var"
diff --git a/bundle/regal/rules/idiomatic/use-strings-count/use_strings_count.rego b/bundle/regal/rules/idiomatic/use-strings-count/use_strings_count.rego
index 50d38523..3026ac83 100644
--- a/bundle/regal/rules/idiomatic/use-strings-count/use_strings_count.rego
+++ b/bundle/regal/rules/idiomatic/use-strings-count/use_strings_count.rego
@@ -15,7 +15,7 @@ notices contains result.notice(rego.metadata.chain()) if not capabilities.has_st
 # METADATA
 # description: flag calls to `count` where the first argument is a call to `indexof_n`
 report contains violation if {
-	ref := ast.found.refs[_][_]
+	ref := ast.found.calls[_][_]
 
 	ref[0].value[0].type == "var"
 	ref[0].value[0].value == "count"
diff --git a/bundle/regal/rules/imports/prefer-package-imports/prefer_package_imports.rego b/bundle/regal/rules/imports/prefer-package-imports/prefer_package_imports.rego
index 4109093f..bb436d7c 100644
--- a/bundle/regal/rules/imports/prefer-package-imports/prefer_package_imports.rego
+++ b/bundle/regal/rules/imports/prefer-package-imports/prefer_package_imports.rego
@@ -19,7 +19,7 @@ aggregate contains entry if {
 
 		# Special case for custom rules, where we don't want to flag e.g. `import data.regal.ast`
 		# as unknown, even though it's not a package included in evaluation.
-		not _custom_regal_package_and_import(ast.package_path, path)
+		not _custom_regal_package_and_import(ast.package_path, path[0])
 
 		imp := object.union(result.location(_import), {"path": path})
 	]
@@ -30,10 +30,9 @@ aggregate contains entry if {
 	})
 }
 
-_custom_regal_package_and_import(pkg_path, path) if {
+_custom_regal_package_and_import(pkg_path, "regal") if {
 	pkg_path[0] == "custom"
 	pkg_path[1] == "regal"
-	path[0] == "regal"
 }
 
 # METADATA
diff --git a/bundle/regal/rules/imports/unresolved-import/unresolved_import.rego b/bundle/regal/rules/imports/unresolved-import/unresolved_import.rego
index dc74dba8..57c8c2f9 100644
--- a/bundle/regal/rules/imports/unresolved-import/unresolved_import.rego
+++ b/bundle/regal/rules/imports/unresolved-import/unresolved_import.rego
@@ -20,7 +20,7 @@ aggregate contains entry if {
 
 		# Special case for custom rules, where we don't want to flag e.g. `import data.regal.ast`
 		# as unknown, even though it's not a package included in evaluation.
-		not _custom_regal_package_and_import(ast.package_path, path)
+		not _custom_regal_package_and_import(ast.package_path, path[0])
 
 		imp := object.union(result.location(_import), {"path": path})
 	]
@@ -63,25 +63,16 @@ aggregate_report contains violation if {
 	violation := result.fail(rego.metadata.chain(), result.location(imp))
 }
 
-_custom_regal_package_and_import(pkg_path, path) if {
+_custom_regal_package_and_import(pkg_path, "regal") if {
 	pkg_path[0] == "custom"
 	pkg_path[1] == "regal"
-	path[0] == "regal"
 }
 
 # the package part will always be included exported refs
 # but if we have a rule like foo.bar.baz
 # we'll want to include both foo.bar and foo.bar.baz
 _to_paths(pkg_path, ref) := util.all_paths(_to_path(pkg_path, ref)) if count(ref) < 3
-
-_to_paths(pkg_path, ref) := paths if {
-	count(ref) > 2
-
-	paths := [path |
-		some p in util.all_paths(ref)
-		path := _to_path(pkg_path, p)
-	]
-}
+_to_paths(pkg_path, ref) := [_to_path(pkg_path, p) | some p in util.all_paths(ref)] if count(ref) > 2
 
 _to_path(pkg_path, ref) := array.concat(pkg_path, [str |
 	some i, part in ref
diff --git a/bundle/regal/rules/performance/defer-assignment/defer_assignment.rego b/bundle/regal/rules/performance/defer-assignment/defer_assignment.rego
index 3822258f..666555aa 100644
--- a/bundle/regal/rules/performance/defer-assignment/defer_assignment.rego
+++ b/bundle/regal/rules/performance/defer-assignment/defer_assignment.rego
@@ -9,7 +9,7 @@ report contains violation if {
 	some i, rule in input.rules
 	some j, expr in rule.body
 
-	[var, rhs] := ast.assignment_terms(expr)
+	[var, rhs] := ast.assignment_terms(expr.terms)
 
 	not _ref_with_vars(rhs)
 
@@ -19,7 +19,7 @@ report contains violation if {
 
 	next := rule.body[j + 1]
 
-	not ast.is_assignment(next)
+	not ast.is_assignment(next.terms[0])
 	not ast.var_in_head(rule.head, var.value)
 	not _var_value_used_in_expression(var.value, next)
 	not _iteration_expression(next.terms)
@@ -91,6 +91,7 @@ _iteration_expression(terms) if {
 	terms[0].value[0].value == "walk"
 }
 
+# regal ignore:narrow-argument
 _print_call(terms) if {
 	terms[0].value[0].type == "var"
 	terms[0].value[0].value == "print"
diff --git a/bundle/regal/rules/performance/with-outside-test-context/with_outside_test_context.rego b/bundle/regal/rules/performance/with-outside-test-context/with_outside_test_context.rego
index d9792224..ce9262c4 100644
--- a/bundle/regal/rules/performance/with-outside-test-context/with_outside_test_context.rego
+++ b/bundle/regal/rules/performance/with-outside-test-context/with_outside_test_context.rego
@@ -6,8 +6,8 @@ import data.regal.ast
 import data.regal.result
 
 report contains violation if {
-	some rule_index, rule in input.rules
-	some expr_index, expr in ast.exprs[rule_index]
+	some i, rule in input.rules
+	some expr in ast.found.expressions[ast.rule_index_strings[i]]
 
 	expr["with"]
 	not strings.any_prefix_match(ast.ref_to_string(rule.head.ref), {"test_", "todo_test"})
diff --git a/bundle/regal/rules/style/comprehension-term-assignment/comprehension_term_assignment.rego b/bundle/regal/rules/style/comprehension-term-assignment/comprehension_term_assignment.rego
index 12949556..2a5fc0f2 100644
--- a/bundle/regal/rules/style/comprehension-term-assignment/comprehension_term_assignment.rego
+++ b/bundle/regal/rules/style/comprehension-term-assignment/comprehension_term_assignment.rego
@@ -25,7 +25,7 @@ report contains violation if {
 
 	some expr in comp.body
 
-	[lhs, rhs] := ast.assignment_terms(expr)
+	[lhs, rhs] := ast.assignment_terms(expr.terms)
 
 	lhs.type == comp.term.type
 	lhs.value == comp.term.value
@@ -55,7 +55,7 @@ report contains violation if {
 
 	some expr in comp.body
 
-	[lhs, rhs] := ast.assignment_terms(expr)
+	[lhs, rhs] := ast.assignment_terms(expr.terms)
 
 	some kind in ["key", "value"]
 
diff --git a/bundle/regal/rules/style/double-negative/double_negative.rego b/bundle/regal/rules/style/double-negative/double_negative.rego
index 99d8f793..cd723706 100644
--- a/bundle/regal/rules/style/double-negative/double_negative.rego
+++ b/bundle/regal/rules/style/double-negative/double_negative.rego
@@ -11,8 +11,9 @@ import data.regal.ast
 import data.regal.result
 
 report contains violation if {
-	some node
-	ast.negated_expressions[_][node].terms.type == "var"
+	some node, i
+	ast.found.expressions[i][node].negated
+	ast.found.expressions[i][node].terms.type == "var"
 
 	strings.any_prefix_match(node.terms.value, {
 		"cannot_",
diff --git a/bundle/regal/rules/style/pointless-reassignment/pointless_reassignment.rego b/bundle/regal/rules/style/pointless-reassignment/pointless_reassignment.rego
index b8853932..b15fc83d 100644
--- a/bundle/regal/rules/style/pointless-reassignment/pointless_reassignment.rego
+++ b/bundle/regal/rules/style/pointless-reassignment/pointless_reassignment.rego
@@ -23,7 +23,7 @@ report contains violation if {
 
 	not expr["with"]
 
-	[lhs, rhs] := ast.assignment_terms(expr)
+	[lhs, rhs] := ast.assignment_terms(expr.terms)
 
 	lhs.type == "var"
 	rhs.type == "var"
diff --git a/bundle/regal/rules/style/rule-length/rule_length.rego b/bundle/regal/rules/style/rule-length/rule_length.rego
index 493e4c76..d0dd876d 100644
--- a/bundle/regal/rules/style/rule-length/rule_length.rego
+++ b/bundle/regal/rules/style/rule-length/rule_length.rego
@@ -14,7 +14,7 @@ report contains violation if {
 	rule_location := util.to_location_object(rule.location)
 	lines := split(rule_location.text, "\n")
 
-	_line_count(cfg, rule_location.row, lines) > cfg[_max_length_property(rule.head)]
+	_line_count(cfg, rule_location.row, lines) > cfg[_max_length_property(rule.head.ref[0].value)]
 
 	not _no_body_exception(cfg, rule)
 
@@ -28,7 +28,7 @@ _no_body_exception(cfg, rule) if {
 
 default _max_length_property(_) := "max-rule-length"
 
-_max_length_property(head) := "max-test-rule-length" if startswith(head.ref[0].value, "test_")
+_max_length_property(value) := "max-test-rule-length" if startswith(value, "test_")
 
 _line_count(cfg, _, lines) := count(lines) if cfg["count-comments"] == true
 
diff --git a/bundle/regal/rules/style/unconditional-assignment/unconditional_assignment.rego b/bundle/regal/rules/style/unconditional-assignment/unconditional_assignment.rego
index 13170f58..65f1d9ec 100644
--- a/bundle/regal/rules/style/unconditional-assignment/unconditional_assignment.rego
+++ b/bundle/regal/rules/style/unconditional-assignment/unconditional_assignment.rego
@@ -47,17 +47,20 @@ report contains violation if {
 	# Multi-value rule
 	rule.head.key.type == "var"
 
-	not rule.body[0]["with"]
+	expr := rule.body[0]
 
-	_assignment_expr(rule.body[0].terms)
+	not expr["with"]
 
-	rule.body[0].terms[1].type == "var"
-	rule.body[0].terms[1].value == rule.head.key.value
+	_assignment_expr(expr.terms)
 
-	violation := result.fail(rego.metadata.chain(), result.infix_expr_location(rule.body[0].terms))
+	expr.terms[1].type == "var"
+	expr.terms[1].value == rule.head.key.value
+
+	violation := result.fail(rego.metadata.chain(), result.infix_expr_location(expr.terms))
 }
 
 # Assignment using either = or :=
+# regal ignore:narrow-argument
 _assignment_expr(terms) if {
 	terms[0].type == "ref"
 	terms[0].value[0].type == "var"
diff --git a/bundle/regal/rules/style/unnecessary-some/unnecessary_some.rego b/bundle/regal/rules/style/unnecessary-some/unnecessary_some.rego
index 0b22ff27..c95fbd25 100644
--- a/bundle/regal/rules/style/unnecessary-some/unnecessary_some.rego
+++ b/bundle/regal/rules/style/unnecessary-some/unnecessary_some.rego
@@ -14,24 +14,20 @@ report contains violation if {
 	symbols[0].type == "call"
 	symbols[0].value[0].type == "ref"
 
-	_some_is_unnecessary(symbols, ast.scalar_types)
+	_some_is_unnecessary(symbols[0].value, ast.scalar_types)
 
 	violation := result.fail(rego.metadata.chain(), result.location(symbols))
 }
 
-_some_is_unnecessary(value, scalar_types) if {
-	ref := value[0].value[0].value
-
-	[ref[0].value, ref[1].value] == ["internal", "member_2"]
-
-	value[0].value[1].type in scalar_types
+_some_is_unnecessary(symbol, scalar_types) if {
+	symbol[0].value[0].value == "internal"
+	symbol[0].value[1].value == "member_2"
+	symbol[1].type in scalar_types
 }
 
-_some_is_unnecessary(value, scalar_types) if {
-	ref := value[0].value[0].value
-
-	[ref[0].value, ref[1].value] == ["internal", "member_3"]
-
-	value[0].value[1].type in scalar_types
-	value[0].value[2].type in scalar_types
+_some_is_unnecessary(symbol, scalar_types) if {
+	symbol[0].value[0].value == "internal"
+	symbol[0].value[1].value == "member_3"
+	symbol[1].type in scalar_types
+	symbol[2].type in scalar_types
 }
diff --git a/bundle/regal/rules/style/yoda-condition/yoda_condition.rego b/bundle/regal/rules/style/yoda-condition/yoda_condition.rego
index 0e4f8cad..b87de5fc 100644
--- a/bundle/regal/rules/style/yoda-condition/yoda_condition.rego
+++ b/bundle/regal/rules/style/yoda-condition/yoda_condition.rego
@@ -6,7 +6,7 @@ import data.regal.ast
 import data.regal.result
 
 report contains violation if {
-	value := ast.found.refs[_][_]
+	value := ast.found.calls[_][_]
 
 	value[0].value[0].type == "var"
 	value[0].value[0].value in {"equal", "neq", "gt", "lt", "gte", "lte"}
diff --git a/bundle/regal/rules/testing/dubious-print-sprintf/dubious_print_sprintf.rego b/bundle/regal/rules/testing/dubious-print-sprintf/dubious_print_sprintf.rego
index 421c39f9..4b23e026 100644
--- a/bundle/regal/rules/testing/dubious-print-sprintf/dubious_print_sprintf.rego
+++ b/bundle/regal/rules/testing/dubious-print-sprintf/dubious_print_sprintf.rego
@@ -9,7 +9,7 @@ report contains violation if {
 	# skip traversing refs if no print calls are registered
 	"print" in ast.builtin_functions_called
 
-	value := ast.found.refs[_][_]
+	value := ast.found.calls[_][_]
 
 	value[0].value[0].type == "var"
 	value[0].value[0].value == "print"
diff --git a/bundle/regal/rules/testing/print-or-trace-call/print_or_trace_call.rego b/bundle/regal/rules/testing/print-or-trace-call/print_or_trace_call.rego
index b36b558a..2ea01dbc 100644
--- a/bundle/regal/rules/testing/print-or-trace-call/print_or_trace_call.rego
+++ b/bundle/regal/rules/testing/print-or-trace-call/print_or_trace_call.rego
@@ -10,7 +10,7 @@ report contains violation if {
 	# skip iteration of refs if no print or trace calls are registered
 	util.intersects(ast.builtin_functions_called, {"print", "trace"})
 
-	ref := ast.found.refs[_][_][0]
+	ref := ast.found.calls[_][_][0]
 
 	ref.value[0].type == "var"
 	ref.value[0].value in {"print", "trace"}
diff --git a/internal/embeds/schemas/regal-ast.json b/internal/embeds/schemas/regal-ast.json
index 98b52e85..aabe858a 100644
--- a/internal/embeds/schemas/regal-ast.json
+++ b/internal/embeds/schemas/regal-ast.json
@@ -232,9 +232,6 @@
     },
     "head": {
       "properties": {
-        "name": {
-          "type": "string"
-        },
         "ref": {
           "$ref": "#/$defs/ref"
         },
diff --git a/internal/lsp/config/watcher.go b/internal/lsp/config/watcher.go
index 47fec7cb..3ed8485d 100644
--- a/internal/lsp/config/watcher.go
+++ b/internal/lsp/config/watcher.go
@@ -64,6 +64,33 @@ func (w *Watcher) Start(ctx context.Context) error {
 	return nil
 }
 
+func (w *Watcher) Watch(configFilePath string) {
+	w.pathUpdates <- configFilePath
+}
+
+func (w *Watcher) Stop() error {
+	if w.fsWatcher != nil {
+		if err := w.fsWatcher.Close(); err != nil {
+			return fmt.Errorf("failed to close fsnotify watcher: %w", err)
+		}
+
+		return nil
+	}
+
+	return nil
+}
+
+func (w *Watcher) IsWatching() bool {
+	w.fsWatcherLock.Lock()
+	defer w.fsWatcherLock.Unlock()
+
+	if w.fsWatcher == nil {
+		return false
+	}
+
+	return len(w.fsWatcher.WatchList()) > 0
+}
+
 func (w *Watcher) loop(ctx context.Context) {
 	for {
 		select {
@@ -109,30 +136,3 @@ func (w *Watcher) loop(ctx context.Context) {
 		}
 	}
 }
-
-func (w *Watcher) Watch(configFilePath string) {
-	w.pathUpdates <- configFilePath
-}
-
-func (w *Watcher) Stop() error {
-	if w.fsWatcher != nil {
-		if err := w.fsWatcher.Close(); err != nil {
-			return fmt.Errorf("failed to close fsnotify watcher: %w", err)
-		}
-
-		return nil
-	}
-
-	return nil
-}
-
-func (w *Watcher) IsWatching() bool {
-	w.fsWatcherLock.Lock()
-	defer w.fsWatcherLock.Unlock()
-
-	if w.fsWatcher == nil {
-		return false
-	}
-
-	return len(w.fsWatcher.WatchList()) > 0
-}
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index 605f72f6..889fa3ed 100644
--- a/internal/lsp/server.go
+++ b/internal/lsp/server.go
@@ -102,32 +102,6 @@ type LanguageServerOptions struct {
 	WorkspaceDiagnosticsPoll time.Duration
 }
 
-func NewLanguageServer(ctx context.Context, opts *LanguageServerOptions) *LanguageServer {
-	c := cache.NewCache()
-	store := NewRegalStore()
-
-	ls := &LanguageServer{
-		cache:                       c,
-		regoStore:                   store,
-		logWriter:                   opts.LogWriter,
-		logLevel:                    opts.LogLevel,
-		lintFileJobs:                make(chan lintFileJob, 10),
-		lintWorkspaceJobs:           make(chan lintWorkspaceJob, 10),
-		builtinsPositionJobs:        make(chan lintFileJob, 10),
-		commandRequest:              make(chan types.ExecuteCommandParams, 10),
-		templateFileJobs:            make(chan lintFileJob, 10),
-		completionsManager:          completions.NewDefaultManager(ctx, c, store),
-		webServer:                   web.NewServer(c, opts.LogWriter, opts.LogLevel),
-		loadedBuiltins:              concurrent.MapOf(make(map[string]map[string]*ast.Builtin)),
-		workspaceDiagnosticsPoll:    opts.WorkspaceDiagnosticsPoll,
-		loadedConfigAllRegoVersions: concurrent.MapOf(make(map[string]ast.RegoVersion)),
-	}
-
-	ls.configWatcher = lsconfig.NewWatcher(&lsconfig.WatcherOpts{LogFunc: ls.logf})
-
-	return ls
-}
-
 type LanguageServer struct {
 	logWriter io.Writer
 	logLevel  log.Level
@@ -182,6 +156,32 @@ type lintWorkspaceJob struct {
 	AggregateReportOnly bool
 }
 
+func NewLanguageServer(ctx context.Context, opts *LanguageServerOptions) *LanguageServer {
+	c := cache.NewCache()
+	store := NewRegalStore()
+
+	ls := &LanguageServer{
+		cache:                       c,
+		regoStore:                   store,
+		logWriter:                   opts.LogWriter,
+		logLevel:                    opts.LogLevel,
+		lintFileJobs:                make(chan lintFileJob, 10),
+		lintWorkspaceJobs:           make(chan lintWorkspaceJob, 10),
+		builtinsPositionJobs:        make(chan lintFileJob, 10),
+		commandRequest:              make(chan types.ExecuteCommandParams, 10),
+		templateFileJobs:            make(chan lintFileJob, 10),
+		completionsManager:          completions.NewDefaultManager(ctx, c, store),
+		webServer:                   web.NewServer(c, opts.LogWriter, opts.LogLevel),
+		loadedBuiltins:              concurrent.MapOf(make(map[string]map[string]*ast.Builtin)),
+		workspaceDiagnosticsPoll:    opts.WorkspaceDiagnosticsPoll,
+		loadedConfigAllRegoVersions: concurrent.MapOf(make(map[string]ast.RegoVersion)),
+	}
+
+	ls.configWatcher = lsconfig.NewWatcher(&lsconfig.WatcherOpts{LogFunc: ls.logf})
+
+	return ls
+}
+
 //nolint:wrapcheck
 func (l *LanguageServer) Handle(ctx context.Context, _ *jsonrpc2.Conn, req *jsonrpc2.Request) (any, error) {
 	l.logf(log.LevelDebug, "received request: %s", req.Method)
@@ -479,59 +479,6 @@ func (l *LanguageServer) StartHoverWorker(ctx context.Context) {
 	}
 }
 
-func (l *LanguageServer) getLoadedConfig() *config.Config {
-	l.loadedConfigLock.RLock()
-	defer l.loadedConfigLock.RUnlock()
-
-	return l.loadedConfig
-}
-
-func (l *LanguageServer) getEnabledNonAggregateRules() []string {
-	l.loadedConfigLock.RLock()
-	defer l.loadedConfigLock.RUnlock()
-
-	return l.loadedConfigEnabledNonAggregateRules
-}
-
-func (l *LanguageServer) getEnabledAggregateRules() []string {
-	l.loadedConfigLock.RLock()
-	defer l.loadedConfigLock.RUnlock()
-
-	return l.loadedConfigEnabledAggregateRules
-}
-
-// loadEnabledRulesFromConfig is used to cache the enabled rules for the current
-// config. These take some time to compute and only change when config changes,
-// so we can store them on the server to speed up diagnostic runs.
-func (l *LanguageServer) loadEnabledRulesFromConfig(ctx context.Context, cfg config.Config) error {
-	lint := linter.NewLinter().WithUserConfig(cfg)
-
-	enabledRules, err := lint.DetermineEnabledRules(ctx)
-	if err != nil {
-		return fmt.Errorf("failed to determine enabled rules: %w", err)
-	}
-
-	enabledAggregateRules, err := lint.DetermineEnabledAggregateRules(ctx)
-	if err != nil {
-		return fmt.Errorf("failed to determine enabled aggregate rules: %w", err)
-	}
-
-	l.loadedConfigLock.Lock()
-	defer l.loadedConfigLock.Unlock()
-
-	l.loadedConfigEnabledNonAggregateRules = []string{}
-
-	for _, r := range enabledRules {
-		if !slices.Contains(enabledAggregateRules, r) {
-			l.loadedConfigEnabledNonAggregateRules = append(l.loadedConfigEnabledNonAggregateRules, r)
-		}
-	}
-
-	l.loadedConfigEnabledAggregateRules = enabledAggregateRules
-
-	return nil
-}
-
 func (l *LanguageServer) StartConfigWorker(ctx context.Context) {
 	if err := l.configWatcher.Start(ctx); err != nil {
 		l.logf(log.LevelMessage, "failed to start config watcher: %s", err)
@@ -1075,6 +1022,59 @@ func (l *LanguageServer) StartWebServer(ctx context.Context) {
 	l.webServer.Start(ctx)
 }
 
+func (l *LanguageServer) getLoadedConfig() *config.Config {
+	l.loadedConfigLock.RLock()
+	defer l.loadedConfigLock.RUnlock()
+
+	return l.loadedConfig
+}
+
+func (l *LanguageServer) getEnabledNonAggregateRules() []string {
+	l.loadedConfigLock.RLock()
+	defer l.loadedConfigLock.RUnlock()
+
+	return l.loadedConfigEnabledNonAggregateRules
+}
+
+func (l *LanguageServer) getEnabledAggregateRules() []string {
+	l.loadedConfigLock.RLock()
+	defer l.loadedConfigLock.RUnlock()
+
+	return l.loadedConfigEnabledAggregateRules
+}
+
+// loadEnabledRulesFromConfig is used to cache the enabled rules for the current
+// config. These take some time to compute and only change when config changes,
+// so we can store them on the server to speed up diagnostic runs.
+func (l *LanguageServer) loadEnabledRulesFromConfig(ctx context.Context, cfg config.Config) error {
+	lint := linter.NewLinter().WithUserConfig(cfg)
+
+	enabledRules, err := lint.DetermineEnabledRules(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to determine enabled rules: %w", err)
+	}
+
+	enabledAggregateRules, err := lint.DetermineEnabledAggregateRules(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to determine enabled aggregate rules: %w", err)
+	}
+
+	l.loadedConfigLock.Lock()
+	defer l.loadedConfigLock.Unlock()
+
+	l.loadedConfigEnabledNonAggregateRules = []string{}
+
+	for _, r := range enabledRules {
+		if !slices.Contains(enabledAggregateRules, r) {
+			l.loadedConfigEnabledNonAggregateRules = append(l.loadedConfigEnabledNonAggregateRules, r)
+		}
+	}
+
+	l.loadedConfigEnabledAggregateRules = enabledAggregateRules
+
+	return nil
+}
+
 func (l *LanguageServer) templateContentsForFile(fileURI string) (string, error) {
 	// this function should not be called with files in the root, but if it is,
 	// then it is an error to prevent unwanted behavior.
diff --git a/pkg/fixer/fixer.go b/pkg/fixer/fixer.go
index efdbf008..68219eb1 100644
--- a/pkg/fixer/fixer.go
+++ b/pkg/fixer/fixer.go
@@ -24,6 +24,14 @@ const (
 	OnConflictRename OnConflictOperation = "rename"
 )
 
+// Fixer must be instantiated via NewFixer.
+type Fixer struct {
+	registeredFixes     map[string]any
+	onConflictOperation OnConflictOperation
+	registeredRoots     []string
+	versionsMap         map[string]ast.RegoVersion
+}
+
 // NewFixer instantiates a Fixer.
 func NewFixer() *Fixer {
 	return &Fixer{
@@ -33,14 +41,6 @@ func NewFixer() *Fixer {
 	}
 }
 
-// Fixer must be instantiated via NewFixer.
-type Fixer struct {
-	registeredFixes     map[string]any
-	onConflictOperation OnConflictOperation
-	registeredRoots     []string
-	versionsMap         map[string]ast.RegoVersion
-}
-
 // SetOnConflictOperation sets the fixer's behavior when a conflict occurs.
 func (f *Fixer) SetOnConflictOperation(operation OnConflictOperation) {
 	f.onConflictOperation = operation