fix(lsp): more accurate PWD: from env -> parent dir of current file (#15470)

# Description

Some editors like neovim will provide "workspace root" as PWD, which can
mess up file completion results.

# User-Facing Changes

bug fix

# Tests + Formatting

adjusted

# After Submitting

crates/nu-lsp/src/completion.rs

@@ -29,7 +29,7 @@ impl LanguageServer {
                 .is_some_and(|c| c.is_whitespace() || "|(){}[]<>,:;".contains(c));
 
         self.need_parse |= need_fallback;
-        let engine_state = Arc::new(self.new_engine_state());
+        let engine_state = Arc::new(self.new_engine_state(Some(&path_uri)));
         let completer = NuCompleter::new(engine_state.clone(), Arc::new(Stack::new()));
         let results = if need_fallback {
             completer.fetch_completions_at(&file_text[..location], location)
@@ -264,10 +264,10 @@ mod tests {
         let resp = send_complete_request(&client_connection, script.clone(), 2, 18);
         assert!(result_from_message(resp).as_array().unwrap().contains(
             &serde_json::json!({
-                "label": "LICENSE",
+                "label": "command.nu",
                 "labelDetails": { "description": "" },
                 "textEdit": { "range": { "start": { "line": 2, "character": 17 }, "end": { "line": 2, "character": 18 }, },
-                    "newText": "LICENSE"
+                    "newText": "command.nu"
                 },
                 "kind": 17
             })
@@ -337,10 +337,10 @@ mod tests {
         let resp = send_complete_request(&client_connection, script, 5, 4);
         assert!(result_from_message(resp).as_array().unwrap().contains(
             &serde_json::json!({
-                "label": "LICENSE",
+                "label": "cell_path.nu",
                 "labelDetails": { "description": "" },
                 "textEdit": { "range": { "start": { "line": 5, "character": 3 }, "end": { "line": 5, "character": 4 }, },
-                    "newText": "LICENSE"
+                    "newText": "cell_path.nu"
                 },
                 "kind": 17
             })

crates/nu-lsp/src/diagnostics.rs

@@ -7,7 +7,7 @@ use miette::{miette, IntoDiagnostic, Result};
 
 impl LanguageServer {
     pub(crate) fn publish_diagnostics_for_file(&mut self, uri: Uri) -> Result<()> {
-        let mut engine_state = self.new_engine_state();
+        let mut engine_state = self.new_engine_state(Some(&uri));
         engine_state.generate_nu_constant();
 
         let Some((_, span, working_set)) = self.parse_file(&mut engine_state, &uri, true) else {

crates/nu-lsp/src/goto.rs

@@ -77,13 +77,12 @@ impl LanguageServer {
         &mut self,
         params: &GotoDefinitionParams,
     ) -> Option<GotoDefinitionResponse> {
-        let mut engine_state = self.new_engine_state();
-
         let path_uri = params
             .text_document_position_params
             .text_document
             .uri
             .to_owned();
+        let mut engine_state = self.new_engine_state(Some(&path_uri));
         let (working_set, id, _, _) = self
             .parse_and_find(
                 &mut engine_state,

crates/nu-lsp/src/hover.rs

@@ -129,13 +129,12 @@ impl LanguageServer {
     }
 
     pub(crate) fn hover(&mut self, params: &HoverParams) -> Option<Hover> {
-        let mut engine_state = self.new_engine_state();
-
         let path_uri = params
             .text_document_position_params
             .text_document
             .uri
             .to_owned();
+        let mut engine_state = self.new_engine_state(Some(&path_uri));
         let (working_set, id, _, _) = self
             .parse_and_find(
                 &mut engine_state,

crates/nu-lsp/src/lib.rs

@@ -13,7 +13,7 @@ use miette::{miette, IntoDiagnostic, Result};
 use nu_protocol::{
     ast::{Block, PathMember},
     engine::{EngineState, StateDelta, StateWorkingSet},
-    DeclId, ModuleId, Span, Type, VarId,
+    DeclId, ModuleId, Span, Type, Value, VarId,
 };
 use std::{
     collections::BTreeMap,
@@ -315,13 +315,26 @@ impl LanguageServer {
         Ok(reset)
     }
 
-    pub(crate) fn new_engine_state(&self) -> EngineState {
+    /// Create a clone of the initial_engine_state with:
+    ///
+    /// * PWD set to the parent directory of given uri. Fallback to `$env.PWD` if None.
+    /// * `StateDelta` cache merged
+    pub(crate) fn new_engine_state(&self, uri: Option<&Uri>) -> EngineState {
         let mut engine_state = self.initial_engine_state.clone();
-        let cwd = std::env::current_dir().expect("Could not get current working directory.");
-        engine_state.add_env_var(
-            "PWD".into(),
-            nu_protocol::Value::test_string(cwd.to_string_lossy()),
-        );
+        match uri {
+            Some(uri) => {
+                let path = uri_to_path(uri);
+                if let Some(path) = path.parent() {
+                    engine_state
+                        .add_env_var("PWD".into(), Value::test_string(path.to_string_lossy()))
+                };
+            }
+            None => {
+                let cwd =
+                    std::env::current_dir().expect("Could not get current working directory.");
+                engine_state.add_env_var("PWD".into(), Value::test_string(cwd.to_string_lossy()));
+            }
+        }
         // merge the cached `StateDelta` if text not changed
         if !self.need_parse {
             engine_state
@@ -350,7 +363,7 @@ impl LanguageServer {
         engine_state: &'a mut EngineState,
         uri: &Uri,
         pos: Position,
-    ) -> Result<(StateWorkingSet<'a>, Id, Span, usize)> {
+    ) -> Result<(StateWorkingSet<'a>, Id, Span, Span)> {
         let (block, file_span, working_set) = self
             .parse_file(engine_state, uri, false)
             .ok_or_else(|| miette!("\nFailed to parse current file"))?;
@@ -365,7 +378,7 @@ impl LanguageServer {
         let location = file.offset_at(pos) as usize + file_span.start;
         let (id, span) = ast::find_id(&block, &working_set, &location)
             .ok_or_else(|| miette!("\nFailed to find current name"))?;
-        Ok((working_set, id, span, file_span.start))
+        Ok((working_set, id, span, file_span))
     }
 
     pub(crate) fn parse_file<'a>(
@@ -458,10 +471,7 @@ mod tests {
         engine_state.generate_nu_constant();
         assert!(load_standard_library(&mut engine_state).is_ok());
         let cwd = std::env::current_dir().expect("Could not get current working directory.");
-        engine_state.add_env_var(
-            "PWD".into(),
-            nu_protocol::Value::test_string(cwd.to_string_lossy()),
-        );
+        engine_state.add_env_var("PWD".into(), Value::test_string(cwd.to_string_lossy()));
         if let Some(code) = nu_config_code {
             assert!(merge_input(code.as_bytes(), &mut engine_state, &mut Stack::new()).is_ok());
         }

crates/nu-lsp/src/signature.rs

@@ -78,7 +78,7 @@ impl LanguageServer {
         let file_text = file.get_content(None).to_owned();
         drop(docs);
 
-        let engine_state = self.new_engine_state();
+        let engine_state = self.new_engine_state(Some(&path_uri));
         let mut working_set = StateWorkingSet::new(&engine_state);
 
         // NOTE: in case the cursor is at the end of the call expression

crates/nu-lsp/src/symbols.rs

@@ -270,8 +270,8 @@ impl LanguageServer {
         &mut self,
         params: &DocumentSymbolParams,
     ) -> Option<DocumentSymbolResponse> {
-        let engine_state = self.new_engine_state();
         let uri = params.text_document.uri.to_owned();
+        let engine_state = self.new_engine_state(Some(&uri));
         let docs = self.docs.lock().ok()?;
         self.symbol_cache.update(&uri, &engine_state, &docs);
         self.symbol_cache
@@ -284,7 +284,7 @@ impl LanguageServer {
         params: &WorkspaceSymbolParams,
     ) -> Option<WorkspaceSymbolResponse> {
         if self.symbol_cache.any_dirty() {
-            let engine_state = self.new_engine_state();
+            let engine_state = self.new_engine_state(None);
             let docs = self.docs.lock().ok()?;
             self.symbol_cache.update_all(&engine_state, &docs);
         }

crates/nu-lsp/src/workspace.rs

@@ -1,5 +1,5 @@
 use crate::{
-    ast::{find_id, find_reference_by_id},
+    ast::{self, find_id, find_reference_by_id},
     path_to_uri, span_to_range, uri_to_path, Id, LanguageServer,
 };
 use lsp_textdocument::FullTextDocument;
@@ -46,6 +46,26 @@ fn find_nu_scripts_in_folder(folder_uri: &Uri) -> Result<nu_glob::Paths> {
     nu_glob::glob(&pattern, Uninterruptible).into_diagnostic()
 }
 
+/// HACK: when current file is imported (use keyword) by others in the workspace,
+/// it will get parsed a second time via `parse_module_block`, so that its definitions'
+/// ids are renewed, making it harder to track the references.
+///
+/// FIXME: cross-file shadowing can still cause false-positive/false-negative cases
+///
+/// This is a workaround to track the new id
+struct IDTracker {
+    /// ID to search, renewed on `parse_module_block`
+    pub id: Id,
+    /// Span of the original instance under the cursor
+    pub span: Span,
+    /// Name of the definition
+    pub name: String,
+    /// 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 LanguageServer {
     /// Get initial workspace folders from initialization response
     pub(crate) fn initialize_workspace_folders(
@@ -66,12 +86,12 @@ impl LanguageServer {
         &mut self,
         params: &DocumentHighlightParams,
     ) -> Option<Vec<DocumentHighlight>> {
-        let mut engine_state = self.new_engine_state();
         let path_uri = params
             .text_document_position_params
             .text_document
             .uri
             .to_owned();
+        let mut engine_state = self.new_engine_state(Some(&path_uri));
         let (block, file_span, working_set) =
             self.parse_file(&mut engine_state, &path_uri, false)?;
         let docs = &self.docs.lock().ok()?;
@@ -137,31 +157,38 @@ impl LanguageServer {
         timeout: u128,
     ) -> Option<Vec<Location>> {
         self.occurrences = BTreeMap::new();
-        let mut engine_state = self.new_engine_state();
         let path_uri = params.text_document_position.text_document.uri.to_owned();
-        let (_, id, span, _) = self
+        let mut engine_state = self.new_engine_state(Some(&path_uri));
+
+        let (working_set, id, span, file_span) = self
             .parse_and_find(
                 &mut engine_state,
                 &path_uri,
                 params.text_document_position.position,
             )
             .ok()?;
-        // have to clone it again in order to move to another thread
-        let engine_state = self.new_engine_state();
         let current_workspace_folder = self.get_workspace_folder_by_uri(&path_uri)?;
         let token = params
             .work_done_progress_params
             .work_done_token
             .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,
+        };
+
         self.channels = self
             .find_reference_in_workspace(
                 engine_state,
                 current_workspace_folder,
-                id,
-                span,
                 token.clone(),
                 "Finding references ...".to_string(),
+                id_tracker,
             )
             .ok();
         // TODO: WorkDoneProgress -> PartialResults for quicker response
@@ -200,10 +227,10 @@ impl LanguageServer {
             serde_json::from_value(request.params).into_diagnostic()?;
         self.occurrences = BTreeMap::new();
 
-        let mut engine_state = self.new_engine_state();
         let path_uri = params.text_document.uri.to_owned();
+        let mut engine_state = self.new_engine_state(Some(&path_uri));
 
-        let (working_set, id, span, file_offset) =
+        let (working_set, id, span, file_span) =
             self.parse_and_find(&mut engine_state, &path_uri, params.position)?;
 
         if let Id::Value(_) = id {
@@ -222,7 +249,7 @@ impl LanguageServer {
         let file = docs
             .get_document(&path_uri)
             .ok_or_else(|| miette!("\nFailed to get document"))?;
-        let range = span_to_range(&span, file, file_offset);
+        let range = span_to_range(&span, file, file_span.start);
         let response = PrepareRenameResponse::Range(range);
         self.connection
             .sender
@@ -233,20 +260,24 @@ impl LanguageServer {
             }))
             .into_diagnostic()?;
 
-        // have to clone it again in order to move to another thread
-        let engine_state = self.new_engine_state();
         let current_workspace_folder = self
             .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,
+        };
         self.channels = self
             .find_reference_in_workspace(
                 engine_state,
                 current_workspace_folder,
-                id,
-                span,
                 ProgressToken::Number(0),
                 "Preparing rename ...".to_string(),
+                id_tracker,
             )
             .ok();
         Ok(())
@@ -256,7 +287,7 @@ impl LanguageServer {
         working_set: &mut StateWorkingSet,
         file: &FullTextDocument,
         fp: &Path,
-        id: &Id,
+        id_tracker: &mut IDTracker,
     ) -> Option<Vec<Span>> {
         let block = nu_parser::parse(
             working_set,
@@ -264,7 +295,25 @@ impl LanguageServer {
             file.get_content(None).as_bytes(),
             false,
         );
-        let references: Vec<Span> = find_reference_by_id(&block, working_set, id);
+        // NOTE: Renew the id if there's a module with the same span as the original file.
+        // This requires that the initial parsing results get merged in the engine_state,
+        // typically they're cached with diagnostics before the prepare_rename/references requests,
+        // so that we don't need to clone and merge delta again.
+        if (!id_tracker.renewed)
+            && working_set
+                .find_module_by_span(id_tracker.file_span)
+                .is_some()
+        {
+            if let Some(new_block) = working_set.find_block_by_span(id_tracker.file_span) {
+                if let Some((new_id, _)) =
+                    ast::find_id(&new_block, working_set, &id_tracker.span.start)
+                {
+                    id_tracker.id = new_id;
+                }
+            }
+            id_tracker.renewed = true;
+        }
+        let references: Vec<Span> = find_reference_by_id(&block, working_set, &id_tracker.id);
 
         // add_block to avoid repeated parsing
         working_set.add_block(block);
@@ -304,10 +353,9 @@ impl LanguageServer {
         &self,
         engine_state: EngineState,
         current_workspace_folder: WorkspaceFolder,
-        id: Id,
-        span: Span,
         token: ProgressToken,
         message: String,
+        mut id_tracker: IDTracker,
     ) -> Result<(
         crossbeam_channel::Sender<bool>,
         Arc<crossbeam_channel::Receiver<InternalMessage>>,
@@ -333,7 +381,7 @@ impl LanguageServer {
                 .filter_map(|p| p.ok())
                 .collect();
             let len = scripts.len();
-            let definition_span = Self::find_definition_span_by_id(&working_set, &id);
+            let definition_span = Self::find_definition_span_by_id(&working_set, &id_tracker.id);
 
             for (i, fp) in scripts.iter().enumerate() {
                 #[cfg(test)]
@@ -363,9 +411,7 @@ impl LanguageServer {
                     };
                     // skip if the file does not contain what we're looking for
                     let content_string = String::from_utf8_lossy(&bytes);
-                    let text_to_search =
-                        String::from_utf8_lossy(working_set.get_span_contents(span));
-                    if !content_string.contains(text_to_search.as_ref()) {
+                    if !content_string.contains(&id_tracker.name) {
                         // progress without any data
                         data_sender
                             .send(InternalMessage::OnGoing(token.clone(), percentage))
@@ -374,17 +420,17 @@ impl LanguageServer {
                     }
                     &FullTextDocument::new("nu".to_string(), 0, content_string.into())
                 };
-                let _ = Self::find_reference_in_file(&mut working_set, file, fp, &id).map(
-                    |mut refs| {
+                let _ = Self::find_reference_in_file(&mut working_set, file, fp, &mut id_tracker)
+                    .map(|mut refs| {
                         let file_span = working_set
                             .get_span_for_filename(fp.to_string_lossy().as_ref())
                             .unwrap_or(Span::unknown());
                         if let Some(extra_span) = Self::reference_not_in_ast(
-                            &id,
+                            &id_tracker.id,
                             &working_set,
                             definition_span,
                             file_span,
-                            span,
+                            id_tracker.span,
                         ) {
                             if !refs.contains(&extra_span) {
                                 refs.push(extra_span)
@@ -400,8 +446,7 @@ impl LanguageServer {
                         data_sender
                             .send(InternalMessage::OnGoing(token.clone(), percentage))
                             .ok();
-                    },
-                );
+                    });
             }
             data_sender
                 .send(InternalMessage::Finished(token.clone()))

tests/fixtures/lsp/completion/command.nu

@@ -1,6 +1,6 @@
 config n
 config n foo bar -
-config n foo bar l --l
+config n foo bar c --l
 
 # detail
 def "config n foo bar" [

tests/fixtures/lsp/completion/fallback.nu

@@ -3,6 +3,6 @@ let greeting = "Hello"
 echo $gre
 | st
 
-ls l
+ls c
 
 $greeting not-h