Package ref import suggestions in Rego (#1698)

Rewrite the last Go completion provider in Rego!

The next step is to get tid of the now pointless Manager/Policy
Go completion handlers, and integrate Rego completions in the main
Rego router. This is a little more involved though, as the input
format currently used for completion providers is quite different
from what other Rego handlers are provided. Which in itself is a
good reason for why we'll want to fix it... but will follow up with
another PR for that later, so that this can be reviewed in isolation.

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

Author: anderseknert

bundle/regal/lsp/completion/main.rego

@@ -5,6 +5,11 @@
 #   under regal.lsp.completion.providers
 package regal.lsp.completion
 
+import data.regal.util
+
+# regal ignore:pointless-import
+import data.regal.lsp.completion.location
+
 # METADATA
 # entrypoint: true
 result["response"] := items
@@ -13,6 +18,28 @@ result["response"] := items
 # description: main entry point for completion suggestions
 # entrypoint: true
 items contains object.union(completion, {"_regal": {"provider": provider}}) if {
+	# exit early if caret position is inside a comment. We currently don't have any provider
+	# where doing completions inside of a comment makes sense. Behavior is also editor-specific:
+	# - Zed: always on, with no way to disable
+	# - VSCode: disabled but can be enabled with "editor.quickSuggestions.comments" setting
+	not inside_comment
+
 	some provider, completion
 	data.regal.lsp.completion.providers[provider].items[completion]
 }
+
+# METADATA
+# description: |
+#   checks if the current position is inside a comment
+inside_comment if {
+	position := location.to_position(input.regal.context.location)
+
+	# avoid unmarshalling every comment location but only one that starts
+	# with the line number of the current position
+	line := sprintf("%d:", [position.line + 1])
+
+	some comment in data.workspace.parsed[input.regal.file.uri].comments
+
+	startswith(comment.location, line)
+	util.to_location_no_text(comment.location).col <= position.character + 1
+}

bundle/regal/lsp/completion/main_test.rego

@@ -7,3 +7,29 @@ test_completion_entrypoint if {
 
 	items == {{"_regal": {"provider": "test"}, "foo": "bar"}}
 }
+
+test_inside_comment if {
+	_data := {"file:///example.rego": {"comments": [
+		{"location": "2:1:2:10"},
+		{"location": "4:1:4:10"},
+	]}}
+	_input := {"regal": {
+		"context": {"location": {"row": 4, "col": 5}},
+		"file": {"uri": "file:///example.rego"},
+	}}
+
+	completion.inside_comment with input as _input with data.workspace.parsed as _data
+}
+
+test_not_inside_comment if {
+	_data := {"file:///example.rego": {"comments": [
+		{"location": "2:1:2:10"},
+		{"location": "4:8:4:10"},
+	]}}
+	_input := {"regal": {
+		"context": {"location": {"row": 4, "col": 5}},
+		"file": {"uri": "file:///example.rego"},
+	}}
+
+	not completion.inside_comment with input as _input with data.workspace.parsed as _data
+}

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

@@ -0,0 +1,64 @@
+# METADATA
+# description: Completion suggestions for importing available packages
+package regal.lsp.completion.providers.packagerefs
+
+import data.regal.ast
+
+import data.regal.lsp.completion.kind
+import data.regal.lsp.completion.location
+
+# METADATA
+# description: suggest packages matching typed import ref
+items contains item if {
+	position := location.to_position(input.regal.context.location)
+	line := input.regal.file.lines[position.line]
+
+	startswith(line, "import ")
+
+	ref := location.ref_at(line, input.regal.context.location.col)
+
+	startswith(ref.text, "d")
+
+	some i, path in _paths_sorted
+
+	startswith(path, ref.text)
+
+	item := {
+		"label": path,
+		"kind": kind.module,
+		"detail": "package",
+		"textEdit": {
+			"range": location.word_range(ref, position),
+			"newText": path,
+		},
+		# tell clients to sort paths first by the number of path components (shortest first),
+		# and only then alphabetically (done in _paths_by_num_parts and _paths_sorted)
+		"sortText": sprintf("%03d", [i]),
+	}
+}
+
+_package_paths contains str if {
+	some uri
+	path := data.workspace.parsed[uri].package.path
+
+	uri != input.regal.file.uri # don't suggest the package of the current file
+	not endswith(regal.last(path).value, "_test") # importing tests makes no sense
+
+	str := ast.ref_to_string(path)
+}
+
+# regal ignore:prefer-set-or-object-rule
+_paths_by_num_parts := {num_parts: sort(paths) |
+	some i
+	num_parts := strings.count(_package_paths[i], ".")
+	paths := [path |
+		some j
+		strings.count(_package_paths[j], ".") == num_parts
+		path := _package_paths[j]
+	]
+}
+
+_paths_sorted := [path |
+	some i in sort(object.keys(_paths_by_num_parts))
+	some path in _paths_by_num_parts[i]
+]

bundle/regal/lsp/completion/providers/packagerefs/packagerefs_test.rego

@@ -0,0 +1,75 @@
+package regal.lsp.completion.providers.packagerefs_test
+
+import data.regal.lsp.completion.providers.packagerefs
+
+test_all_package_refs_sugggested_for_import if {
+	items := packagerefs.items with data.workspace.parsed as _workspace_parsed with input as {"regal": {
+		"context": {"location": {"row": 3, "col": 9}},
+		"file": {
+			"uri": "file:///example.rego",
+			"lines": [
+				"package foo.bar",
+				"",
+				"import d",
+			],
+		},
+	}}
+
+	# 6 suggestions minus the current package and one test package
+	# also note how the sortText attribute hints to the client to sort not by
+	# the label but by the value calculated based on the number of path components
+	# (shortest first) and then alphabetically
+	items == {
+		_suggestion("data.bar", "000", [2, 7, 2, 8]),
+		_suggestion("data.baz", "001", [2, 7, 2, 8]),
+		_suggestion("data.foo.baz", "002", [2, 7, 2, 8]),
+		_suggestion("data.foo.baz.again", "003", [2, 7, 2, 8]),
+	}
+}
+
+test_matching_package_refs_sugggested_for_import if {
+	items := packagerefs.items with data.workspace.parsed as _workspace_parsed with input as {"regal": {
+		"context": {"location": {"row": 3, "col": 14}},
+		"file": {
+			"uri": "file:///example.rego",
+			"lines": [
+				"package foo.bar",
+				"",
+				"import data.f",
+			],
+		},
+	}}
+
+	items == {
+		_suggestion("data.foo.baz", "002", [2, 7, 2, 13]),
+		_suggestion("data.foo.baz.again", "003", [2, 7, 2, 13]),
+	}
+}
+
+_to_path(str) := [{"type": _type(i), "value": value} | some i, value in split(str, ".")]
+
+_type(0) := "var"
+_type(x) := "string" if x > 0
+
+_suggestion(label, sort_text, range) := {
+	"detail": "package",
+	"kind": 9,
+	"label": label,
+	"sortText": sort_text,
+	"textEdit": {
+		"newText": label,
+		"range": {
+			"start": {"line": range[0], "character": range[1]},
+			"end": {"line": range[2], "character": range[3]},
+		},
+	},
+}
+
+_workspace_parsed := {
+	"file:///example.rego": {"package": {"path": _to_path("data.foo.bar")}},
+	"file:///other.rego": {"package": {"path": _to_path("data.foo.baz")}},
+	"file:///again.rego": {"package": {"path": _to_path("data.foo.baz.again")}},
+	"file:///test.rego": {"package": {"path": _to_path("data.foo.baz.again_test")}},
+	"file:///more.rego": {"package": {"path": _to_path("data.bar")}},
+	"file:///last.rego": {"package": {"path": _to_path("data.baz")}},
+}

internal/lsp/cache/cache.go

@@ -49,11 +49,6 @@ type Cache struct {
 	// when a file is successfully parsed, the number of lines in the file is stored
 	// here. This is used to gracefully fail when exiting unparsable files.
 	successfulParseLineCounts *concurrent.Map[string, int]
-
-	// fileRefs is a map of file URI to refs that are defined in that file. These are
-	// intended to be used for completions in other files.
-	// fileRefs is expected to be updated when a file is successfully parsed.
-	fileRefs *concurrent.Map[string, map[string]types.Ref]
 }
 
 func NewCache() *Cache {
@@ -66,7 +61,6 @@ func NewCache() *Cache {
 		diagnosticsParseErrors:    concurrent.MapOf(make(map[string][]types.Diagnostic)),
 		builtinPositionsFile:      concurrent.MapOf(make(map[string]map[uint][]types.BuiltinPosition)),
 		keywordLocationsFile:      concurrent.MapOf(make(map[string]map[uint][]types.KeywordLocation)),
-		fileRefs:                  concurrent.MapOf(make(map[string]map[string]types.Ref)),
 		successfulParseLineCounts: concurrent.MapOf(make(map[string]int)),
 	}
 }
@@ -166,11 +160,6 @@ func (c *Cache) Rename(oldKey, newKey string) {
 		c.keywordLocationsFile.Delete(oldKey)
 	}
 
-	if refs, ok := c.fileRefs.Get(oldKey); ok {
-		c.fileRefs.Set(newKey, refs)
-		c.fileRefs.Delete(oldKey)
-	}
-
 	if lineCount, ok := c.successfulParseLineCounts.Get(oldKey); ok {
 		c.successfulParseLineCounts.Set(newKey, lineCount)
 		c.successfulParseLineCounts.Delete(oldKey)
@@ -282,20 +271,6 @@ func (c *Cache) GetKeywordLocations(fileURI string) (map[uint][]types.KeywordLoc
 	return c.keywordLocationsFile.Get(fileURI)
 }
 
-func (c *Cache) SetFileRefs(fileURI string, items map[string]types.Ref) {
-	c.fileRefs.Set(fileURI, items)
-}
-
-func (c *Cache) GetFileRefs(fileURI string) map[string]types.Ref {
-	refs, _ := c.fileRefs.Get(fileURI)
-
-	return refs
-}
-
-func (c *Cache) GetAllFileRefs() map[string]map[string]types.Ref {
-	return c.fileRefs.Clone()
-}
-
 func (c *Cache) GetSuccessfulParseLineCount(fileURI string) (int, bool) {
 	return c.successfulParseLineCounts.Get(fileURI)
 }
@@ -315,7 +290,6 @@ func (c *Cache) Delete(fileURI string) {
 	c.diagnosticsParseErrors.Delete(fileURI)
 	c.builtinPositionsFile.Delete(fileURI)
 	c.keywordLocationsFile.Delete(fileURI)
-	c.fileRefs.Delete(fileURI)
 	c.successfulParseLineCounts.Delete(fileURI)
 }
 

internal/lsp/completions/manager.go

@@ -8,83 +8,35 @@ import (
 
 	"github.com/open-policy-agent/regal/internal/lsp/cache"
 	"github.com/open-policy-agent/regal/internal/lsp/completions/providers"
-	"github.com/open-policy-agent/regal/internal/lsp/rego"
 	"github.com/open-policy-agent/regal/internal/lsp/rego/query"
 	"github.com/open-policy-agent/regal/internal/lsp/types"
+	"github.com/open-policy-agent/regal/internal/util"
 )
 
 type Manager struct {
-	c         *cache.Cache
-	opts      *ManagerOptions
-	providers []Provider
-}
-
-type ManagerOptions struct{}
-
-type Provider interface {
-	Run(context.Context, *cache.Cache, types.CompletionParams, *providers.Options) ([]types.CompletionItem, error)
-	Name() string
-}
-
-func NewManager(c *cache.Cache, opts *ManagerOptions) *Manager {
-	return &Manager{c: c, opts: opts}
+	c      *cache.Cache
+	policy *providers.Policy
 }
 
 func NewDefaultManager(ctx context.Context, c *cache.Cache, store storage.Store, qc *query.Cache) *Manager {
-	m := NewManager(c, &ManagerOptions{})
-
-	m.RegisterProvider(&providers.PackageRefs{})
-	m.RegisterProvider(providers.NewPolicy(ctx, store, qc))
-
-	return m
+	return &Manager{c: c, policy: providers.NewPolicy(ctx, store, qc)}
 }
 
 func (m *Manager) Run(
 	ctx context.Context,
 	params types.CompletionParams,
 	opts *providers.Options,
 ) ([]types.CompletionItem, error) {
-	if m.isInsideOfComment(params) {
-		// Exit early if caret position is inside a comment. We currently don't have any provider
-		// where doing completions inside of a comment makes much sense. Behavior is also editor-specific:
-		// - Zed: always on, with no way to disable
-		// - VSCode: disabled but can be enabled with "editor.quickSuggestions.comments" setting
-		return []types.CompletionItem{}, nil
+	completions, err := m.policy.Run(ctx, m.c, params, opts)
+	if err != nil {
+		return nil, fmt.Errorf("error running completion provider: %w", err)
 	}
 
-	var completionsList []types.CompletionItem
-
-	for _, provider := range m.providers {
-		providerCompletions, err := provider.Run(ctx, m.c, params, opts)
-		if err != nil {
-			return nil, fmt.Errorf("error running completion provider: %w", err)
-		}
-
-		for _, completion := range providerCompletions {
-			completion.Regal = nil
-			completionsList = append(completionsList, completion)
-		}
-	}
-
-	return completionsList, nil
+	return util.Map(completions, removeMetadata), nil
 }
 
-func (m *Manager) RegisterProvider(provider Provider) {
-	m.providers = append(m.providers, provider)
-}
-
-func (m *Manager) isInsideOfComment(params types.CompletionParams) bool {
-	if module, ok := m.c.GetModule(params.TextDocument.URI); ok {
-		for _, comment := range module.Comments {
-			cp := rego.PositionFromLocation(comment.Location)
-
-			if cp.Line == params.Position.Line {
-				if cp.Character <= params.Position.Character {
-					return true
-				}
-			}
-		}
-	}
+func removeMetadata(item types.CompletionItem) types.CompletionItem {
+	item.Regal = nil
 
-	return false
+	return item
 }