Completions: support import aliases in rulerefs provider

[charlesdaniels]

Consider the following example:

# foo/bar/bar.rego
package foo.bar

hello := "world"

# a/b/b.rego
package a.b
import data.foo.bar as baz

allow if {
    baz.h # tab complete here
}

Currently, this will not work, as the LSP is not import alias aware. However, right now if the as baz was removed from the import statement, then bar.h would correctly complete to bar.hello.

It would be desirable for Regal to also support import aliases, in addition to normal imports which are already supported.

Background

I have investigated a little in terms of what would be required to support this, and my findings are outlined below.

bundle/regal/lsp/completion/providers/rulerefs/rulerefs.rego is where the relevant code lives. Currently, the list of imports for the current file is determined by _current_file_imports. This line:

	some imp in _parsed_current_file.imports

Will iterate over values that look similar to these:

{"location": "3:1:aW1wb3J0", "path": {"location": "3:8:cmVnby52MQ==", "type": "ref", "value": [{"location": "3:8:cmVnbw==", "type": "var", "value": "rego"}, {"location": "3:13:djE=", "type": "string", "value": "v1"}]}}

{"alias": "s", "location": "4:1:aW1wb3J0", "path": {"location": "4:8:ZGF0YS5mb28uYmFyLnE=", "type": "ref", "value": [{"location": "4:8:ZGF0YQ==", "type": "var", "value": "data"}, {"location": "4:13:Zm9v", "type": "string", "value": "foo"}, {"location": "4:17:YmFy", "type": "string", "value": "bar"}, {"location": "4:21:cQ==", "type": "string", "value": "q"}]}}

Tip

To be able to see print statements from this policy during testing, you can add regoArgs = append(regoArgs, rego.EnablePrintStatements(true)) into prepareQuery in internal/lsp/completions/providers/policy.go. This method is how the samples above were obtained.

In the case of imports which have an as alias, a top level alias field will be present in the import object. A naive solution might be to simply mangle the value passed to ast.ref_to_string(pathval) to reflect the alias values, perhaps like this:

_current_file_imports contains ref if {
	some imp in _parsed_current_file.imports

	pathval := [
		v
		|
		some i, elem in imp.path.value ;
		v := {
			false: elem,
			true: object.union(elem, {"value": object.get(imp, "alias", elem)}),
		}[i == (count(imp.path.value)-1)]
	]

	ref := ast.ref_to_string(pathval)}

The trouble comes in when matching imported references in _imported_package_refs. There, ref would take on a value like bar.hello, because this accurately reflects the terminal path element of the package (bar) where hello was declared. This means that when baz.h is encountered, it will not in fact share a prefix with bar.hello.

To properly address this, the _workspace_rule_refs rule will need to be modified to rewrite references for imported packages that have aliases to accurately reflect those aliases, this way when _imported_package_refs prefix matches against the workspace references, the prefixes will be correct.

It is possible that other considerations may come up as well, this only reflects the problems I have encountered so far.