Completion suggestions for variables in local scope (#840)

This is the first provider that is using Rego policy to determine suggestions!

That makes the PR more extensive than normally, but much of this will be reusable
for future providers, or even existing ones we may choose to convert.

Fixes #792

Signed-off-by: Anders Eknert <[email protected]>

bundle/regal/ast/ast.rego

@@ -291,16 +291,24 @@ _before_location(var, location) if {
 	var.location.col < location.col
 }
 
+# METADATA
+# description: find *only* names in the local scope, and not e.g. rule names
+find_names_in_local_scope(rule, location) := names if {
+	fn_arg_names := _function_arg_names(rule)
+	var_names := {var.value | some var in find_vars_in_local_scope(rule, location)}
+
+	names := fn_arg_names | var_names
+}
+
 # METADATA
 # description: |
 #   similar to `find_vars_in_local_scope`, but returns all variable names in scope
 #   of the given location *and* the rule names present in the scope (i.e. module)
 find_names_in_scope(rule, location) := names if {
-	fn_arg_names := _function_arg_names(rule)
-	var_names := {var.value | some var in find_vars_in_local_scope(rule, location)}
+	locals := find_names_in_local_scope(rule, location)
 
 	# parens below added by opa-fmt :)
-	names := ((rule_names | imported_identifiers) | fn_arg_names) | var_names
+	names := (rule_names | imported_identifiers) | locals
 }
 
 # METADATA

bundle/regal/lsp/completion/kind.rego

@@ -0,0 +1,53 @@
+package regal.lsp.completion.kind
+
+import rego.v1
+
+text := 1
+
+method := 2
+
+function := 3
+
+constructor := 4
+
+field := 5
+
+variable := 6
+
+class := 7
+
+interface := 8
+
+module := 9
+
+property := 10
+
+unit := 11
+
+value := 12
+
+enum := 13
+
+keyword := 14
+
+snippet := 15
+
+color := 16
+
+file := 17
+
+reference := 18
+
+folder := 19
+
+enum_member := 20
+
+constant := 21
+
+struct := 22
+
+event := 23
+
+operator := 24
+
+type_parameter := 25

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

@@ -0,0 +1,56 @@
+# METADATA
+# description: various rules and functions related to location and position
+package regal.lsp.completion.location
+
+import rego.v1
+
+import data.regal.ast
+
+# METADATA
+# description: best-effort helper to determine if the current line is in a rule body
+# scope: document
+in_rule_body(line) if contains(line, " if ")
+
+in_rule_body(line) if contains(line, " contains ")
+
+in_rule_body(line) if contains(line, " else ")
+
+in_rule_body(line) if contains(line, "= ")
+
+in_rule_body(line) if regex.match(`^\s+`, line)
+
+# METADATA
+# description: converts OPA location to LSP position
+to_position(location) := {
+	"line": location.row - 1,
+	"character": location.col - 1,
+}
+
+# METADATA
+# description: |
+#   estimate where the location "ends" based on its text attribute,
+#   both line and column
+end_location_estimate(location) := end if {
+	lines := split(base64.decode(location.text), "\n")
+	end := {
+		"row": (location.row + count(lines)) - 1,
+		"col": count(regal.last(lines)),
+	}
+}
+
+# METADATA
+# description: |
+#   find and return rule at provided location
+#   undefined if provided location is not within the range of a rule
+find_rule(rules, location) := [rule |
+	some i, rule in rules
+	end_location := end_location_estimate(rule.location)
+	location.row >= rule.location.row
+	location.row <= end_location.row
+][0]
+
+# METADATA
+# description: |
+#   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)

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

@@ -0,0 +1,63 @@
+package regal.lsp.completion.providers.location_test
+
+import rego.v1
+
+import data.regal.ast
+
+import data.regal.lsp.completion.location
+
+test_find_rule_from_location if {
+	module := regal.parse_module("p.rego", `package p
+
+import rego.v1
+
+rule1 if {
+	x := 1
+}
+
+rule2 if {
+	y := 2
+}
+
+rule3 if {
+	z := 3
+}
+`)
+	not location.find_rule(module.rules, {"row": 2, "col": 6})
+
+	r1 := location.find_rule(module.rules, {"row": 5, "col": 6})
+	ast.ref_to_string(r1.head.ref) == "rule1"
+
+	r2 := location.find_rule(module.rules, {"row": 9, "col": 11})
+	ast.ref_to_string(r2.head.ref) == "rule2"
+
+	r3 := location.find_rule(module.rules, {"row": 15, "col": 0})
+	ast.ref_to_string(r3.head.ref) == "rule3"
+}
+
+test_find_locals_at_location if {
+	module := regal.parse_module("p.rego", `package p
+
+import rego.v1
+
+rule if {
+	x := 1
+}
+
+function(a, b) if {
+	c := 3
+}
+
+another if {
+	some x, y in collection
+	z := x + y
+}
+`)
+
+	location.find_locals(module.rules, {"row": 6, "col": 1}) == set()
+	location.find_locals(module.rules, {"row": 6, "col": 10}) == {"x"}
+	location.find_locals(module.rules, {"row": 10, "col": 1}) == {"a", "b"}
+	location.find_locals(module.rules, {"row": 10, "col": 6}) == {"a", "b", "c"}
+	location.find_locals(module.rules, {"row": 15, "col": 1}) == {"x", "y"}
+	location.find_locals(module.rules, {"row": 16, "col": 1}) == {"x", "y", "z"}
+}

bundle/regal/lsp/completion/main.rego

@@ -0,0 +1,13 @@
+# METADATA
+# description: |
+#   base package for completion suggestion provider policies, and acts
+#   like a router that'll collection suggestions from all provider policies
+#   under regal.lsp.completion.providers
+package regal.lsp.completion
+
+import rego.v1
+
+# METADATA
+# description: main entry point for completion suggestions
+# entrypoint: true
+items contains data.regal.lsp.completion.providers[_].items[_]

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

@@ -0,0 +1,36 @@
+package regal.lsp.completion.providers.locals
+
+import rego.v1
+
+import data.regal.lsp.completion.kind
+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]
+	line != ""
+	location.in_rule_body(line)
+
+	last_word := regal.last(regex.split(`\s+`, trim_space(line)))
+
+	some local in location.find_locals(input.rules, input.regal.context.location)
+
+	startswith(local, last_word)
+
+	item := {
+		"label": local,
+		"kind": kind.variable,
+		"detail": "local variable",
+		"textEdit": {
+			"range": {
+				"start": {
+					"line": position.line,
+					"character": position.character - count(last_word),
+				},
+				"end": position,
+			},
+			"newText": local,
+		},
+	}
+}

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

