Rule: defer-assignment (#1215)

Finally a new rule in the `performance` category! And this is quite
a fun one. I was curious to see how many violations this would find
in the Regal codebase, if any. It did find a few, and good ones to
fix!

In order to keep it simple, this first implementation is perhaps a
bit too cautious about what it considers deferrable. Check the
implementation to see what exactly, but pretty much only assignments
where the next line is a not an assignment, a potential loop, or
something else that *might* make deferring the assignment undesirable.

Fixes #1147

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

README.md

@@ -253,6 +253,7 @@ The following rules are currently available:
 | imports     | [redundant-data-import](https://docs.styra.com/regal/rules/imports/redundant-data-import)               | Redundant import of data                                  |
 | imports     | [unresolved-import](https://docs.styra.com/regal/rules/imports/unresolved-import)                       | Unresolved import                                         |
 | imports     | [use-rego-v1](https://docs.styra.com/regal/rules/imports/use-rego-v1)                                   | Use `import rego.v1`                                      |
+| performance | [defer-assignment](https://docs.styra.com/regal/rules/performance/defer-assignment)                     | Assignment can be deferred                                |
 | performance | [with-outside-test-context](https://docs.styra.com/regal/rules/performance/with-outside-test-context)   | `with` used outside test context                          |
 | style       | [avoid-get-and-list-prefix](https://docs.styra.com/regal/rules/style/avoid-get-and-list-prefix)         | Avoid `get_` and `list_` prefix for rules and functions   |
 | style       | [chained-rule-body](https://docs.styra.com/regal/rules/style/chained-rule-body)                         | Avoid chaining rule bodies                                |

bundle/regal/ast/ast.rego

@@ -363,13 +363,36 @@ is_chained_rule_body(rule, lines) if {
 }
 
 # METADATA
-# description: returns the terms in an assignment expression, or undefined if not assignment
-assignment_terms(expr) := [expr.terms[1], expr.terms[2]] if {
+# description: answers wether variable of `name` is found anywhere in provided rule `head`
+# scope: document
+var_in_head(head, name) if {
+	head.value.value == name
+} else if {
+	head.key.value == name
+} else if {
+	some var in find_term_vars(head.value.value)
+	var.value == name
+} else if {
+	some var in find_term_vars(head.key.value)
+	var.value == name
+} else if {
+	some i, var in head.ref
+	i > 0
+	var.value == name
+}
+
+# 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"
 }
 
+# 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)
+
 # METADATA
 # description: |
 #   For a given rule head name, this rule contains a list of locations where

bundle/regal/ast/search.rego

@@ -228,7 +228,7 @@ found.symbols[rule_index] contains value.symbols if {
 }
 
 # METADATA
-# description: all comprehensions foundd in module
+# description: all comprehensions found in module
 found.comprehensions[rule_index] contains value if {
 	some i, rule in _rules
 
@@ -268,10 +268,12 @@ _before_location(rule, _, location) if {
 	loc := util.to_location_object(location)
 
 	value_start := util.to_location_object(rule.head.value.location)
-	value_end := _end_location(util.to_location_object(rule.head.value.location))
 
 	loc.row >= value_start.row
 	loc.col >= value_start.col
+
+	value_end := _end_location(util.to_location_object(rule.head.value.location))
+
 	loc.row <= value_end.row
 	loc.col <= value_end.col
 }

bundle/regal/config/provided/data.yaml

@@ -118,6 +118,8 @@ rules:
     use-rego-v1:
       level: error
   performance:
+    defer-assignment:
+      level: error
     with-outside-test-context:
       level: error
   style:
@@ -159,8 +161,8 @@ rules:
     prefer-snake-case:
       level: error
     prefer-some-in-iteration:
-      ignore-nesting-level: 2
       ignore-if-sub-attribute: true
+      ignore-nesting-level: 2
       level: error
     rule-length:
       count-comments: false

bundle/regal/lsp/completion/providers/import/import.rego

@@ -12,10 +12,11 @@ import data.regal.lsp.completion.location
 items contains item if {
 	position := location.to_position(input.regal.context.location)
 	line := input.regal.file.lines[position.line]
-	word := location.word_at(line, input.regal.context.location.col)
 
 	startswith("import", line)
 
+	word := location.word_at(line, input.regal.context.location.col)
+
 	item := {
 		"label": "import",
 		"kind": kind.keyword,

bundle/regal/lsp/completion/providers/snippet/snippet.rego

@@ -19,10 +19,11 @@ import data.regal.lsp.completion.location
 items contains item if {
 	position := location.to_position(input.regal.context.location)
 	line := input.regal.file.lines[position.line]
-	word := location.word_at(line, input.regal.context.location.col)
 
 	location.in_rule_body(line)
 
+	word := location.word_at(line, input.regal.context.location.col)
+
 	some label, snippet in _snippets
 
 	strings.any_prefix_match(snippet.prefix, word.text)

bundle/regal/result/result.rego

@@ -120,19 +120,17 @@ fail(metadata, details) := _fail_annotated(metadata, details)
 #   ignored because the set capabilities does not include a dependency (like a built-in function)
 #   needed by the rule
 notice(metadata) := result if {
-	is_array(metadata)
 	rule_meta := metadata[0]
-	annotations := rule_meta.annotations
 
 	some category, title
 	["regal", "rules", category, title, "notices"] = rule_meta.path
 
 	result := {
 		"category": category,
-		"description": annotations.description,
+		"description": rule_meta.annotations.description,
 		"level": "notice",
 		"title": title,
-		"severity": annotations.custom.severity,
+		"severity": rule_meta.annotations.custom.severity,
 	}
 }
 

bundle/regal/rules/bugs/sprintf-arguments-mismatch/sprintf_arguments_mismatch.rego

@@ -24,6 +24,8 @@ report contains violation if {
 	some rule_index, fn
 	ast.function_calls[rule_index][fn].name == "sprintf"
 
+	fn.args[1].type == "array" # can only check static arrays, not vars
+
 	# this could come either from a term directly (the common case):
 	#     sprintf("%d", [1])
 	# or a variable (quite uncommon):
@@ -35,8 +37,6 @@ report contains violation if {
 	# to not report anything erroneously.
 	format_term := _first_arg_value(rule_index, fn.args[0])
 
-	fn.args[1].type == "array" # can only check static arrays, not vars
-
 	values_in_arr := count(fn.args[1].value)
 	str_no_escape := replace(format_term.value, "%%", "") # don't include '%%' as it's used to "escape" %
 	values_in_str := strings.count(str_no_escape, "%") - _repeated_explicit_argument_indexes(str_no_escape)

bundle/regal/rules/bugs/unused-output-variable/unused_output_variable.rego

@@ -29,7 +29,7 @@ report contains violation if {
 
 	some var in var_refs
 
-	not _var_in_head(input.rules[to_number(rule_index)].head, var.value)
+	not ast.var_in_head(input.rules[to_number(rule_index)].head, var.value)
 	not _var_in_call(ast.function_calls, rule_index, var.value)
 	not _ref_base_vars[rule_index][var.value]
 
@@ -56,28 +56,6 @@ _ref_base_vars[rule_index][term.value] contains term if {
 	not startswith(term.value, "$")
 }
 
-_var_in_head(head, name) if head.value.value == name
-
-_var_in_head(head, name) if {
-	some var in ast.find_term_vars(head.value.value)
-
-	var.value == name
-}
-
-_var_in_head(head, name) if head.key.value == name
-
-_var_in_head(head, name) if {
-	some var in ast.find_term_vars(head.key.value)
-
-	var.value == name
-}
-
-_var_in_head(head, name) if {
-	some i, var in head.ref
-	i > 0
-	var.value == name
-}
-
 _var_in_call(calls, rule_index, name) if _var_in_arg(calls[rule_index][_].args[_], name)
 
 _var_in_arg(arg, name) if {

bundle/regal/rules/idiomatic/custom-has-key-construct/custom_has_key_construct.rego

@@ -17,8 +17,6 @@ notices contains result.notice(rego.metadata.chain()) if not capabilities.has_ob
 report contains violation if {
 	some rule in ast.functions
 
-	arg_names := ast.function_arg_names(rule)
-
 	count(rule.body) == 1
 
 	terms := rule.body[0].terms
@@ -29,6 +27,9 @@ report contains violation if {
 	[_, ref] := _normalize_eq_terms(terms)
 
 	ref.value[0].type == "var"
+
+	arg_names := ast.function_arg_names(rule)
+
 	ref.value[0].value in arg_names
 	ref.value[1].type == "var"
 	ref.value[1].value in arg_names

bundle/regal/rules/idiomatic/custom-in-construct/custom_in_construct.rego

@@ -10,8 +10,6 @@ import data.regal.result
 report contains violation if {
 	some rule in ast.functions
 
-	arg_names := ast.function_arg_names(rule)
-
 	# while there could be more convoluted ways of doing this
 	# we'll settle for the likely most common case (`item == coll[_]`)
 	count(rule.body) == 1
@@ -23,6 +21,8 @@ report contains violation if {
 
 	[var, ref] := _normalize_eq_terms(terms)
 
+	arg_names := ast.function_arg_names(rule)
+
 	var.value in arg_names
 	ref.value[0].value in arg_names
 	ref.value[1].type == "var"

bundle/regal/rules/idiomatic/equals-pattern-matching/equals_pattern_matching.rego

@@ -45,11 +45,6 @@ report contains violation if {
 	fn.body
 	not fn["else"]
 
-	arg_var_names := {arg.value |
-		some arg in fn.head.args
-		arg.type == "var"
-	}
-
 	# FOR NOW: Limit to a lone comparison
 	# More elaborate cases are certainly doable,
 	# but we'd need to keep track of whatever else
@@ -64,6 +59,11 @@ report contains violation if {
 	expr.terms[0].value[0].value == "equal"
 
 	terms := _normalize_eq_terms(expr.terms, ast.scalar_types)
+	arg_var_names := {arg.value |
+		some arg in fn.head.args
+		arg.type == "var"
+	}
+
 	terms[0].value in arg_var_names
 
 	violation := result.fail(rego.metadata.chain(), result.location(fn))

bundle/regal/rules/performance/defer-assignment/defer_assignment.rego

@@ -0,0 +1,97 @@
+# METADATA
+# description: Assignment can be deferred
+package regal.rules.performance["defer-assignment"]
+
+import rego.v1
+
+import data.regal.ast
+import data.regal.result
+
+report contains violation if {
+	some i, rule in input.rules
+	some j, expr in rule.body
+
+	[var, rhs] := ast.assignment_terms(expr)
+
+	not _ref_with_vars(rhs)
+
+	# for now, only simple var assignment counts.. later we can
+	# consider checking the contents of arrays here
+	var.type == "var"
+
+	next := rule.body[j + 1]
+
+	not ast.is_assignment(next)
+	not ast.var_in_head(rule.head, var.value)
+	not _var_used_in_expression(var, next)
+	not _iteration_expression(next)
+	not _print_call(next)
+
+	violation := result.fail(rego.metadata.chain(), result.location(expr))
+}
+
+_ref_with_vars(node) if {
+	node.type == "ref"
+
+	some i, term in node.value
+
+	i > 0
+	term.type == "var"
+}
+
+_var_used_in_expression(var, expr) if {
+	not expr.terms.symbols
+
+	is_array(expr.terms)
+
+	some term in expr.terms
+
+	walk(term, [_, value])
+
+	value.type == "var"
+	value.value == var.value
+}
+
+_var_used_in_expression(var, expr) if {
+	some w in expr["with"]
+
+	w.value.type == "var"
+	w.value.value == var.value
+}
+
+_var_used_in_expression(var, expr) if {
+	# `not x`
+	is_object(expr.terms)
+
+	expr.terms.type == "var"
+	expr.terms.value == var.value
+} else if {
+	# `not x.y`
+	is_object(expr.terms)
+
+	some term in expr.terms.value
+
+	walk(term, [_, value])
+
+	value.type == "var"
+	value.value == var.value
+}
+
+# while not technically checking of use here:
+# the next expression having symbols indicate iteration, and
+# we don't want to defer assignment into a loop
+_iteration_expression(expr) if expr.terms.symbols
+
+# likewise with every
+_iteration_expression(expr) if expr.terms.domain
+
+# and walk
+_iteration_expression(expr) if {
+	expr.terms[0].value[0].type == "var"
+	expr.terms[0].value[0].value == "walk"
+}
+
+_print_call(expr) if {
+	expr.terms[0].value[0].type == "var"
+	expr.terms[0].value[0].value == "print"
+}