Refactor full highlight request

Author: pfoerster

src/features/semantic_tokens.rs

@@ -6,54 +6,40 @@ use lsp_types::{
     Position, Range, SemanticToken, SemanticTokenModifier, SemanticTokenType, SemanticTokens,
     SemanticTokensLegend, Url,
 };
-use rowan::TextRange;
+use rowan::{TextLen, TextRange};
 
 use crate::{
-    db::Workspace,
+    db::{Document, Workspace},
     util::{line_index::LineIndex, line_index_ext::LineIndexExt},
     Db,
 };
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
 #[repr(u32)]
-pub enum TokenKind {
+enum TokenKind {
     Label = 0,
     MathDelimiter = 1,
 }
 
 bitflags! {
     #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
-    pub struct TokenModifiers: u32 {
+    struct TokenModifiers: u32 {
         const NONE = 0;
         const UNDEFINED = 1;
         const UNUSED = 2;
         const DEPRECATED = 4;
     }
 }
 
-pub fn legend() -> SemanticTokensLegend {
-    SemanticTokensLegend {
-        token_types: vec![
-            SemanticTokenType::new("label"),
-            SemanticTokenType::new("mathDelimiter"),
-        ],
-        token_modifiers: vec![
-            SemanticTokenModifier::new("undefined"),
-            SemanticTokenModifier::new("unused"),
-            SemanticTokenModifier::new("deprecated"),
-        ],
-    }
-}
-
 #[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
-pub struct Token {
-    pub range: TextRange,
-    pub kind: TokenKind,
-    pub modifiers: TokenModifiers,
+struct Token {
+    range: TextRange,
+    kind: TokenKind,
+    modifiers: TokenModifiers,
 }
 
 #[derive(Debug, Default)]
-pub struct TokenBuilder {
+struct TokenBuilder {
     tokens: Vec<Token>,
 }
 
@@ -106,12 +92,45 @@ impl TokenBuilder {
     }
 }
 
-pub fn find_all(db: &dyn Db, uri: &Url, viewport: Range) -> Option<SemanticTokens> {
+#[derive(Clone, Copy)]
+struct Context<'db> {
+    db: &'db dyn Db,
+    document: Document,
+    viewport: TextRange,
+}
+
+pub fn legend() -> SemanticTokensLegend {
+    SemanticTokensLegend {
+        token_types: vec![
+            SemanticTokenType::new("label"),
+            SemanticTokenType::new("mathDelimiter"),
+        ],
+        token_modifiers: vec![
+            SemanticTokenModifier::new("undefined"),
+            SemanticTokenModifier::new("unused"),
+            SemanticTokenModifier::new("deprecated"),
+        ],
+    }
+}
+
+pub fn find_all(db: &dyn Db, uri: &Url, viewport: Option<Range>) -> Option<SemanticTokens> {
     let workspace = Workspace::get(db);
     let document = workspace.lookup_uri(db, uri)?;
-    let viewport = document.line_index(db).offset_lsp_range(viewport);
+    let viewport = viewport.map_or_else(
+        || TextRange::new(0.into(), document.text(db).text_len()),
+        |range| document.line_index(db).offset_lsp_range(range),
+    );
+
+    let context = Context {
+        db,
+        document,
+        viewport,
+    };
+
     let mut builder = TokenBuilder::default();
-    label::find(db, document, viewport, &mut builder);
-    math_delimiter::find(db, document, viewport, &mut builder);
+
+    label::find(context, &mut builder);
+    math_delimiter::find(context, &mut builder);
+
     Some(builder.finish(document.line_index(db)))
 }

src/features/semantic_tokens/label.rs

@@ -1,34 +1,28 @@
-use rowan::TextRange;
-
 use crate::{
     db::{analysis::label, Document, Workspace},
     Db,
 };
 
-use super::{Token, TokenBuilder, TokenKind, TokenModifiers};
+use super::{Context, Token, TokenBuilder, TokenKind, TokenModifiers};
 
-pub fn find(
-    db: &dyn Db,
-    document: Document,
-    viewport: TextRange,
-    builder: &mut TokenBuilder,
-) -> Option<()> {
-    let labels = document.parse(db).as_tex()?.analyze(db).labels(db);
+pub(super) fn find(context: Context, builder: &mut TokenBuilder) -> Option<()> {
+    let db = context.db;
+    let labels = context.document.parse(db).as_tex()?.analyze(db).labels(db);
     for label in labels
         .iter()
-        .filter(|label| viewport.intersect(label.range(db)).is_some())
+        .filter(|label| context.viewport.intersect(label.range(db)).is_some())
     {
         let name = label.name(db).text(db);
         let modifiers = match label.origin(db) {
             label::Origin::Definition(_) => {
-                if !is_label_referenced(db, document, name) {
+                if !is_label_referenced(db, context.document, name) {
                     TokenModifiers::UNUSED
                 } else {
                     TokenModifiers::NONE
                 }
             }
             label::Origin::Reference(_) | label::Origin::ReferenceRange(_) => {
-                if !is_label_defined(db, document, name) {
+                if !is_label_defined(db, context.document, name) {
                     TokenModifiers::UNDEFINED
                 } else {
                     TokenModifiers::NONE

src/features/semantic_tokens/math_delimiter.rs

@@ -1,19 +1,13 @@
-use rowan::TextRange;
+use crate::syntax::latex;
 
-use crate::{db::Document, syntax::latex, Db};
+use super::{Context, Token, TokenBuilder, TokenKind, TokenModifiers};
 
-use super::{Token, TokenBuilder, TokenKind, TokenModifiers};
-
-pub fn find(
-    db: &dyn Db,
-    document: Document,
-    viewport: TextRange,
-    builder: &mut TokenBuilder,
-) -> Option<()> {
-    let root = document.parse(db).as_tex()?.root(db);
+pub(super) fn find(context: Context, builder: &mut TokenBuilder) -> Option<()> {
+    let db = context.db;
+    let root = context.document.parse(db).as_tex()?.root(db);
 
     for token in root
-        .covering_element(viewport)
+        .covering_element(context.viewport)
         .as_node()?
         .descendants_with_tokens()
         .filter_map(|elem| elem.into_token())

src/server.rs

@@ -13,7 +13,7 @@ use log::{error, info};
 use lsp_server::{Connection, ErrorCode, Message, RequestId};
 use lsp_types::{notification::*, request::*, *};
 use once_cell::sync::Lazy;
-use rowan::{ast::AstNode, TextLen, TextRange, TextSize};
+use rowan::{ast::AstNode, TextSize};
 use rustc_hash::FxHashSet;
 use serde::{Deserialize, Serialize};
 use serde_repr::{Deserialize_repr, Serialize_repr};
@@ -816,9 +816,7 @@ impl Server {
 
     fn semantic_tokens_full(&mut self, id: RequestId, params: SemanticTokensParams) -> Result<()> {
         self.run_with_db(id, move |db| {
-            let Some(document) = Workspace::get(db).lookup_uri(db, &params.text_document.uri) else { return None };
-            let range = document.line_index(db).line_col_lsp_range(TextRange::new(0.into(), document.text(db).text_len()));
-            semantic_tokens::find_all(db, &params.text_document.uri, range)
+            semantic_tokens::find_all(db, &params.text_document.uri, None)
         });
 
         Ok(())
@@ -830,7 +828,7 @@ impl Server {
         params: SemanticTokensRangeParams,
     ) -> Result<()> {
         self.run_with_db(id, move |db| {
-            semantic_tokens::find_all(db, &params.text_document.uri, params.range)
+            semantic_tokens::find_all(db, &params.text_document.uri, Some(params.range))
         });
 
         Ok(())