diff --git a/kclvm/sema/src/core/global_state.rs b/kclvm/sema/src/core/global_state.rs index 26ce04fe4..f640841d8 100644 --- a/kclvm/sema/src/core/global_state.rs +++ b/kclvm/sema/src/core/global_state.rs @@ -45,6 +45,10 @@ impl GlobalState { pub fn get_packages_mut(&mut self) -> &mut PackageDB { &mut self.packages } + + pub fn get_sema_db(&self) -> &SemanticDB { + &self.sema_db + } } impl GlobalState { diff --git a/kclvm/sema/src/core/semantic_information.rs b/kclvm/sema/src/core/semantic_information.rs index efaac7dfe..517e53c81 100644 --- a/kclvm/sema/src/core/semantic_information.rs +++ b/kclvm/sema/src/core/semantic_information.rs @@ -11,6 +11,12 @@ pub struct SemanticDB { pub(crate) file_sema_map: IndexMap, } +impl SemanticDB { + pub fn get_file_sema(&self, file: &String) -> Option<&FileSemanticInfo> { + self.file_sema_map.get(file) + } +} + #[allow(unused)] #[derive(Debug, Clone)] pub struct FileSemanticInfo { @@ -47,6 +53,10 @@ impl FileSemanticInfo { } } } + + pub fn get_symbols(&self) -> &Vec { + &self.symbols + } } #[derive(Debug, Eq, PartialEq, Clone)] diff --git a/kclvm/tools/src/LSP/src/capabilities.rs b/kclvm/tools/src/LSP/src/capabilities.rs index fb9190187..6f0f0086c 100644 --- a/kclvm/tools/src/LSP/src/capabilities.rs +++ b/kclvm/tools/src/LSP/src/capabilities.rs @@ -1,14 +1,30 @@ use lsp_types::{ ClientCapabilities, CodeActionKind, CodeActionOptions, CodeActionProviderCapability, - CompletionOptions, HoverProviderCapability, OneOf, ServerCapabilities, - TextDocumentSyncCapability, TextDocumentSyncKind, WorkDoneProgressOptions, + CompletionOptions, HoverProviderCapability, OneOf, SemanticTokensFullOptions, + SemanticTokensLegend, SemanticTokensOptions, ServerCapabilities, TextDocumentSyncCapability, + TextDocumentSyncKind, WorkDoneProgressOptions, }; +use crate::semantic_token::LEGEND_TYPE; + /// Returns the capabilities of this LSP server implementation given the capabilities of the client. #[allow(dead_code)] pub fn server_capabilities(client_caps: &ClientCapabilities) -> ServerCapabilities { ServerCapabilities { text_document_sync: Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::FULL)), + semantic_tokens_provider: Some( + lsp_types::SemanticTokensServerCapabilities::SemanticTokensOptions( + SemanticTokensOptions { + work_done_progress_options: WorkDoneProgressOptions::default(), + legend: SemanticTokensLegend { + token_types: LEGEND_TYPE.into(), + token_modifiers: vec![], + }, + range: Some(false), + full: Some(SemanticTokensFullOptions::Bool(true)), + }, + ), + ), document_symbol_provider: Some(OneOf::Left(true)), completion_provider: Some(CompletionOptions { resolve_provider: None, diff --git a/kclvm/tools/src/LSP/src/lib.rs b/kclvm/tools/src/LSP/src/lib.rs index 5c112b5af..8aaf01612 100644 --- a/kclvm/tools/src/LSP/src/lib.rs +++ b/kclvm/tools/src/LSP/src/lib.rs @@ -15,6 +15,7 @@ mod notification; mod quick_fix; pub mod rename; mod request; +mod semantic_token; mod state; #[cfg(test)] mod tests; diff --git a/kclvm/tools/src/LSP/src/main.rs b/kclvm/tools/src/LSP/src/main.rs index 501bee468..1c3e58b85 100644 --- a/kclvm/tools/src/LSP/src/main.rs +++ b/kclvm/tools/src/LSP/src/main.rs @@ -17,6 +17,7 @@ mod main_loop; mod notification; mod quick_fix; mod request; +mod semantic_token; mod state; mod to_lsp; mod util; diff --git a/kclvm/tools/src/LSP/src/request.rs b/kclvm/tools/src/LSP/src/request.rs index 67a5bf63c..5c4d76cc2 100644 --- a/kclvm/tools/src/LSP/src/request.rs +++ b/kclvm/tools/src/LSP/src/request.rs @@ -3,7 +3,7 @@ use crossbeam_channel::Sender; use kclvm_config::modfile::KCL_FILE_SUFFIX; use kclvm_sema::info::is_valid_kcl_name; -use lsp_types::{Location, TextEdit}; +use lsp_types::{Location, SemanticTokensResult, TextEdit}; use ra_ap_vfs::VfsPath; use std::collections::HashMap; use std::time::Instant; @@ -18,6 +18,7 @@ use crate::{ from_lsp::{self, file_path_from_url, kcl_pos}, goto_def::goto_definition_with_gs, hover, quick_fix, + semantic_token::semantic_tokens_full, state::{log_message, LanguageServerSnapshot, LanguageServerState, Task}, util::{parse_param_and_compile, Param}, }; @@ -57,6 +58,7 @@ impl LanguageServerState { .on::(handle_formatting)? .on::(handle_range_formatting)? .on::(handle_rename)? + .on::(handle_semantic_tokens_full)? .finish(); Ok(()) @@ -91,6 +93,21 @@ impl LanguageServerSnapshot { } } +pub(crate) fn handle_semantic_tokens_full( + snapshot: LanguageServerSnapshot, + params: lsp_types::SemanticTokensParams, + sender: Sender, +) -> anyhow::Result> { + let file = file_path_from_url(¶ms.text_document.uri)?; + let path = from_lsp::abs_path(¶ms.text_document.uri)?; + if !snapshot.verify_request_path(&path.clone().into(), &sender) { + return Ok(None); + } + let db = snapshot.get_db(&path.clone().into())?; + let res = semantic_tokens_full(&file, &db.gs); + Ok(res) +} + pub(crate) fn handle_formatting( _snapshot: LanguageServerSnapshot, params: lsp_types::DocumentFormattingParams, diff --git a/kclvm/tools/src/LSP/src/semantic_token.rs b/kclvm/tools/src/LSP/src/semantic_token.rs new file mode 100644 index 000000000..02906e6be --- /dev/null +++ b/kclvm/tools/src/LSP/src/semantic_token.rs @@ -0,0 +1,159 @@ +use std::vec; + +use kclvm_error::Position; +use kclvm_sema::core::{global_state::GlobalState, symbol::SymbolKind}; +use lsp_types::{SemanticToken, SemanticTokenType, SemanticTokens, SemanticTokensResult}; + +pub const LEGEND_TYPE: &[SemanticTokenType] = &[ + SemanticTokenType::VARIABLE, + SemanticTokenType::STRUCT, + SemanticTokenType::PROPERTY, + SemanticTokenType::NAMESPACE, + SemanticTokenType::TYPE, + SemanticTokenType::MACRO, +]; + +pub(crate) struct KCLSemanticToken { + pub start: Position, + pub kind: u32, + pub length: u32, +} + +pub(crate) fn semantic_tokens_full(file: &str, gs: &GlobalState) -> Option { + let mut kcl_tokens: Vec = vec![]; + let sema_db = gs.get_sema_db(); + 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 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) + } + SymbolKind::Rule => type_index(SemanticTokenType::MACRO), + }; + 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 + }, + }); + } + } + } + + Some(SemanticTokensResult::Tokens(SemanticTokens { + result_id: None, + data: kcl_semantic_tokens_to_semantic_tokens(&mut kcl_tokens), + })) +} + +pub(crate) fn type_index(ty: SemanticTokenType) -> u32 { + LEGEND_TYPE.iter().position(|it| *it == ty).unwrap() as u32 +} + +pub(crate) fn kcl_semantic_tokens_to_semantic_tokens( + tokens: &mut Vec, +) -> Vec { + tokens.sort_by(|a, b| { + if a.start.line == b.start.line { + a.start + .column + .unwrap_or(0) + .cmp(&b.start.column.unwrap_or(0)) + } else { + a.start.line.cmp(&b.start.line) + } + }); + let mut pre_line = 0; + let mut pre_start = 0; + + let semantic_tokens: Vec = tokens + .iter() + .map(|obj| { + // ref: https://github.com/microsoft/vscode-extension-samples/blob/5ae1f7787122812dcc84e37427ca90af5ee09f14/semantic-tokens-sample/vscode.proposed.d.ts#L71 + // A file can contain many tokens, perhaps even hundreds of thousands of tokens. Therefore, to improve + // the memory consumption around describing semantic tokens, we have decided to avoid allocating an object + // for each token and we represent tokens from a file as an array of integers. Furthermore, the position + // of each token is expressed relative to the token before it because most tokens remain stable relative to + // each other when edits are made in a file. + let line = obj.start.line - 1; + let start = obj.start.column.unwrap_or(0); + + let delta_line: u32 = (line - pre_line) as u32; + let delta_start: u32 = (if delta_line == 0 { + start - pre_start + } else { + start + }) as u32; + let length = obj.length; + let ret = SemanticToken { + delta_line, + delta_start, + length, + token_type: obj.kind, + token_modifiers_bitset: 0, + }; + pre_line = line; + pre_start = start; + ret + }) + .collect(); + semantic_tokens +} + +#[cfg(test)] +mod tests { + use crate::tests::compile_test_file; + use proc_macro_crate::bench_test; + + use super::semantic_tokens_full; + + #[test] + #[bench_test] + 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), + ]; + let res = semantic_tokens_full(&file, &gs); + if let Some(tokens) = res { + match &tokens { + lsp_types::SemanticTokensResult::Tokens(tokens) => { + let get: Vec<(u32, u32, u32, u32)> = tokens + .data + .iter() + .map(|token| { + ( + token.delta_line as u32, + token.delta_start as u32, + token.length as u32, + token.token_type as u32, + ) + }) + .collect(); + assert_eq!(get, expected); + } + lsp_types::SemanticTokensResult::Partial(_) => { + panic!("test failed") + } + } + } + } +} diff --git a/kclvm/tools/src/LSP/src/test_data/sema_token.k b/kclvm/tools/src/LSP/src/test_data/sema_token.k new file mode 100644 index 000000000..2b64e71ed --- /dev/null +++ b/kclvm/tools/src/LSP/src/test_data/sema_token.k @@ -0,0 +1,8 @@ +type num = int | float +schema Persons: + name: str + +p5: Persons = Persons { + name: "alice" +} +n: num = 1 \ No newline at end of file