Completion of locals in more places (#847)

Locals are now suggested anywhere inside of a word boundary, including
function a call argument list, map keys/values, arrays, etc.

Added a couple of useful helpers as part of this that we should be able
to use for other completion providers too.

Fixes #844

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

bundle/regal/lsp/completion/location/location.rego

@@ -51,6 +51,42 @@ find_rule(rules, location) := [rule |
 
 # METADATA
 # description: |
-#   find local variables (declared via function arguments, some/every declarations or assignment)
-#   at the given location
+#   find local variables (declared via function arguments, some/every
+#   declarations or assignment) at the given location
 find_locals(rules, location) := ast.find_names_in_local_scope(find_rule(rules, location), location)
+
+# METADATA
+# description: |
+#   return the range for a word object (as return by `word_at`)
+word_range(word, position) := {
+	"start": {
+		"line": position.line,
+		"character": position.character - word.offset_before,
+	},
+	"end": {
+		"line": position.line,
+		"character": position.character + word.offset_after,
+	},
+}
+
+# METADATA
+# description: |
+#   find word at column in line, and return its text, and the offset
+#   from the position (before and after)
+word_at(line, col) := word if {
+	text_before := substring(line, 0, col - 1)
+	word_before := _to_string(regex.find_n(`[a-zA-Z_]+$`, text_before, 1))
+
+	text_after := substring(line, col - 1, count(line))
+	word_after := _to_string(regex.find_n(`^[a-zA-Z_]+`, text_after, 1))
+
+	word := {
+		"offset_before": count(word_before),
+		"offset_after": count(word_after),
+		"text": sprintf("%s%s", [word_before, word_after]),
+	}
+}
+
+_to_string(arr) := "" if count(arr) == 0
+
+_to_string(arr) := arr[0] if count(arr) > 0

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

@@ -12,25 +12,31 @@ items contains item if {
 	line != ""
 	location.in_rule_body(line)
 
-	last_word := regal.last(regex.split(`\s+`, trim_space(line)))
+	word := location.word_at(line, input.regal.context.location.col)
+
+	not excluded(line, position)
 
 	some local in location.find_locals(input.rules, input.regal.context.location)
 
-	startswith(local, last_word)
+	startswith(local, word.text)
 
 	item := {
 		"label": local,
 		"kind": kind.variable,
 		"detail": "local variable",
 		"textEdit": {
-			"range": {
-				"start": {
-					"line": position.line,
-					"character": position.character - count(last_word),
-				},
-				"end": position,
-			},
+			"range": location.word_range(word, position),
 			"newText": local,
 		},
 	}
 }
+
+# exclude local suggestions in function args definition,
+# as those would recursively contribute to themselves
+excluded(line, position) if _function_args_position(substring(line, 0, position.character))
+
+_function_args_position(text) if {
+	text == trim_left(text, " \t")
+	contains(text, "(")
+	not contains(text, "=")
+}

bundle/regal/lsp/completion/providers/locals/locals_test.rego

@@ -52,15 +52,43 @@ function(bar) if {
 		},
 		"context": {"location": {
 			"row": 9,
-			"col": 7,
+			"col": 10,
 		}},
 	}})
 	items := locals.items with input as module
 
 	count(items) == 2
+	expect_item(items, "bar", {"end": {"character": 9, "line": 8}, "start": {"character": 8, "line": 8}})
+	expect_item(items, "baz", {"end": {"character": 9, "line": 8}, "start": {"character": 8, "line": 8}})
+}
+
+test_locals_in_completion_items_function_call if {
+	policy := `package policy
+
+import rego.v1
+
+foo := 1
 
-	expect_item(items, "bar", {"end": {"character": 6, "line": 8}, "start": {"character": 5, "line": 8}})
-	expect_item(items, "baz", {"end": {"character": 6, "line": 8}, "start": {"character": 5, "line": 8}})
+function(bar) if {
+	baz := 1
+	qux := other_function(b)
+}
+`
+	module := object.union(regal.parse_module("p.rego", policy), {"regal": {
+		"file": {
+			"name": "p.rego",
+			"lines": split(policy, "\n"),
+		},
+		"context": {"location": {
+			"row": 9,
+			"col": 25,
+		}},
+	}})
+	items := locals.items with input as module
+
+	count(items) == 2
+	expect_item(items, "bar", {"end": {"character": 24, "line": 8}, "start": {"character": 23, "line": 8}})
+	expect_item(items, "baz", {"end": {"character": 24, "line": 8}, "start": {"character": 23, "line": 8}})
 }
 
 expect_item(items, label, range) if {