Use different token types depending on labeled obj

latex-lsp · Mar 21, 2023 · 0d5b0aa · 0d5b0aa
1 parent e1c4a20
commit 0d5b0aa
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 24 deletions.
diff --git a/src/features/semantic_tokens.rs b/src/features/semantic_tokens.rs
@@ -18,9 +18,14 @@ use crate::{
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
 #[repr(u32)]
 enum TokenKind {
-    Label = 0,
-    Citation = 2,
-    MathDelimiter = 3,
+    GenericLabel = 0,
+    SectionLabel,
+    FloatLabel,
+    TheoremLabel,
+    EquationLabel,
+    EnumItemLabel,
+    Citation,
+    MathDelimiter,
 }
 
 bitflags! {
@@ -104,7 +109,12 @@ struct Context<'db> {
 pub fn legend() -> SemanticTokensLegend {
     SemanticTokensLegend {
         token_types: vec![
-            SemanticTokenType::new("label"),
+            SemanticTokenType::new("genericLabel"),
+            SemanticTokenType::new("sectionLabel"),
+            SemanticTokenType::new("floatLabel"),
+            SemanticTokenType::new("theoremLabel"),
+            SemanticTokenType::new("equationLabel"),
+            SemanticTokenType::new("enumItemLabel"),
             SemanticTokenType::new("citation"),
             SemanticTokenType::new("mathDelimiter"),
         ],

diff --git a/src/features/semantic_tokens/label.rs b/src/features/semantic_tokens/label.rs
@@ -1,6 +1,6 @@
 use crate::{
     db::{analysis::label, Document, Workspace},
-    Db,
+    util, Db,
 };
 
 use super::{Context, Token, TokenBuilder, TokenKind, TokenModifiers};
@@ -11,36 +11,64 @@ pub(super) fn find(context: Context, builder: &mut TokenBuilder) -> Option<()> {
     for label in labels
         .iter()
         .filter(|label| context.viewport.intersect(label.range(db)).is_some())
+        .copied()
     {
-        let name = label.name(db).text(db);
-        let modifiers = match label.origin(db) {
-            label::Origin::Definition(_) => {
-                if !is_label_referenced(db, context.document, name) {
-                    TokenModifiers::UNUSED
-                } else {
-                    TokenModifiers::NONE
-                }
-            }
-            label::Origin::Reference(_) | label::Origin::ReferenceRange(_) => {
-                if !is_label_defined(db, context.document, name) {
-                    TokenModifiers::UNDEFINED
-                } else {
-                    TokenModifiers::NONE
-                }
-            }
-        };
-
+        let kind = token_type(context, label);
+        let modifiers = token_modifiers(context, label);
         let range = label.range(db);
         builder.push(Token {
             range,
-            kind: TokenKind::Label,
+            kind,
             modifiers,
         });
     }
 
     Some(())
 }
 
+fn token_type(context: Context, label: label::Name) -> TokenKind {
+    let db = context.db;
+    let definition = match label.origin(db) {
+        label::Origin::Definition(_) => Some((context.document, label)),
+        label::Origin::Reference(_) | label::Origin::ReferenceRange(_) => {
+            util::label::find_label_definition(db, context.document, label.name(db))
+        }
+    };
+
+    match definition
+        .and_then(|(doc, label)| util::label::render(db, doc, label))
+        .map(|label| label.object)
+    {
+        Some(util::label::LabeledObject::Section { .. }) => TokenKind::SectionLabel,
+        Some(util::label::LabeledObject::Float { .. }) => TokenKind::FloatLabel,
+        Some(util::label::LabeledObject::EnumItem { .. }) => TokenKind::EnumItemLabel,
+        Some(util::label::LabeledObject::Equation { .. }) => TokenKind::EquationLabel,
+        Some(util::label::LabeledObject::Theorem { .. }) => TokenKind::TheoremLabel,
+        None => TokenKind::GenericLabel,
+    }
+}
+
+fn token_modifiers(context: Context, label: label::Name) -> TokenModifiers {
+    let db = context.db;
+    let name = label.name(db).text(db);
+    match label.origin(db) {
+        label::Origin::Definition(_) => {
+            if !is_label_referenced(db, context.document, name) {
+                TokenModifiers::UNUSED
+            } else {
+                TokenModifiers::NONE
+            }
+        }
+        label::Origin::Reference(_) | label::Origin::ReferenceRange(_) => {
+            if !is_label_defined(db, context.document, name) {
+                TokenModifiers::UNDEFINED
+            } else {
+                TokenModifiers::NONE
+            }
+        }
+    }
+}
+
 fn is_label_defined(db: &dyn Db, child: Document, name: &str) -> bool {
     Workspace::get(db)
         .related(db, child)