@@ -0,0 +1,78 @@
+package regal.lsp.completion.providers.locals_test
+
+import rego.v1
+
+import data.regal.lsp.completion.providers.locals
+
+test_no_locals_in_completion_items if {
+	policy := `package policy
+
+import rego.v1
+
+foo := 1
+
+bar if {
+	foo == 1
+}
+`
+
+	module := regal.parse_module("p.rego", policy)
+	regal_module := object.union(module, {"regal": {
+		"file": {
+			"name": "p.rego",
+			"lines": split(policy, "\n"),
+		},
+		"context": {"location": {
+			"row": 8,
+			"col": 9,
+		}},
+	}})
+	items := locals.items with input as regal_module
+
+	count(items) == 0
+}
+
+test_locals_in_completion_items if {
+	policy := `package policy
+
+import rego.v1
+
+foo := 1
+
+function(bar) if {
+	baz := 1
+	qux := b
+}
+`
+
+	module := object.union(regal.parse_module("p.rego", policy), {"regal": {
+		"file": {
+			"name": "p.rego",
+			"lines": split(policy, "\n"),
+		},
+		"context": {"location": {
+			"row": 9,
+			"col": 7,
+		}},
+	}})
+	items := locals.items with input as module
+
+	count(items) == 2
+
+	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}})
+}
+
+expect_item(items, label, range) if {
+	expected := {"detail": "local variable", "kind": 6}
+
+	item := object.union(expected, {
+		"label": label,
+		"textEdit": {
+			"newText": label,
+			"range": range,
+		},
+	})
+
+	item in items
+}

internal/embeds/schemas/regal-ast.json

@@ -395,6 +395,10 @@
             }
           },
           "type": "object"
+        },
+        "context": {
+          "description": "extra attributes provided in the specific evaluation context",
+          "type": "object"
         }
       },
       "type": "object",

internal/lsp/completions/manager.go

@@ -42,6 +42,8 @@ func NewDefaultManager(c *cache.Cache) *Manager {
 	m.RegisterProvider(&providers.CommonRule{})
 	m.RegisterProvider(&providers.UsedRefs{})
 
+	m.RegisterProvider(providers.NewPolicy())
+
 	return m
 }
 

internal/lsp/completions/providers/builtins.go

@@ -42,7 +42,7 @@ func (*BuiltIns) Run(c *cache.Cache, params types.CompletionParams, _ *Options)
 			items = append(items, types.CompletionItem{
 				Label:  key,
 				Kind:   completion.Function,
-				Detail: "",
+				Detail: "built-in function",
 				Documentation: &types.MarkupContent{
 					Kind:  "markdown",
 					Value: hover.CreateHoverContent(builtIn),

internal/lsp/completions/providers/input.go

@@ -54,9 +54,7 @@ You can also experiment with input in the [Rego Playground](https://play.openpol
 						Line:      params.Position.Line,
 						Character: params.Position.Character - uint(len(lastWord)),
 					},
-					End: types.Position{
-						Line: params.Position.Line, Character: params.Position.Character,
-					},
+					End: params.Position,
 				},
 				NewText: "input",
 			},