From 33461babf2dd5992288eff391fdb4b9dbf1abeb1 Mon Sep 17 00:00:00 2001
From: LazyTanuki <43273245+lazytanuki@users.noreply.github.com>
Date: Fri, 10 Jun 2022 19:45:04 +0200
Subject: [PATCH] fix: shorter description and code style

---
 helix-term/src/commands.rs     |  2 +-
 helix-term/src/commands/lsp.rs | 56 +++++++++++++++-------------------
 2 files changed, 26 insertions(+), 32 deletions(-)

diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index 7e93839ef0321..69cf6e3a6aff0 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -263,7 +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 references to symbol under cursor",
+        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 6f0780107b687..35525cdb87df3 100644
--- a/helix-term/src/commands/lsp.rs
+++ b/helix-term/src/commands/lsp.rs
@@ -6,7 +6,7 @@ use helix_lsp::{
 
 use super::{align_view, push_jump, Align, Context, Editor};
 
-use helix_core::{Selection, SmallVec};
+use helix_core::Selection;
 use helix_view::editor::Action;
 
 use crate::{
@@ -734,37 +734,31 @@ pub fn select_references_to_symbol_under_cursor(cx: &mut Context) {
     cx.callback(
         future,
         move |editor, _compositor, response: Option<Vec<lsp::DocumentHighlight>>| {
-            if let Some(document_highlights) = response {
-                if !document_highlights.is_empty() {
-                    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;
-
-                    let mut ranges = SmallVec::with_capacity(document_highlights.len());
-
-                    // We must find the range that contains our primary cursor to prevent our primary cursor to move
-                    let mut primary_index = 0;
-                    let mut i = 0;
-
-                    for document_highlight in document_highlights {
-                        let lsp_range: lsp::Range = document_highlight.range;
-                        let range =
-                            helix_lsp::util::lsp_range_to_range(text, lsp_range, offset_encoding);
-                        if let Some(range) = range {
-                            if range.contains(pos) || range.head == pos {
-                                primary_index = i;
-                            }
-                            ranges.push(range);
-                            i += 1;
-                        }
-                    }
+            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;
 
-                    let selection = Selection::new(ranges, primary_index);
-                    doc.set_selection(view.id, selection);
-                }
-            }
+            // 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);
         },
     );
 }