fix: fix lsp semantic token(sema hight) of import stmt and unresolved…

… token Signed-off-by: he1pa <[email protected]>
kcl-lang · Jan 17, 2024 · 22c96e7 · 22c96e7
1 parent 5fa77a8
commit 22c96e7
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 33 deletions.
diff --git a/kclvm/sema/src/advanced_resolver/node.rs b/kclvm/sema/src/advanced_resolver/node.rs
@@ -119,7 +119,12 @@ impl<'ctx> MutSelfTypedResultWalker<'ctx> for AdvancedResolver<'ctx> {
 
     fn walk_import_stmt(&mut self, import_stmt: &'ctx ast::ImportStmt) -> Self::Result {
         let ast_id = self.ctx.cur_node.clone();
-        let (start_pos, end_pos) = (self.ctx.start_pos.clone(), self.ctx.end_pos.clone());
+        let (start_pos, end_pos) = import_stmt
+            .asname
+            .clone()
+            .unwrap_or(import_stmt.path.clone())
+            .get_span_pos();
+
         let mut unresolved =
             UnresolvedSymbol::new(import_stmt.path.node.clone(), start_pos, end_pos, None);
         let package_symbol = self

diff --git a/kclvm/tools/src/LSP/src/semantic_token.rs b/kclvm/tools/src/LSP/src/semantic_token.rs
@@ -1,7 +1,10 @@
 use std::vec;
 
 use kclvm_error::Position;
-use kclvm_sema::core::{global_state::GlobalState, symbol::SymbolKind};
+use kclvm_sema::core::{
+    global_state::GlobalState,
+    symbol::{KCLSymbol, SymbolKind},
+};
 use lsp_types::{SemanticToken, SemanticTokenType, SemanticTokens, SemanticTokensResult};
 
 pub const LEGEND_TYPE: &[SemanticTokenType] = &[
@@ -25,31 +28,22 @@ pub(crate) fn semantic_tokens_full(file: &str, gs: &GlobalState) -> Option<Seman
     if let Some(file_sema) = sema_db.get_file_sema(&file.to_string()) {
         let symbols = file_sema.get_symbols();
         for symbol_ref in symbols {
-            if symbol_ref.get_kind() == SymbolKind::Expression {
-                continue;
-            }
             if let Some(symbol) = gs.get_symbols().get_symbol(*symbol_ref) {
                 let (start, end) = symbol.get_range();
-                let kind = match symbol_ref.get_kind() {
-                    SymbolKind::Schema => type_index(SemanticTokenType::STRUCT),
-                    SymbolKind::Attribute => type_index(SemanticTokenType::PROPERTY),
-                    SymbolKind::Package => type_index(SemanticTokenType::NAMESPACE),
-                    SymbolKind::TypeAlias => type_index(SemanticTokenType::TYPE),
-                    SymbolKind::Value | SymbolKind::Unresolved => {
-                        type_index(SemanticTokenType::VARIABLE)
+                match get_kind(symbol_ref.get_kind(), symbol, gs) {
+                    Some(kind) => {
+                        kcl_tokens.push(KCLSemanticToken {
+                            start: start.clone(),
+                            kind,
+                            length: if start.line == end.line {
+                                (end.column.unwrap_or(0) - start.column.unwrap_or(0)) as u32
+                            } else {
+                                symbol.get_name().len() as u32
+                            },
+                        });
                     }
-                    SymbolKind::Rule => type_index(SemanticTokenType::MACRO),
-                    SymbolKind::Expression => unreachable!(),
-                };
-                kcl_tokens.push(KCLSemanticToken {
-                    start: start.clone(),
-                    kind,
-                    length: if start.line == end.line {
-                        (end.column.unwrap_or(0) - start.column.unwrap_or(0)) as u32
-                    } else {
-                        symbol.get_name().len() as u32
-                    },
-                });
+                    None => continue,
+                }
             }
         }
     }
@@ -60,6 +54,25 @@ pub(crate) fn semantic_tokens_full(file: &str, gs: &GlobalState) -> Option<Seman
     }))
 }
 
+pub(crate) fn get_kind(ty: SymbolKind, symbol: &KCLSymbol, gs: &GlobalState) -> Option<u32> {
+    match ty {
+        SymbolKind::Schema => Some(type_index(SemanticTokenType::STRUCT)),
+        SymbolKind::Attribute => Some(type_index(SemanticTokenType::PROPERTY)),
+        SymbolKind::Package => Some(type_index(SemanticTokenType::NAMESPACE)),
+        SymbolKind::TypeAlias => Some(type_index(SemanticTokenType::TYPE)),
+        SymbolKind::Value => Some(type_index(SemanticTokenType::VARIABLE)),
+        SymbolKind::Rule => Some(type_index(SemanticTokenType::MACRO)),
+        SymbolKind::Unresolved => match &symbol.get_definition() {
+            Some(def_ref) => match gs.get_symbols().get_symbol(*def_ref) {
+                Some(symbol) => get_kind(def_ref.get_kind(), symbol, gs),
+                None => Some(type_index(SemanticTokenType::VARIABLE)),
+            },
+            None => Some(type_index(SemanticTokenType::VARIABLE)),
+        },
+        SymbolKind::Expression => None,
+    }
+}
+
 pub(crate) fn type_index(ty: SemanticTokenType) -> u32 {
     LEGEND_TYPE.iter().position(|it| *it == ty).unwrap() as u32
 }
@@ -126,15 +139,16 @@ mod tests {
     fn semantic_tokens_full_test() {
         let (file, _, _, _, gs) = compile_test_file("src/test_data/sema_token.k");
         let expected = [
-            (0, 5, 3, 4),
-            (1, 7, 7, 1),
-            (1, 4, 4, 2),
-            (2, 0, 2, 0),
-            (0, 4, 7, 0),
-            (0, 10, 7, 0),
-            (1, 4, 4, 0),
-            (2, 0, 1, 0),
-            (0, 3, 3, 0),
+            (0, 15, 1, 3), // m
+            (1, 5, 3, 4),  // num
+            (1, 7, 7, 1),  // Persons
+            (1, 4, 4, 2),  // name
+            (2, 0, 2, 0),  // p5
+            (0, 4, 7, 1),  // Persons
+            (0, 10, 7, 1), // Persons
+            (1, 4, 4, 2),  // name
+            (2, 0, 1, 0),  // n
+            (0, 3, 3, 4),  // num
         ];
         let res = semantic_tokens_full(&file, &gs);
         if let Some(tokens) = res {

diff --git a/kclvm/tools/src/LSP/src/test_data/sema_token.k b/kclvm/tools/src/LSP/src/test_data/sema_token.k
@@ -1,3 +1,4 @@
+import math as m
 type num = int | float
 schema Persons:
     name: str