From 096abdd19bb422cba607d43d05b195b37a7f95e9 Mon Sep 17 00:00:00 2001
From: lazytanuki <43273245+lazytanuki@users.noreply.github.com>
Date: Mon, 27 Jun 2022 13:19:56 +0200
Subject: [PATCH] feat: highlight / select symbol under cursor using LSP
 textDocument/documentHighlight (#2738)

* feat: highlight symbol under cursor using LSP textDocument/documentHighlight

* fix: please clippy

* fix: shorter description and code style
---
 helix-lsp/src/client.rs          | 20 +++++++++++++++
 helix-term/src/commands.rs       |  1 +
 helix-term/src/commands/lsp.rs   | 43 ++++++++++++++++++++++++++++++++
 helix-term/src/keymap/default.rs |  1 +
 4 files changed, 65 insertions(+)

diff --git a/helix-lsp/src/client.rs b/helix-lsp/src/client.rs
index cf6706b5..aa833541 100644
--- a/helix-lsp/src/client.rs
+++ b/helix-lsp/src/client.rs
@@ -759,6 +759,26 @@ impl Client {
         Ok(response.unwrap_or_default())
     }
 
+    pub fn text_document_document_highlight(
+        &self,
+        text_document: lsp::TextDocumentIdentifier,
+        position: lsp::Position,
+        work_done_token: Option<lsp::ProgressToken>,
+    ) -> impl Future<Output = Result<Value>> {
+        let params = lsp::DocumentHighlightParams {
+            text_document_position_params: lsp::TextDocumentPositionParams {
+                text_document,
+                position,
+            },
+            work_done_progress_params: lsp::WorkDoneProgressParams { work_done_token },
+            partial_result_params: lsp::PartialResultParams {
+                partial_result_token: None,
+            },
+        };
+
+        self.call::<lsp::request::DocumentHighlightRequest>(params)
+    }
+
     fn goto_request<
         T: lsp::request::Request<
             Params = lsp::GotoDefinitionParams,
diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index 43bc0c7c..0a28444b 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -263,6 +263,7 @@ impl MappableCommand {
         code_action, "Perform code action",
         buffer_picker, "Open buffer picker",
         symbol_picker, "Open symbol picker",
+        select_references_to_symbol_under_cursor, "Select symbol references",
         workspace_symbol_picker, "Open workspace symbol picker",
         last_picker, "Open last picker",
         prepend_to_line, "Insert at start of line",
diff --git a/helix-term/src/commands/lsp.rs b/helix-term/src/commands/lsp.rs
index 6f75fb49..c1cdb164 100644
--- a/helix-term/src/commands/lsp.rs
+++ b/helix-term/src/commands/lsp.rs
@@ -655,6 +655,7 @@ pub fn signature_help(cx: &mut Context) {
         },
     );
 }
+
 pub fn hover(cx: &mut Context) {
     let (view, doc) = current!(cx.editor);
     let language_server = language_server!(cx.editor, doc);
@@ -704,6 +705,7 @@ pub fn hover(cx: &mut Context) {
         },
     );
 }
+
 pub fn rename_symbol(cx: &mut Context) {
     ui::prompt(
         cx,
@@ -729,3 +731,44 @@ pub fn rename_symbol(cx: &mut Context) {
         },
     );
 }
+
+pub fn select_references_to_symbol_under_cursor(cx: &mut Context) {
+    let (view, doc) = current!(cx.editor);
+    let language_server = language_server!(cx.editor, doc);
+    let offset_encoding = language_server.offset_encoding();
+
+    let pos = doc.position(view.id, offset_encoding);
+
+    let future = language_server.text_document_document_highlight(doc.identifier(), pos, None);
+
+    cx.callback(
+        future,
+        move |editor, _compositor, response: Option<Vec<lsp::DocumentHighlight>>| {
+            let document_highlights = match response {
+                Some(highlights) if !highlights.is_empty() => highlights,
+                _ => return,
+            };
+            let (view, doc) = current!(editor);
+            let language_server = language_server!(editor, doc);
+            let offset_encoding = language_server.offset_encoding();
+            let text = doc.text();
+            let pos = doc.selection(view.id).primary().head;
+
+            // We must find the range that contains our primary cursor to prevent our primary cursor to move
+            let mut primary_index = 0;
+            let ranges = document_highlights
+                .iter()
+                .filter_map(|highlight| lsp_range_to_range(text, highlight.range, offset_encoding))
+                .enumerate()
+                .map(|(i, range)| {
+                    if range.contains(pos) {
+                        primary_index = i;
+                    }
+                    range
+                })
+                .collect();
+            let selection = Selection::new(ranges, primary_index);
+            doc.set_selection(view.id, selection);
+        },
+    );
+}
diff --git a/helix-term/src/keymap/default.rs b/helix-term/src/keymap/default.rs
index c3695117..e5bb0482 100644
--- a/helix-term/src/keymap/default.rs
+++ b/helix-term/src/keymap/default.rs
@@ -259,6 +259,7 @@ pub fn default() -> HashMap<Mode, Keymap> {
             "/" => global_search,
             "k" => hover,
             "r" => rename_symbol,
+            "h" => select_references_to_symbol_under_cursor,
             "?" => command_palette,
         },
         "z" => { "View"