From 8b6d6708e08fdc1c3e4c8f47f2031bfa62662cc1 Mon Sep 17 00:00:00 2001 From: blindfs Date: Wed, 9 Apr 2025 13:23:42 +0800 Subject: [PATCH] fix: module record var shadowing issue, perf optimization --- Cargo.lock | 1 + crates/nu-lsp/Cargo.toml | 1 + crates/nu-lsp/src/ast.rs | 77 +++++++++++++++----------- crates/nu-lsp/src/goto.rs | 8 +-- crates/nu-lsp/src/hover.rs | 2 +- crates/nu-lsp/src/lib.rs | 4 +- crates/nu-lsp/src/symbols.rs | 8 +-- crates/nu-lsp/src/workspace.rs | 83 ++++++++++++++++------------- tests/fixtures/lsp/workspace/baz.nu | 2 + 9 files changed, 107 insertions(+), 79 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f6b9a92593..5ad5c1890b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3821,6 +3821,7 @@ dependencies = [ "lsp-server", "lsp-textdocument", "lsp-types", + "memchr", "miette", "nu-cli", "nu-cmd-lang", diff --git a/crates/nu-lsp/Cargo.toml b/crates/nu-lsp/Cargo.toml index fbc444b418..be0d17ca7f 100644 --- a/crates/nu-lsp/Cargo.toml +++ b/crates/nu-lsp/Cargo.toml @@ -18,6 +18,7 @@ crossbeam-channel = { workspace = true } lsp-server = { workspace = true } lsp-textdocument = { workspace = true } lsp-types = { workspace = true } +memchr = { workspace = true } miette = { workspace = true } nucleo-matcher = { workspace = true } serde = { workspace = true } diff --git a/crates/nu-lsp/src/ast.rs b/crates/nu-lsp/src/ast.rs index c1a2a1c88d..c504817709 100644 --- a/crates/nu-lsp/src/ast.rs +++ b/crates/nu-lsp/src/ast.rs @@ -7,7 +7,7 @@ use nu_protocol::{ use std::sync::Arc; /// Adjust span if quoted -fn strip_quotes(span: Span, working_set: &StateWorkingSet) -> (Vec, Span) { +fn strip_quotes(span: Span, working_set: &StateWorkingSet) -> (Box<[u8]>, Span) { let text = working_set.get_span_contents(span); if text.len() > 1 && ((text.starts_with(b"\"") && text.ends_with(b"\"")) @@ -16,11 +16,11 @@ fn strip_quotes(span: Span, working_set: &StateWorkingSet) -> (Vec, Span) { ( text.get(1..text.len() - 1) .expect("Invalid quoted span!") - .to_vec(), + .into(), Span::new(span.start.saturating_add(1), span.end.saturating_sub(1)), ) } else { - (text.to_vec(), span) + (text.into(), span) } } @@ -120,8 +120,8 @@ fn try_find_id_in_mod( return None; } } - let check_location = |span: &Span| location.is_none_or(|pos| span.contains(*pos)); + let check_location = |span: &Span| location.is_none_or(|pos| span.contains(*pos)); call.arguments.first().and_then(|arg| { if !check_location(&arg.span()) { return None; @@ -147,7 +147,7 @@ fn try_find_id_in_mod( }) .map(ModuleId::new) })?; - let found_id = Id::Module(module_id, name.as_bytes().to_vec()); + let found_id = Id::Module(module_id, name.as_bytes().into()); let found_span = strip_quotes(arg.span(), working_set).1; id_ref .is_none_or(|id_r| found_id == *id_r) @@ -185,37 +185,42 @@ fn try_find_id_in_use( let find_by_name = |name: &[u8]| { let module = working_set.get_module(module_id); match id_ref { - Some(Id::Variable(var_id_ref)) => module + Some(Id::Variable(var_id_ref, name_ref)) => module .constants .get(name) .cloned() .or_else(|| { - // TODO: This is for the module record variable: + // NOTE: This is for the module record variable: // https://www.nushell.sh/book/modules/using_modules.html#importing-constants // The definition span is located at the head of the `use` command. - // This method won't work if the variable is later shadowed by a new definition. - let var_id = working_set.find_variable(name)?; - call.span() - .contains_span(working_set.get_variable(var_id).declaration_span) - .then_some(var_id) + (name_ref.as_ref() == name + && call + .span() + .contains_span(working_set.get_variable(*var_id_ref).declaration_span)) + .then_some(*var_id_ref) }) - .and_then(|var_id| (var_id == *var_id_ref).then_some(Id::Variable(var_id))), + .and_then(|var_id| { + (var_id == *var_id_ref).then_some(Id::Variable(var_id, name.into())) + }), Some(Id::Declaration(decl_id_ref)) => module.decls.get(name).and_then(|decl_id| { (*decl_id == *decl_id_ref).then_some(Id::Declaration(*decl_id)) }), // this is only for argument `members` Some(Id::Module(module_id_ref, name_ref)) => { module.submodules.get(name).and_then(|module_id| { - (*module_id == *module_id_ref && name_ref == name) - .then_some(Id::Module(*module_id, name.to_vec())) + (*module_id == *module_id_ref && name_ref.as_ref() == name) + .then_some(Id::Module(*module_id, name.into())) }) } None => module .submodules .get(name) - .map(|id| Id::Module(*id, name.to_vec())) + .map(|id| Id::Module(*id, name.into())) .or(module.decls.get(name).cloned().map(Id::Declaration)) - .or(module.constants.get(name).cloned().map(Id::Variable)), + .or(module + .constants + .get(name) + .map(|id| Id::Variable(*id, name.into()))), _ => None, } }; @@ -227,14 +232,14 @@ fn try_find_id_in_use( let (span_content, clean_span) = strip_quotes(span, working_set); if let Some(Id::Module(id_ref, name_ref)) = id_ref { // still need to check the rest, if id not matched - if module_id == *id_ref && *name_ref == span_content { - return Some((Id::Module(module_id, span_content.to_vec()), clean_span)); + if module_id == *id_ref && name_ref == &span_content { + return Some((Id::Module(module_id, span_content), clean_span)); } } if let Some(pos) = location { // first argument of `use`/`hide` should always be module name if span.contains(*pos) { - return Some((Id::Module(module_id, span_content.to_vec()), clean_span)); + return Some((Id::Module(module_id, span_content), clean_span)); } } @@ -308,7 +313,7 @@ fn try_find_id_in_overlay( else { return None; }; - let found_id = Id::Module(*module_id, name.as_bytes().to_vec()); + let found_id = Id::Module(*module_id, name.as_bytes().into()); id_ref .is_none_or(|id_r| found_id == *id_r) .then_some((found_id, strip_quotes(span, working_set).1)) @@ -317,7 +322,7 @@ fn try_find_id_in_overlay( let module_from_overlay_name = |name: &str, span: Span| { let found_id = Id::Module( working_set.find_overlay(name.as_bytes())?.origin, - name.as_bytes().to_vec(), + name.as_bytes().into(), ); id_ref .is_none_or(|id_r| found_id == *id_r) @@ -350,6 +355,19 @@ fn try_find_id_in_overlay( None } +/// For variable references `$foo` +fn strip_dollar_sign(span: Span, working_set: &StateWorkingSet<'_>) -> (Box<[u8]>, Span) { + let content = working_set.get_span_contents(span); + if content.starts_with(b"$") { + ( + content[1..].into(), + Span::new(span.start.saturating_add(1), span.end), + ) + } else { + (content.into(), span) + } +} + fn find_id_in_expr( expr: &Expression, working_set: &StateWorkingSet, @@ -361,12 +379,11 @@ fn find_id_in_expr( } let span = expr.span; match &expr.expr { - Expr::VarDecl(var_id) => FindMapResult::Found((Id::Variable(*var_id), span)), // trim leading `$` sign - Expr::Var(var_id) => FindMapResult::Found(( - Id::Variable(*var_id), - Span::new(span.start.saturating_add(1), span.end), - )), + Expr::VarDecl(var_id) | Expr::Var(var_id) => { + let (name, clean_span) = strip_dollar_sign(span, working_set); + FindMapResult::Found((Id::Variable(*var_id, name), clean_span)) + } Expr::Call(call) => { if call.head.contains(*location) { FindMapResult::Found((Id::Declaration(call.decl_id), call.head)) @@ -408,7 +425,7 @@ fn find_id_in_expr( } } Expr::Overlay(Some(module_id)) => { - FindMapResult::Found((Id::Module(*module_id, vec![]), span)) + FindMapResult::Found((Id::Module(*module_id, [].into()), span)) } // terminal value expressions Expr::Bool(_) @@ -445,12 +462,12 @@ fn find_reference_by_id_in_expr( ) -> Option> { let closure = |e| find_reference_by_id_in_expr(e, working_set, id); match (&expr.expr, id) { - (Expr::Var(vid1), Id::Variable(vid2)) if *vid1 == *vid2 => Some(vec![Span::new( + (Expr::Var(vid1), Id::Variable(vid2, _)) if *vid1 == *vid2 => Some(vec![Span::new( // we want to exclude the `$` sign for renaming expr.span.start.saturating_add(1), expr.span.end, )]), - (Expr::VarDecl(vid1), Id::Variable(vid2)) if *vid1 == *vid2 => Some(vec![expr.span]), + (Expr::VarDecl(vid1), Id::Variable(vid2, _)) if *vid1 == *vid2 => Some(vec![expr.span]), // also interested in `var_id` in call.arguments of `use` command // and `module_id` in `module` command (Expr::Call(call), _) => { diff --git a/crates/nu-lsp/src/goto.rs b/crates/nu-lsp/src/goto.rs index 995b05ae21..1a232c5986 100644 --- a/crates/nu-lsp/src/goto.rs +++ b/crates/nu-lsp/src/goto.rs @@ -51,7 +51,7 @@ impl LanguageServer { let block_id = working_set.get_decl(*decl_id).block_id()?; working_set.get_block(block_id).span } - Id::Variable(var_id) => { + Id::Variable(var_id, _) => { let var = working_set.get_variable(*var_id); Some(var.declaration_span) } @@ -147,7 +147,7 @@ mod tests { let mut none_existent_path = root(); none_existent_path.push("none-existent.nu"); let script = path_to_uri(&none_existent_path); - let resp = send_goto_definition_request(&client_connection, script.clone(), 0, 0); + let resp = send_goto_definition_request(&client_connection, script, 0, 0); assert_json_eq!(result_from_message(resp), serde_json::json!(null)); } @@ -195,7 +195,7 @@ mod tests { serde_json::json!({ "line": 1, "character": 10 }) ); - let resp = send_goto_definition_request(&client_connection, script.clone(), 2, 9); + let resp = send_goto_definition_request(&client_connection, script, 2, 9); assert_json_eq!( result_from_message(resp).pointer("/range/start").unwrap(), serde_json::json!({ "line": 1, "character": 17 }) @@ -451,7 +451,7 @@ mod tests { serde_json::json!({ "line": 0, "character": 0 }) ); - let resp = send_goto_definition_request(&client_connection, script.clone(), 2, 30); + let resp = send_goto_definition_request(&client_connection, script, 2, 30); assert_json_eq!( result_from_message(resp).pointer("/range/start").unwrap(), serde_json::json!({ "line": 0, "character": 0 }) diff --git a/crates/nu-lsp/src/hover.rs b/crates/nu-lsp/src/hover.rs index de5b7715ea..83b64d6f64 100644 --- a/crates/nu-lsp/src/hover.rs +++ b/crates/nu-lsp/src/hover.rs @@ -142,7 +142,7 @@ impl LanguageServer { }; match id { - Id::Variable(var_id) => { + Id::Variable(var_id, _) => { let var = working_set.get_variable(var_id); let value = var .const_val diff --git a/crates/nu-lsp/src/lib.rs b/crates/nu-lsp/src/lib.rs index 9c98c2d500..822d7a0849 100644 --- a/crates/nu-lsp/src/lib.rs +++ b/crates/nu-lsp/src/lib.rs @@ -40,10 +40,10 @@ mod workspace; #[derive(Debug, Clone, PartialEq)] pub(crate) enum Id { - Variable(VarId), + Variable(VarId, Box<[u8]>), Declaration(DeclId), Value(Type), - Module(ModuleId, Vec), + Module(ModuleId, Box<[u8]>), CellPath(VarId, Vec), External(String), } diff --git a/crates/nu-lsp/src/symbols.rs b/crates/nu-lsp/src/symbols.rs index b3b6e81c13..9bb3666dee 100644 --- a/crates/nu-lsp/src/symbols.rs +++ b/crates/nu-lsp/src/symbols.rs @@ -110,7 +110,7 @@ impl SymbolCache { range: span_to_range(&span, doc, doc_span.start), }) } - Id::Variable(var_id) => { + Id::Variable(var_id, _) => { let var = working_set.get_variable(var_id); let span = var.declaration_span; if !doc_span.contains(span.start) || span.end == span.start { @@ -157,7 +157,7 @@ impl SymbolCache { .chain((0..working_set.num_vars()).filter_map(|id| { Self::get_symbol_by_id( working_set, - Id::Variable(VarId::new(id)), + Id::Variable(VarId::new(id), [].into()), doc, &cached_file.covered_span, ) @@ -165,7 +165,7 @@ impl SymbolCache { .chain((0..working_set.num_modules()).filter_map(|id| { Self::get_symbol_by_id( working_set, - Id::Module(ModuleId::new(id), vec![]), + Id::Module(ModuleId::new(id), [].into()), doc, &cached_file.covered_span, ) @@ -361,7 +361,7 @@ mod tests { let script = path_to_uri(&script); open_unchecked(&client_connection, script.clone()); - let resp = send_document_symbol_request(&client_connection, script.clone()); + let resp = send_document_symbol_request(&client_connection, script); assert_json_eq!(result_from_message(resp), serde_json::json!([])); } diff --git a/crates/nu-lsp/src/workspace.rs b/crates/nu-lsp/src/workspace.rs index b581d0a9ef..d5e7cebe21 100644 --- a/crates/nu-lsp/src/workspace.rs +++ b/crates/nu-lsp/src/workspace.rs @@ -8,6 +8,7 @@ use lsp_types::{ PrepareRenameResponse, ProgressToken, Range, ReferenceParams, RenameParams, TextDocumentPositionParams, TextEdit, Uri, WorkspaceEdit, WorkspaceFolder, }; +use memchr::memmem; use miette::{miette, IntoDiagnostic, Result}; use nu_glob::Uninterruptible; use nu_protocol::{ @@ -59,13 +60,29 @@ struct IDTracker { /// Span of the original instance under the cursor pub span: Span, /// Name of the definition - pub name: String, + pub name: Box<[u8]>, /// Span of the original file where the request comes from pub file_span: Span, /// The redundant parsing should only happen once pub renewed: bool, } +impl IDTracker { + fn new(id: Id, span: Span, file_span: Span, working_set: &StateWorkingSet) -> Self { + let name = match &id { + Id::Variable(_, name) | Id::Module(_, name) => name.clone(), + _ => working_set.get_span_contents(span).into(), + }; + Self { + id, + span, + name, + file_span, + renewed: false, + } + } +} + impl LanguageServer { /// Get initial workspace folders from initialization response pub(crate) fn initialize_workspace_folders( @@ -101,14 +118,9 @@ impl LanguageServer { let (id, cursor_span) = find_id(&block, &working_set, &location)?; let mut refs = find_reference_by_id(&block, &working_set, &id); let definition_span = Self::find_definition_span_by_id(&working_set, &id); - if let Some(extra_span) = Self::reference_not_in_ast( - &id, - &String::from_utf8_lossy(working_set.get_span_contents(cursor_span)), - &working_set, - definition_span, - file_span, - cursor_span, - ) { + if let Some(extra_span) = + Self::reference_not_in_ast(&id, &working_set, definition_span, file_span, cursor_span) + { if !refs.contains(&extra_span) { refs.push(extra_span); } @@ -181,16 +193,7 @@ impl LanguageServer { .to_owned() .unwrap_or(ProgressToken::Number(1)); - let id_tracker = IDTracker { - id, - span, - file_span, - name: String::from_utf8_lossy(working_set.get_span_contents(span)).to_string(), - renewed: false, - }; - // make sure the parsing result of current file is merged in the state - let engine_state = self.new_engine_state(Some(&path_uri)); - + let id_tracker = IDTracker::new(id, span, file_span, &working_set); self.channels = self .find_reference_in_workspace( engine_state, @@ -275,16 +278,7 @@ impl LanguageServer { .get_workspace_folder_by_uri(&path_uri) .ok_or_else(|| miette!("\nCurrent file is not in any workspace"))?; // now continue parsing on other files in the workspace - let id_tracker = IDTracker { - id, - span, - file_span, - name: String::from_utf8_lossy(working_set.get_span_contents(span)).to_string(), - renewed: false, - }; - // make sure the parsing result of current file is merged in the state - let engine_state = self.new_engine_state(Some(&path_uri)); - + let id_tracker = IDTracker::new(id, span, file_span, &working_set); self.channels = self .find_reference_in_workspace( engine_state, @@ -338,13 +332,12 @@ impl LanguageServer { /// which is not covered in the AST fn reference_not_in_ast( id: &Id, - name_ref: &str, working_set: &StateWorkingSet, definition_span: Option, file_span: Span, sample_span: Span, ) -> Option { - if let (Id::Variable(_), Some(decl_span)) = (&id, definition_span) { + if let (Id::Variable(_, name_ref), Some(decl_span)) = (&id, definition_span) { if file_span.contains_span(decl_span) && decl_span.end > decl_span.start { let content = working_set.get_span_contents(decl_span); let leading_dashes = content @@ -354,7 +347,7 @@ impl LanguageServer { .count(); let start = decl_span.start + leading_dashes; return content.get(leading_dashes..).and_then(|name| { - name.starts_with(name_ref.as_bytes()).then_some(Span { + name.starts_with(name_ref).then_some(Span { start, end: start + sample_span.end - sample_span.start, }) @@ -419,7 +412,7 @@ impl LanguageServer { let file = if let Some(file) = docs.get_document(&uri) { file } else { - let bytes = match fs::read(fp) { + let file_content_byte = match fs::read(fp) { Ok(it) => it, Err(_) => { // continue on fs error @@ -427,15 +420,18 @@ impl LanguageServer { } }; // skip if the file does not contain what we're looking for - let content_string = String::from_utf8_lossy(&bytes); - if !content_string.contains(&id_tracker.name) { + if memmem::find(&file_content_byte, id_tracker.name.as_ref()).is_none() { // progress without any data data_sender .send(InternalMessage::OnGoing(token.clone(), percentage)) .into_diagnostic()?; continue; } - &FullTextDocument::new("nu".to_string(), 0, content_string.into()) + &FullTextDocument::new( + "nu".to_string(), + 0, + String::from_utf8_lossy(&file_content_byte).into(), + ) }; let _ = Self::find_reference_in_file(&mut working_set, file, fp, &mut id_tracker) .map(|mut refs| { @@ -444,7 +440,6 @@ impl LanguageServer { .unwrap_or(Span::unknown()); if let Some(extra_span) = Self::reference_not_in_ast( &id_tracker.id, - &id_tracker.name, &working_set, definition_span, file_span, @@ -1141,7 +1136,7 @@ mod tests { let (client_connection, _recv) = initialize_language_server(None, None); open_unchecked(&client_connection, script.clone()); - let message = send_document_highlight_request(&client_connection, script, 8, 0); + let message = send_document_highlight_request(&client_connection, script.clone(), 8, 0); let Message::Response(r) = message else { panic!("unexpected message type"); }; @@ -1152,5 +1147,17 @@ mod tests { { "range": { "start": { "line": 8, "character": 1 }, "end": { "line": 8, "character": 8 } }, "kind": 1 }, ]), ); + + let message = send_document_highlight_request(&client_connection, script, 10, 7); + let Message::Response(r) = message else { + panic!("unexpected message type"); + }; + assert_json_eq!( + r.result, + serde_json::json!([ + { "range": { "start": { "line": 10, "character": 4 }, "end": { "line": 10, "character": 12 } }, "kind": 1 }, + { "range": { "start": { "line": 11, "character": 1 }, "end": { "line": 11, "character": 8 } }, "kind": 1 }, + ]), + ); } } diff --git a/tests/fixtures/lsp/workspace/baz.nu b/tests/fixtures/lsp/workspace/baz.nu index 2f1aea642b..c06848e75d 100644 --- a/tests/fixtures/lsp/workspace/baz.nu +++ b/tests/fixtures/lsp/workspace/baz.nu @@ -8,3 +8,5 @@ use ./foo.nu [ "mod name" cst_mod ] $cst_mod."sub module".var_name mod name sub module cmd name +let $cst_mod = 1 +$cst_mod