Add support for documentHighlight

6 files changed, 78 insertions, 3 deletions

diff --git a/README.md b/README.md
index 05ee71f..0ea80d7 100644
--- a/README.md
+++ b/README.md
@@ -78,6 +78,8 @@ Options:
 
 - `enable_references`: Turns on finding references for a symbol. _(Enabled by default)_
 
+- `enable_document_highlights`: Turns on highlighting of symbol references in file. _(Enabled by default)_
+
 - `enable_document_links`: Follow links when opening documentation. This is usually done via `<ctrl+click>` and will open the documentation in a browser (or similar). _(Enabled by default)_
 
 - `enable_completion_matching`: Attempt to match types and pointers when passing arguments to procedures. _(Enabled by default)_
diff --git a/misc/ols.schema.json b/misc/ols.schema.json
index 898914e..c3ad997 100644
--- a/misc/ols.schema.json
+++ b/misc/ols.schema.json
@@ -72,6 +72,11 @@
 			"description": "Turns on finding references for a symbol.",
 			"default": true
 		},
+		"enable_document_highlights": {
+			"type": "boolean",
+			"description": "Turns on highlighting of symbol references in file.",
+			"default": true
+		},
 		"enable_completion_matching": {
 			"type": "boolean",
 			"description": "Attempt to match types and pointers when passing arguments to procedures.",
diff --git a/src/common/config.odin b/src/common/config.odin
index b44623a..4aaefac 100644
--- a/src/common/config.odin
+++ b/src/common/config.odin
@@ -28,6 +28,7 @@ Config :: struct {
 	enable_procedure_context:           bool,
 	enable_snippets:                    bool,
 	enable_references:                  bool,
+	enable_document_highlights:         bool,
 	enable_label_details:               bool,
 	enable_std_references:              bool,
 	enable_import_fixer:                bool,
diff --git a/src/server/references.odin b/src/server/references.odin
index 339be89..735d5d2 100644
--- a/src/server/references.odin
+++ b/src/server/references.odin
@@ -236,6 +236,7 @@ resolve_references :: proc(
 	document: ^Document,
 	ast_context: ^AstContext,
 	position_context: ^DocumentPositionContext,
+	current_file_only := false,
 ) -> (
 	[]common.Location,
 	bool,
@@ -269,7 +270,7 @@ resolve_references :: proc(
 
 	fullpaths := slice.unique(fullpaths[:])
 
-	if .Local not_in symbol.flags {
+	if .Local not_in symbol.flags && !current_file_only {
 		for fullpath in fullpaths {
 			dir := filepath.dir(fullpath)
 			base := filepath.base(dir)
@@ -384,7 +385,7 @@ resolve_references :: proc(
 	return locations[:], true
 }
 
-get_references :: proc(document: ^Document, position: common.Position) -> ([]common.Location, bool) {
+get_references :: proc(document: ^Document, position: common.Position, current_file_only := false) -> ([]common.Location, bool) {
 	ast_context := make_ast_context(
 		document.ast,
 		document.imports,
@@ -410,7 +411,7 @@ get_references :: proc(document: ^Document, position: common.Position) -> ([]com
 		get_locals(document.ast, position_context.function, &ast_context, &position_context)
 	}
 
-	locations, ok2 := resolve_references(document, &ast_context, &position_context)
+	locations, ok2 := resolve_references(document, &ast_context, &position_context, current_file_only)
 
 	temp_locations := make([dynamic]common.Location, 0, context.temp_allocator)
 
diff --git a/src/server/requests.odin b/src/server/requests.odin
index 7253063..a775850 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/documentHighlight"    = request_highlights,
 	"textDocument/codeAction"           = request_code_action,
 	"window/progress"                   = request_noop,
 	"workspace/symbol"                  = request_workspace_symbols,
@@ -367,6 +368,7 @@ read_ols_initialize_options :: proc(config: ^common.Config, ols_config: OlsConfi
 		ols_config.enable_procedure_context.(bool) or_else config.enable_procedure_context
 	config.enable_snippets = ols_config.enable_snippets.(bool) or_else config.enable_snippets
 	config.enable_references = ols_config.enable_references.(bool) or_else config.enable_references
+	config.enable_document_highlights = ols_config.enable_document_highlights.(bool) or_else config.enable_document_highlights
 	config.enable_completion_matching =
 		ols_config.enable_completion_matching.(bool) or_else config.enable_completion_matching
 	config.enable_document_links = ols_config.enable_document_links.(bool) or_else config.enable_document_links
@@ -619,6 +621,7 @@ request_initialize :: proc(
 	config.enable_procedure_context = false
 	config.enable_snippets = false
 	config.enable_references = true
+	config.enable_document_highlights = true
 	config.enable_completion_matching = true
 	config.enable_document_links = true
 	config.verbose = false
@@ -708,6 +711,7 @@ request_initialize :: proc(
 				renameProvider = RenameOptions{prepareProvider = true},
 				workspaceSymbolProvider = true,
 				referencesProvider = config.enable_references,
+				documentHighlightProvider = config.enable_document_highlights,
 				definitionProvider = true,
 				typeDefinitionProvider = true,
 				completionProvider = CompletionOptions {
@@ -1493,6 +1497,49 @@ request_references :: proc(
 	return .None
 }
 
+request_highlights :: proc(
+	params: json.Value,
+	id: RequestId,
+	config: ^common.Config,
+	writer: ^Writer,
+) -> common.Error {
+	params_object, ok := params.(json.Object)
+
+	if !ok {
+		return .ParseError
+	}
+
+	highlight_param: HighlightParams
+
+	if unmarshal(params, highlight_param, context.temp_allocator) != nil {
+		return .ParseError
+	}
+
+	document := document_get(highlight_param.textDocument.uri)
+
+	if document == nil {
+		return .InternalError
+	}
+
+	locations: []common.Location
+	locations, ok = get_references(document, highlight_param.position, true)
+
+	if !ok {
+		return .InternalError
+	}
+
+	highlights := make([dynamic]DocumentHighlight, 0, context.temp_allocator)
+	for location in locations {
+		append(&highlights, DocumentHighlight{kind = .Text, range = location.range})
+	}
+
+	response := make_response_message(params = highlights[:], id = id)
+
+	send_response(response, writer)
+
+	return .None
+}
+
 request_code_action :: proc(params: json.Value, id: RequestId, config: ^common.Config, writer: ^Writer) -> common.Error {
 	params_object, ok := params.(json.Object)
 
diff --git a/src/server/types.odin b/src/server/types.odin
index e0a2667..4b4c0bd 100644
--- a/src/server/types.odin
+++ b/src/server/types.odin
@@ -32,6 +32,7 @@ ResponseParams :: union {
 	WorkspaceEdit,
 	common.Range,
 	[]CodeAction,
+	[]DocumentHighlight,
 }
 
 RequestMessage :: struct {
@@ -143,6 +144,7 @@ ServerCapabilities :: struct {
 	inlayHintProvider:          bool,
 	renameProvider:             RenameOptions,
 	referencesProvider:         bool,
+	documentHighlightProvider:  bool,
 	workspaceSymbolProvider:    bool,
 	documentLinkProvider:       DocumentLinkOptions,
 	codeActionProvider:         CodeActionOptions,
@@ -418,6 +420,7 @@ OlsConfig :: struct {
 	enable_document_symbols:           Maybe(bool),
 	enable_fake_methods:               Maybe(bool),
 	enable_references:                 Maybe(bool),
+	enable_document_highlights:        Maybe(bool),
 	enable_document_links:             Maybe(bool),
 	enable_completion_matching:        Maybe(bool),
 	enable_inlay_hints_params:         Maybe(bool),
@@ -560,6 +563,11 @@ ReferenceParams :: struct {
 	position:     common.Position,
 }
 
+HighlightParams :: struct {
+	textDocument: TextDocumentIdentifier,
+	position:     common.Position,
+}
+
 OptionalVersionedTextDocumentIdentifier :: struct {
 	uri:     string,
 	version: Maybe(int),
@@ -587,3 +595,14 @@ WorkspaceSymbol :: struct {
 DidChangeConfigurationParams :: struct {
 	settings: OlsConfig,
 }
+
+DocumentHighlight :: struct {
+	range: common.Range,
+	kind:  DocumentHighlightKind,
+}
+
+DocumentHighlightKind :: enum {
+	Text  = 1,
+	Read  = 2,
+	Write = 3,
+}