Completions: rulerefs optimizations (StyraInc#898)

This one can be quite heavy to process as there are potentially thousands
of suggestions if typing in just `data`. Filter out tests from the list,
and let the client do the sorting. Also some minor optimizations in Rego
code processing.

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

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

@@ -3,6 +3,7 @@ package regal.lsp.completion.providers.rulerefs
 import rego.v1
 
 import data.regal.ast
+
 import data.regal.lsp.completion.kind
 import data.regal.lsp.completion.location
 
@@ -25,9 +26,7 @@ workspace_rule_refs contains ref if {
 
 parsed_current_file := data.workspace.parsed[input.regal.file.uri]
 
-current_file_package := concat(".", [segment.value |
-	some segment in parsed_current_file["package"].path
-])
+current_file_package := ast.ref_to_string(parsed_current_file["package"].path)
 
 current_file_imports contains ref if {
 	some imp in parsed_current_file.imports
@@ -43,18 +42,19 @@ current_package_refs contains ref if {
 
 imported_package_refs contains ref if {
 	some ref in workspace_rule_refs
-	some pkg_ref in current_file_imports
 
 	not ref_is_internal(ref)
 
-	startswith(ref, pkg_ref)
+	strings.any_prefix_match(ref, current_file_imports)
 }
 
 other_package_refs contains ref if {
 	some ref in workspace_rule_refs
 
 	not ref in imported_package_refs
 	not ref in current_package_refs
+
+	not ref_is_internal(ref)
 }
 
 # from the current package
@@ -71,63 +71,47 @@ rule_ref_suggestions contains pkg_ref if {
 
 	startswith(ref, imported_package)
 
-	parts := split(imported_package, ".")
-	prefix := concat(".", array.slice(parts, 0, count(parts) - 1))
+	prefix := regex.replace(imported_package, `\.[^\.]+$`, "")
 	pkg_ref := trim_prefix(ref, sprintf("%s.", [prefix]))
 }
 
 # from any other package
-rule_ref_suggestions contains ref if {
-	some ref in other_package_refs
-
-	not ref_is_internal(ref)
-}
+rule_ref_suggestions contains ref if some ref in other_package_refs
 
 # also suggest the unimported packages themselves
 # e.g. data.foo.rule will also generate data.foo as a suggestion
 rule_ref_suggestions contains pkg if {
 	some ref in other_package_refs
 
-	not ref_is_internal(ref)
-
-	parts := split(ref, ".")
-	pkg := concat(".", array.slice(parts, 0, count(parts) - 1))
+	pkg := regex.replace(ref, `\.[^\.]+$`, "")
 }
 
 matching_rule_ref_suggestions contains ref if {
 	line != ""
 	location.in_rule_body(line)
 
+	# \W is used here to match ( in the case of func() := ..., as well as the space in the case of rule := ...
+	first_word := regex.split(`\W+`, trim_space(line))[0]
 	prefix := determine_ref_prefix(word.text)
 
 	some ref in rule_ref_suggestions
 
 	startswith(ref, prefix)
 
-	# \W is used here to match ( in the case of func() := ..., as well as the space in the case of rule := ...
-	first_word := regex.split(`\W+`, trim_space(line))[0]
-
 	# this is to avoid suggesting a recursive rule, e.g. rule := rule, or func() := func()
 	ref != first_word
 }
 
-grouped_refs[size] contains ref if {
+items contains item if {
 	some ref in matching_rule_ref_suggestions
 
-	size := count(indexof_n(ref, "."))
-}
-
-items := [item |
-	some group in grouped_refs
-	some ref in sort(group)
-
 	item := {
 		"label": ref,
 		"kind": kind.variable,
-		"detail": "rule ref",
+		"detail": "reference",
 		"textEdit": {
 			"range": location.word_range(word, position),
 			"newText": ref,
 		},
 	}
-]
+}

bundle/regal/lsp/completion/providers/rulerefs/rulerefs_test.rego

@@ -44,9 +44,7 @@ parsed_modules[file_uri] := parsed_module if {
 defined_refs[file_uri] contains ref if {
 	some file_uri, parsed_module in parsed_modules
 
-	package_name := concat(".", [path.value |
-		some i, path in parsed_module["package"].path
-	])
+	package_name := ast.ref_to_string(parsed_module["package"].path)
 
 	some rule in parsed_module.rules
 
@@ -75,15 +73,15 @@ another_local_rule := `])
 		with data.workspace.parsed as parsed_modules
 		with data.workspace.defined_refs as defined_refs
 
-	labels := [item.label | some item in items]
+	labels := {item.label | some item in items}
 
-	expected_refs := [
+	expected_refs := {
 		"local_rule",
 		"imported_pkg.another_rule",
 		"imported_pkg_2.another_rule_2",
 		"data.not_imported_pkg.foo.bar", # partial generated from rule below
 		"data.not_imported_pkg.foo.bar.yet_another_rule",
-	]
+	}
 
 	expected_refs == labels
 }

internal/lsp/completions/refs/defined.go

@@ -42,6 +42,11 @@ func DefinedInModule(module *ast.Module) map[string]types.Ref {
 
 	for _, rule := range module.Rules {
 		name := rast.RefToString(rule.Head.Ref())
+
+		if strings.HasPrefix(name, "test_") {
+			continue
+		}
+
 		ruleGroups[name] = append(ruleGroups[name], rule)
 	}