to_proto::semantic_tokens

2 files changed, 33 insertions, 34 deletions

diff --git a/crates/rust-analyzer/src/main_loop/handlers.rs b/crates/rust-analyzer/src/main_loop/handlers.rs
index daa5b4411..4f619654f 100644
--- a/crates/rust-analyzer/src/main_loop/handlers.rs
+++ b/crates/rust-analyzer/src/main_loop/handlers.rs
@@ -35,7 +35,6 @@ use crate::{
     diagnostics::DiagnosticTask,
     from_json, from_proto,
     req::{self, InlayHint, InlayHintsParams},
-    semantic_tokens::SemanticTokensBuilder,
     to_proto,
     world::WorldSnapshot,
     LspError, Result,
@@ -1147,23 +1146,9 @@ pub fn handle_semantic_tokens(
     let text = world.analysis().file_text(file_id)?;
     let line_index = world.analysis().file_line_index(file_id)?;
 
-    let mut builder = SemanticTokensBuilder::default();
-
-    for highlight_range in world.analysis().highlight(file_id)?.into_iter() {
-        let (token_index, modifier_bitset) =
-            to_proto::token_type_index_modifiers_bitself(highlight_range.highlight);
-        for mut range in line_index.lines(highlight_range.range) {
-            if text[range].ends_with('\n') {
-                range = TextRange::new(range.start(), range.end() - TextSize::of('\n'));
-            }
-            let range = to_proto::range(&line_index, range);
-            builder.push(range, token_index, modifier_bitset);
-        }
-    }
-
-    let tokens = builder.build();
-
-    Ok(Some(tokens.into()))
+    let highlights = world.analysis().highlight(file_id)?;
+    let semantic_tokens = to_proto::semantic_tokens(&text, &line_index, highlights);
+    Ok(Some(semantic_tokens.into()))
 }
 
 pub fn handle_semantic_tokens_range(
@@ -1173,18 +1158,10 @@ pub fn handle_semantic_tokens_range(
     let _p = profile("handle_semantic_tokens_range");
 
     let frange = from_proto::file_range(&world, params.text_document, params.range)?;
+    let text = world.analysis().file_text(frange.file_id)?;
     let line_index = world.analysis().file_line_index(frange.file_id)?;
 
-    let mut builder = SemanticTokensBuilder::default();
-
-    for highlight_range in world.analysis().highlight_range(frange)?.into_iter() {
-        let (token_type, token_modifiers) =
-            to_proto::token_type_index_modifiers_bitself(highlight_range.highlight);
-        let range = to_proto::range(&line_index, highlight_range.range);
-        builder.push(range, token_type, token_modifiers);
-    }
-
-    let tokens = builder.build();
-
-    Ok(Some(tokens.into()))
+    let highlights = world.analysis().highlight_range(frange)?;
+    let semantic_tokens = to_proto::semantic_tokens(&text, &line_index, highlights);
+    Ok(Some(semantic_tokens.into()))
 }
diff --git a/crates/rust-analyzer/src/to_proto.rs b/crates/rust-analyzer/src/to_proto.rs
index 39f2e6d4d..29b47608e 100644
--- a/crates/rust-analyzer/src/to_proto.rs
+++ b/crates/rust-analyzer/src/to_proto.rs
@@ -3,8 +3,8 @@ use ra_db::{FileId, FileRange};
 use ra_ide::{
     translate_offset_with_edit, Assist, CompletionItem, CompletionItemKind, Documentation,
     FileSystemEdit, Fold, FoldKind, FunctionSignature, Highlight, HighlightModifier, HighlightTag,
-    InlayHint, InlayKind, InsertTextFormat, LineIndex, NavigationTarget, ReferenceAccess, Severity,
-    SourceChange, SourceFileEdit,
+    HighlightedRange, InlayHint, InlayKind, InsertTextFormat, LineIndex, NavigationTarget,
+    ReferenceAccess, Severity, SourceChange, SourceFileEdit,
 };
 use ra_syntax::{SyntaxKind, TextRange, TextSize};
 use ra_text_edit::{Indel, TextEdit};
@@ -227,8 +227,30 @@ pub(crate) fn inlay_int(line_index: &LineIndex, inlay_hint: InlayHint) -> req::I
     }
 }
 
-// TODO: this is wrong
-pub(crate) fn token_type_index_modifiers_bitself(highlight: Highlight) -> (u32, u32) {
+pub(crate) fn semantic_tokens(
+    text: &str,
+    line_index: &LineIndex,
+    highlights: Vec<HighlightedRange>,
+) -> lsp_types::SemanticTokens {
+    let mut builder = semantic_tokens::SemanticTokensBuilder::default();
+
+    for highlight_range in highlights {
+        let (token_index, modifier_bitset) =
+            token_type_index_modifiers_bitself(highlight_range.highlight);
+        for mut text_range in line_index.lines(highlight_range.range) {
+            if text[text_range].ends_with('\n') {
+                text_range =
+                    TextRange::new(text_range.start(), text_range.end() - TextSize::of('\n'));
+            }
+            let range = range(&line_index, text_range);
+            builder.push(range, token_index, modifier_bitset);
+        }
+    }
+
+    builder.build()
+}
+
+fn token_type_index_modifiers_bitself(highlight: Highlight) -> (u32, u32) {
     let mut mods = semantic_tokens::ModifierSet::default();
     let type_ = match highlight.tag {
         HighlightTag::Struct => lsp_types::SemanticTokenType::STRUCT,