From f51a79181a3826c3c69d9bc7314c1ef80f97b2e2 Mon Sep 17 00:00:00 2001 From: zc he Date: Tue, 25 Feb 2025 21:14:48 +0800 Subject: [PATCH] feat(lsp): semantic tokens for highlighting internal commands with spaces (#15173) # Description We decided to move that specific highlighting task from tree-sitter-nu to lsp for various reasons. https://github.com/nushell/tree-sitter-nu/pull/184 # User-Facing Changes Before: image After: image # Tests + Formatting +1 # After Submitting --- crates/nu-lsp/src/lib.rs | 30 +++- crates/nu-lsp/src/semantic_tokens.rs | 168 ++++++++++++++++++ .../fixtures/lsp/semantic_tokens/internals.nu | 15 ++ 3 files changed, 210 insertions(+), 3 deletions(-) create mode 100644 crates/nu-lsp/src/semantic_tokens.rs create mode 100644 tests/fixtures/lsp/semantic_tokens/internals.nu diff --git a/crates/nu-lsp/src/lib.rs b/crates/nu-lsp/src/lib.rs index 95fac87f83..b6a638e2d7 100644 --- a/crates/nu-lsp/src/lib.rs +++ b/crates/nu-lsp/src/lib.rs @@ -6,7 +6,8 @@ use lsp_types::{ CompletionItem, CompletionItemKind, CompletionItemLabelDetails, CompletionParams, CompletionResponse, CompletionTextEdit, Documentation, Hover, HoverContents, HoverParams, InlayHint, Location, MarkupContent, MarkupKind, OneOf, Position, Range, ReferencesOptions, - RenameOptions, ServerCapabilities, TextDocumentSyncKind, TextEdit, Uri, + RenameOptions, SemanticToken, SemanticTokenType, SemanticTokensLegend, SemanticTokensOptions, + SemanticTokensServerCapabilities, ServerCapabilities, TextDocumentSyncKind, TextEdit, Uri, WorkDoneProgressOptions, WorkspaceFolder, WorkspaceFoldersServerCapabilities, WorkspaceServerCapabilities, }; @@ -33,6 +34,7 @@ mod diagnostics; mod goto; mod hints; mod notification; +mod semantic_tokens; mod symbols; mod workspace; @@ -53,6 +55,7 @@ pub struct LanguageServer { initial_engine_state: EngineState, symbol_cache: SymbolCache, inlay_hints: BTreeMap>, + semantic_tokens: BTreeMap>, workspace_folders: BTreeMap, /// for workspace wide requests occurrences: BTreeMap>, @@ -106,6 +109,7 @@ impl LanguageServer { initial_engine_state: engine_state, symbol_cache: SymbolCache::new(), inlay_hints: BTreeMap::new(), + semantic_tokens: BTreeMap::new(), workspace_folders: BTreeMap::new(), occurrences: BTreeMap::new(), channels: None, @@ -143,6 +147,17 @@ impl LanguageServer { ..Default::default() }), workspace_symbol_provider: Some(OneOf::Left(true)), + semantic_tokens_provider: Some( + SemanticTokensServerCapabilities::SemanticTokensOptions(SemanticTokensOptions { + // NOTE: only internal command names with space supported for now + legend: SemanticTokensLegend { + token_types: vec![SemanticTokenType::FUNCTION], + token_modifiers: vec![], + }, + full: Some(lsp_types::SemanticTokensFullOptions::Bool(true)), + ..Default::default() + }), + ), ..Default::default() }) .expect("Must be serializable"); @@ -202,6 +217,11 @@ impl LanguageServer { request::InlayHintRequest::METHOD => { Self::handle_lsp_request(request, |params| self.get_inlay_hints(params)) } + request::SemanticTokensFullRequest::METHOD => { + Self::handle_lsp_request(request, |params| { + self.get_semantic_tokens(params) + }) + } request::PrepareRenameRequest::METHOD => { let id = request.id.clone(); if let Err(e) = self.prepare_rename(request) { @@ -336,7 +356,7 @@ impl LanguageServer { &mut self, engine_state: &'a mut EngineState, uri: &Uri, - need_hints: bool, + need_extra_info: bool, ) -> Option<(Arc, Span, StateWorkingSet<'a>)> { let mut working_set = StateWorkingSet::new(engine_state); let docs = self.docs.lock().ok()?; @@ -347,10 +367,14 @@ impl LanguageServer { let _ = working_set.files.push(file_path.clone(), Span::unknown()); let block = nu_parser::parse(&mut working_set, Some(file_path_str), contents, false); let span = working_set.get_span_for_filename(file_path_str)?; - if need_hints { + if need_extra_info { let file_inlay_hints = Self::extract_inlay_hints(&working_set, &block, span.start, file); self.inlay_hints.insert(uri.clone(), file_inlay_hints); + let file_semantic_tokens = + Self::extract_semantic_tokens(&working_set, &block, span.start, file); + self.semantic_tokens + .insert(uri.clone(), file_semantic_tokens); } if self.need_parse { // TODO: incremental parsing diff --git a/crates/nu-lsp/src/semantic_tokens.rs b/crates/nu-lsp/src/semantic_tokens.rs new file mode 100644 index 0000000000..2d6e6e939d --- /dev/null +++ b/crates/nu-lsp/src/semantic_tokens.rs @@ -0,0 +1,168 @@ +use std::sync::Arc; + +use lsp_textdocument::FullTextDocument; +use lsp_types::{SemanticToken, SemanticTokens, SemanticTokensParams}; +use nu_protocol::{ + ast::{Block, Expr, Expression, Traverse}, + engine::StateWorkingSet, + Span, +}; + +use crate::{span_to_range, LanguageServer}; + +/// Important for keep spans in increasing order, +/// since `SemanticToken`s are created by relative positions +/// to one's previous token +/// +/// Currently supported types: +/// 1. internal command names with space +fn extract_semantic_tokens_from_expression( + expr: &Expression, + working_set: &StateWorkingSet, +) -> Option> { + let closure = |e| extract_semantic_tokens_from_expression(e, working_set); + match &expr.expr { + Expr::Call(call) => { + let command_name_bytes = working_set.get_span_contents(call.head); + let head_span = if command_name_bytes.contains(&b' ') + // Some keywords that are already highlighted properly, e.g. by tree-sitter-nu + && !command_name_bytes.starts_with(b"export") + && !command_name_bytes.starts_with(b"overlay") + { + vec![call.head] + } else { + vec![] + }; + let spans = head_span + .into_iter() + .chain( + call.arguments + .iter() + .filter_map(|arg| arg.expr()) + .flat_map(|e| e.flat_map(working_set, &closure)), + ) + .collect(); + Some(spans) + } + _ => None, + } +} + +impl LanguageServer { + pub(crate) fn get_semantic_tokens( + &mut self, + params: &SemanticTokensParams, + ) -> Option { + self.semantic_tokens + .get(¶ms.text_document.uri) + .map(|vec| SemanticTokens { + result_id: None, + data: vec.clone(), + }) + } + + pub(crate) fn extract_semantic_tokens( + working_set: &StateWorkingSet, + block: &Arc, + offset: usize, + file: &FullTextDocument, + ) -> Vec { + let spans = block.flat_map(working_set, &|e| { + extract_semantic_tokens_from_expression(e, working_set) + }); + let mut last_token_line = 0; + let mut last_token_char = 0; + let mut last_span = Span::unknown(); + let mut tokens = vec![]; + for sp in spans { + let range = span_to_range(&sp, file, offset); + // shouldn't happen + if sp < last_span { + continue; + } + // in case the start position is at the end of lastline + let real_start_char = if range.end.line != range.start.line { + 0 + } else { + range.start.character + }; + let mut delta_start = real_start_char; + if range.end.line == last_token_line { + delta_start -= last_token_char; + } + tokens.push(SemanticToken { + delta_start, + delta_line: range.end.line.saturating_sub(last_token_line), + length: range.end.character.saturating_sub(real_start_char), + // 0 means function in semantic_token_legend + token_type: 0, + token_modifiers_bitset: 0, + }); + last_span = sp; + last_token_line = range.end.line; + last_token_char = real_start_char; + } + tokens + } +} + +#[cfg(test)] +mod tests { + use crate::path_to_uri; + use crate::tests::{initialize_language_server, open_unchecked, result_from_message}; + use assert_json_diff::assert_json_eq; + use lsp_server::{Connection, Message}; + use lsp_types::{ + request::{Request, SemanticTokensFullRequest}, + TextDocumentIdentifier, Uri, WorkDoneProgressParams, + }; + use lsp_types::{PartialResultParams, SemanticTokensParams}; + use nu_test_support::fs::fixtures; + + fn send_semantic_token_request(client_connection: &Connection, uri: Uri) -> Message { + client_connection + .sender + .send(Message::Request(lsp_server::Request { + id: 1.into(), + method: SemanticTokensFullRequest::METHOD.to_string(), + params: serde_json::to_value(SemanticTokensParams { + text_document: TextDocumentIdentifier { uri }, + work_done_progress_params: WorkDoneProgressParams::default(), + partial_result_params: PartialResultParams::default(), + }) + .unwrap(), + })) + .unwrap(); + client_connection + .receiver + .recv_timeout(std::time::Duration::from_secs(2)) + .unwrap() + } + + #[test] + fn semantic_token_internals() { + let (client_connection, _recv) = initialize_language_server(None); + + let mut script = fixtures(); + script.push("lsp"); + script.push("semantic_tokens"); + script.push("internals.nu"); + let script = path_to_uri(&script); + + open_unchecked(&client_connection, script.clone()); + let resp = send_semantic_token_request(&client_connection, script.clone()); + + assert_json_eq!( + result_from_message(resp), + serde_json::json!( + { "data": [ + // delta_line, delta_start, length, token_type, token_modifiers_bitset + 0, 0, 13, 0, 0, + 1, 2, 10, 0, 0, + 7, 15, 13, 0, 0, + 0, 20, 10, 0, 0, + 4, 0, 7, 0, 0 + ]}) + ); + } +} diff --git a/tests/fixtures/lsp/semantic_tokens/internals.nu b/tests/fixtures/lsp/semantic_tokens/internals.nu new file mode 100644 index 0000000000..66cccb05be --- /dev/null +++ b/tests/fixtures/lsp/semantic_tokens/internals.nu @@ -0,0 +1,15 @@ +str substring 1.. +| ansi strip + +# User defined one +export def "foo bar" [] { + # inside a block + ( + # same line + "🤔🤖🐘" | str substring 1.. | ansi strip + ) +} + +foo bar + +overlay use foo