Add code actions for importing packages from collections

4 files changed, 187 insertions, 4 deletions

diff --git a/src/server/action.odin b/src/server/action.odin
index ba1b367..4a337a3 100644
--- a/src/server/action.odin
+++ b/src/server/action.odin
@@ -1,11 +1,19 @@
 package server
 
-CodeActionKind :: struct {}
+import "core:fmt"
+import "core:log"
+import "core:odin/ast"
+import path "core:path/slashpath"
+import "core:strings"
+
+import "src:common"
+
+CodeActionKind :: string
 
 CodeActionClientCapabilities :: struct {
 	codeActionLiteralSupport: struct {
 		codeActionKind: struct {
-			valueSet: []CodeActionKind,
+			valueSet: [dynamic]CodeActionKind,
 		},
 	},
 }
@@ -14,3 +22,108 @@ CodeActionOptions :: struct {
 	codeActionKinds: []CodeActionKind,
 	resolveProvider: bool,
 }
+
+CodeActionParams :: struct {
+	textDocument: TextDocumentIdentifier,
+	range:        common.Range,
+}
+
+CodeAction :: struct {
+	title:       string,
+	kind:        CodeActionKind,
+	isPreferred: bool,
+	edit:        WorkspaceEdit,
+}
+
+get_code_actions :: proc(document: ^Document, range: common.Range, config: ^common.Config) -> ([]CodeAction, bool) {
+	ast_context := make_ast_context(
+		document.ast,
+		document.imports,
+		document.package_name,
+		document.uri.uri,
+		document.fullpath,
+		context.temp_allocator,
+	)
+
+	position_context, ok := get_document_position_context(document, range.start, .Hover)
+	if !ok {
+		log.warn("Failed to get position context")
+		return {}, false
+	}
+
+	ast_context.position_hint = position_context.hint
+
+	get_globals(document.ast, &ast_context)
+
+	ast_context.current_package = ast_context.document_package
+
+	if position_context.function != nil {
+		get_locals(document.ast, position_context.function, &ast_context, &position_context)
+	}
+
+	actions := make([dynamic]CodeAction, 0, context.allocator)
+
+	if position_context.selector_expr != nil {
+		if selector, ok := position_context.selector_expr.derived.(^ast.Selector_Expr); ok {
+			add_missing_imports(&ast_context, selector, strings.clone(document.uri.uri), config, &actions)
+		}
+	}
+
+	return actions[:], true
+}
+
+add_missing_imports :: proc(
+	ast_context: ^AstContext,
+	selector: ^ast.Selector_Expr,
+	uri: string,
+	config: ^common.Config,
+	actions: ^[dynamic]CodeAction,
+) {
+	if name, ok := selector.expr.derived.(^ast.Ident); ok {
+		for collection, pkgs in build_cache.pkg_aliases {
+			for pkg in pkgs {
+				fullpath := path.join({config.collections[collection], pkg})
+				found := false
+
+				for doc_pkg in ast_context.imports {
+					if fullpath == doc_pkg.name {
+						found = true
+					}
+				}
+
+				if found {
+					continue
+				}
+
+				if pkg == name.name {
+					pkg_decl := ast_context.file.pkg_decl
+					log.error(pkg_decl.end.line)
+					import_edit := TextEdit {
+						range = {
+							start = {line = pkg_decl.end.line + 1, character = 0},
+							end = {line = pkg_decl.end.line + 1, character = 0},
+						},
+						newText = fmt.tprintf("import \"%v:%v\"\n", collection, pkg),
+					}
+					textEdits := make([dynamic]TextEdit, context.temp_allocator)
+					append(&textEdits, import_edit)
+
+					workspaceEdit: WorkspaceEdit
+					workspaceEdit.changes = make(map[string][]TextEdit, 0, context.temp_allocator)
+					workspaceEdit.changes[uri] = textEdits[:]
+					append(
+						actions,
+						CodeAction {
+							kind = "refactor.rewrite",
+							isPreferred = true,
+							title = fmt.tprintf(`import package "%v:%v"`, collection, pkg),
+							edit = workspaceEdit,
+						},
+					)
+				}
+			}
+		}
+	}
+
+	return
+}
diff --git a/src/server/requests.odin b/src/server/requests.odin
index 4f090d8..0cbcd99 100644
--- a/src/server/requests.odin
+++ b/src/server/requests.odin
@@ -241,6 +241,7 @@ call_map: map[string]proc(_: json.Value, _: RequestId, _: ^common.Config, _: ^Wr
 	"textDocument/rename"               = request_rename,
 	"textDocument/prepareRename"        = request_prepare_rename,
 	"textDocument/references"           = request_references,
+	"textDocument/codeAction"           = request_code_action,
 	"window/progress"                   = request_noop,
 	"workspace/symbol"                  = request_workspace_symbols,
 	"workspace/didChangeConfiguration"  = notification_workspace_did_change_configuration,
@@ -586,7 +587,8 @@ request_initialize :: proc(
 
 	initialize_params: RequestInitializeParams
 
-	if unmarshal(params, initialize_params, context.temp_allocator) != nil {
+	if err := unmarshal(params, initialize_params, context.temp_allocator); err != nil {
+		log.error("Here?", err, params)
 		return .ParseError
 	}
 
@@ -726,6 +728,7 @@ request_initialize :: proc(
 				hoverProvider = config.enable_hover,
 				documentFormattingProvider = config.enable_format,
 				documentLinkProvider = {resolveProvider = false},
+				codeActionProvider = {resolveProvider = false, codeActionKinds = {"refactor.rewrite"}},
 			},
 		},
 		id = id,
@@ -1483,6 +1486,37 @@ request_references :: proc(
 	return .None
 }
 
+request_code_action :: proc(params: json.Value, id: RequestId, config: ^common.Config, writer: ^Writer) -> common.Error {
+	params_object, ok := params.(json.Object)
+
+	if !ok {
+		return .ParseError
+	}
+
+	code_action_params: CodeActionParams
+
+	if unmarshal(params, code_action_params, context.temp_allocator) != nil {
+		return .ParseError
+	}
+
+	document := document_get(code_action_params.textDocument.uri)
+
+	if document == nil {
+		return .InternalError
+	}
+
+	code_actions: []CodeAction
+	code_actions, ok = get_code_actions(document, code_action_params.range, config)
+	if !ok {
+		return .InternalError
+	}
+	response := make_response_message(params = code_actions, id = id)
+
+	send_response(response, writer)
+
+	return .None
+}
+
 notification_did_change_watched_files :: proc(
 	params: json.Value,
 	id: RequestId,
diff --git a/src/server/types.odin b/src/server/types.odin
index bdbeaa6..bcd0e07 100644
--- a/src/server/types.odin
+++ b/src/server/types.odin
@@ -31,6 +31,7 @@ ResponseParams :: union {
 	[]WorkspaceSymbol,
 	WorkspaceEdit,
 	common.Range,
+	[]CodeAction,
 }
 
 RequestMessage :: struct {
@@ -144,6 +145,7 @@ ServerCapabilities :: struct {
 	referencesProvider:         bool,
 	workspaceSymbolProvider:    bool,
 	documentLinkProvider:       DocumentLinkOptions,
+	codeActionProvider:         CodeActionOptions,
 }
 
 DidChangeWatchedFilesRegistrationOptions :: struct {
@@ -187,6 +189,7 @@ TextDocumentClientCapabilities :: struct {
 	hover:          HoverClientCapabilities,
 	signatureHelp:  SignatureHelpClientCapabilities,
 	documentSymbol: DocumentSymbolClientCapabilities,
+	codeAction:     CodeActionClientCapabilities,
 }
 
 StaleRequestSupport :: struct {
@@ -225,7 +228,6 @@ DidChangeWatchedFilesClientCapabilities :: struct {
 	dynamicRegistration: bool,
 }
 
-
 RangeOptional :: union {
 	common.Range,
 }
@@ -576,3 +578,5 @@ WorkspaceSymbol :: struct {
 DidChangeConfigurationParams :: struct {
 	settings: OlsConfig,
 }
+
+
diff --git a/src/testing/testing.odin b/src/testing/testing.odin
index 1b40c6b..57354e1 100644
--- a/src/testing/testing.odin
+++ b/src/testing/testing.odin
@@ -458,6 +458,38 @@ expect_prepare_rename_range :: proc(t: ^testing.T, src: ^Source, expect_range: c
 	}
 }
 
+
+expect_action :: proc(t: ^testing.T, src: ^Source, expect_action_names: []string) {
+	setup(src)
+	defer teardown(src)
+
+	input_range := common.Range{start=src.position, end=src.position}
+	actions, ok := server.get_code_actions(src.document, input_range, &src.config)
+	if !ok {
+		log.error("Failed to find actions")
+	}
+
+	if len(expect_action_names) == 0 && len(actions) > 0 {
+		log.errorf("Expected empty actions, but received %v", actions)
+	}
+
+	flags := make([]int, len(expect_action_names), context.temp_allocator)
+
+	for name, i in expect_action_names {
+		for action, j in actions {
+			if action.title == name {
+				flags[i] += 1
+			}
+		}
+	}
+
+	for flag, i in flags {
+		if flag != 1 {
+			log.errorf("Expected action %v, but received %v", expect_action_names[i], actions)
+		}
+	}
+}
+
 expect_semantic_tokens :: proc(t: ^testing.T, src: ^Source, expected: []server.SemanticToken) {
 	setup(src)
 	defer teardown(src)