fix(lsp): parser_info based id detection for use/overlay keywords (#15517)

# Description

Now, with PWD correctly set in #15470 , identifiers in
`use/hide/overlay` commands can be identified using a more robust
method, i.e. module_id from `parser_info`.

# User-Facing Changes

bug fix

# Tests + Formatting

+1 (fails without this PR)

# After Submitting

crates/nu-lsp/src/ast.rs

@@ -29,9 +29,7 @@ fn try_find_id_in_misc(
     match call_name {
         "def" | "export def" => try_find_id_in_def(call, working_set, location, id_ref),
         "module" | "export module" => try_find_id_in_mod(call, working_set, location, id_ref),
-        "use" | "export use" | "hide" => {
-            try_find_id_in_use(call, working_set, location, id_ref, call_name)
-        }
+        "use" | "export use" | "hide" => try_find_id_in_use(call, working_set, location, id_ref),
         "overlay use" | "overlay hide" => {
             try_find_id_in_overlay(call, working_set, location, id_ref)
         }
@@ -129,9 +127,6 @@ fn try_find_id_in_mod(
 
 /// Find id in use/hide command
 /// `hide foo.nu bar` or `use foo.nu [bar baz]`
-/// NOTE: `call.parser_info` contains a 'import_pattern' field for `use` commands,
-/// but sometimes it is missing, so fall back to `call_name == "use"` here.
-/// One drawback is that the `module_id` is harder to get
 ///
 /// # Arguments
 /// - `location`: None if no `contains` check required
@@ -141,43 +136,58 @@ fn try_find_id_in_use(
     working_set: &StateWorkingSet,
     location: Option<&usize>,
     id: Option<&Id>,
-    call_name: &str,
 ) -> Option<(Id, Span)> {
-    // TODO: for keyword `hide`, the decl/var is already hidden in working_set,
-    // this function will always return None.
-    let find_by_name = |name: &[u8]| match id {
-        Some(Id::Variable(var_id_ref)) => working_set
-            .find_variable(name)
-            .and_then(|var_id| (var_id == *var_id_ref).then_some(Id::Variable(var_id))),
-        Some(Id::Declaration(decl_id_ref)) => working_set
-            .find_decl(name)
-            .and_then(|decl_id| (decl_id == *decl_id_ref).then_some(Id::Declaration(decl_id))),
-        Some(Id::Module(module_id_ref)) => working_set
-            .find_module(name)
-            .and_then(|module_id| (module_id == *module_id_ref).then_some(Id::Module(module_id))),
-        None => working_set
-            .find_module(name)
-            .map(Id::Module)
-            .or(working_set.find_decl(name).map(Id::Declaration))
-            .or(working_set.find_variable(name).map(Id::Variable)),
-        _ => None,
+    // NOTE: `call.parser_info` contains a 'import_pattern' field for `use`/`hide` commands,
+    // If it's missing, usually it means the PWD env is not correctly set,
+    // checkout `new_engine_state` in lib.rs
+    let Expression {
+        expr: Expr::ImportPattern(import_pattern),
+        ..
+    } = call.get_parser_info("import_pattern")?
+    else {
+        return None;
+    };
+    let module_id = import_pattern.head.id?;
+
+    let find_by_name = |name: &[u8]| {
+        let module = working_set.get_module(module_id);
+        match id {
+            Some(Id::Variable(var_id_ref)) => module
+                .constants
+                .get(name)
+                .and_then(|var_id| (*var_id == *var_id_ref).then_some(Id::Variable(*var_id))),
+            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)) => module.submodules.get(name).and_then(|module_id| {
+                (*module_id == *module_id_ref).then_some(Id::Module(*module_id))
+            }),
+            None => module
+                .submodules
+                .get(name)
+                .cloned()
+                .map(Id::Module)
+                .or(module.decls.get(name).cloned().map(Id::Declaration))
+                .or(module.constants.get(name).cloned().map(Id::Variable)),
+            _ => None,
+        }
     };
     let check_location = |span: &Span| location.is_none_or(|pos| span.contains(*pos));
 
     // Get module id if required
     let module_name = call.arguments.first()?;
     let span = module_name.span();
-    if let Some(Id::Module(_)) = id {
+    if let Some(Id::Module(id_ref)) = id {
         // still need to check the rest, if id not matched
-        if let Some(res) = get_matched_module_id(working_set, span, id) {
-            return Some(res);
+        if module_id == *id_ref {
+            return Some((Id::Module(module_id), strip_quotes(span, working_set)));
         }
     }
     if let Some(pos) = location {
-        // first argument of `use` should always be module name
-        // while it is optional in `hide`
-        if span.contains(*pos) && call_name != "hide" {
-            return get_matched_module_id(working_set, span, id);
+        // first argument of `use`/`hide` should always be module name
+        if span.contains(*pos) {
+            return Some((Id::Module(module_id), strip_quotes(span, working_set)));
         }
     }
 
@@ -196,12 +206,7 @@ fn try_find_id_in_use(
         })
     };
 
-    let arguments = if call_name != "hide" {
-        call.arguments.get(1..)?
-    } else {
-        call.arguments.as_slice()
-    };
-
+    let arguments = call.arguments.get(1..)?;
     for arg in arguments {
         let Argument::Positional(expr) = arg else {
             continue;
@@ -243,11 +248,25 @@ fn try_find_id_in_overlay(
     id: Option<&Id>,
 ) -> Option<(Id, Span)> {
     let check_location = |span: &Span| location.is_none_or(|pos| span.contains(*pos));
+    let module_from_parser_info = |span: Span| {
+        let Expression {
+            expr: Expr::Overlay(Some(module_id)),
+            ..
+        } = call.get_parser_info("overlay_expr")?
+        else {
+            return None;
+        };
+        let found_id = Id::Module(*module_id);
+        id.is_none_or(|id_r| found_id == *id_r)
+            .then_some((found_id, strip_quotes(span, working_set)))
+    };
+    // NOTE: `overlay_expr` doesn't work for `overlay hide`
     let module_from_overlay_name = |name: &str, span: Span| {
         let found_id = Id::Module(working_set.find_overlay(name.as_bytes())?.origin);
         id.is_none_or(|id_r| found_id == *id_r)
             .then_some((found_id, strip_quotes(span, working_set)))
     };
+
     for arg in call.arguments.iter() {
         let Argument::Positional(expr) = arg else {
             continue;
@@ -256,11 +275,12 @@ fn try_find_id_in_overlay(
             continue;
         };
         let matched = match &expr.expr {
-            Expr::String(name) => get_matched_module_id(working_set, expr.span, id)
-                .or(module_from_overlay_name(name, expr.span)),
+            Expr::String(name) => module_from_parser_info(expr.span)
+                .or_else(|| module_from_overlay_name(name, expr.span)),
             // keyword 'as'
             Expr::Keyword(kwd) => match &kwd.expr.expr {
-                Expr::String(name) => module_from_overlay_name(name, kwd.expr.span),
+                Expr::String(name) => module_from_parser_info(kwd.expr.span)
+                    .or_else(|| module_from_overlay_name(name, kwd.expr.span)),
                 _ => None,
             },
             _ => None,
@@ -272,20 +292,6 @@ fn try_find_id_in_overlay(
     None
 }
 
-fn get_matched_module_id(
-    working_set: &StateWorkingSet,
-    span: Span,
-    id: Option<&Id>,
-) -> Option<(Id, Span)> {
-    let span = strip_quotes(span, working_set);
-    let name = String::from_utf8_lossy(working_set.get_span_contents(span));
-    let path = std::path::PathBuf::from(name.as_ref());
-    let stem = path.file_stem().and_then(|fs| fs.to_str()).unwrap_or(&name);
-    let found_id = Id::Module(working_set.find_module(stem.as_bytes())?);
-    id.is_none_or(|id_r| found_id == *id_r)
-        .then_some((found_id, span))
-}
-
 fn find_id_in_expr(
     expr: &Expression,
     working_set: &StateWorkingSet,

crates/nu-lsp/src/goto.rs

@@ -184,18 +184,18 @@ mod tests {
         let mut script = fixtures();
         script.push("lsp");
         script.push("hover");
-        script.push("cell_path.nu");
+        script.push("use.nu");
         let script = path_to_uri(&script);
 
         open_unchecked(&client_connection, script.clone());
 
-        let resp = send_goto_definition_request(&client_connection, script.clone(), 4, 7);
+        let resp = send_goto_definition_request(&client_connection, script.clone(), 2, 7);
         assert_json_eq!(
             result_from_message(resp).pointer("/range/start").unwrap(),
             serde_json::json!({ "line": 1, "character": 10 })
         );
 
-        let resp = send_goto_definition_request(&client_connection, script.clone(), 4, 9);
+        let resp = send_goto_definition_request(&client_connection, script.clone(), 2, 9);
         assert_json_eq!(
             result_from_message(resp).pointer("/range/start").unwrap(),
             serde_json::json!({ "line": 1, "character": 17 })

crates/nu-lsp/src/hover.rs

@@ -246,26 +246,26 @@ mod hover_tests {
         let mut script = fixtures();
         script.push("lsp");
         script.push("hover");
-        script.push("cell_path.nu");
+        script.push("use.nu");
         let script = path_to_uri(&script);
 
         open_unchecked(&client_connection, script.clone());
 
-        let resp = send_hover_request(&client_connection, script.clone(), 4, 3);
+        let resp = send_hover_request(&client_connection, script.clone(), 2, 3);
         let result = result_from_message(resp);
         assert_json_eq!(
             result.pointer("/contents/value").unwrap(),
             serde_json::json!("```\nlist<any>\n```")
         );
 
-        let resp = send_hover_request(&client_connection, script.clone(), 4, 7);
+        let resp = send_hover_request(&client_connection, script.clone(), 2, 7);
         let result = result_from_message(resp);
         assert_json_eq!(
             result.pointer("/contents/value").unwrap(),
             serde_json::json!("```\nrecord<bar: int>\n```")
         );
 
-        let resp = send_hover_request(&client_connection, script.clone(), 4, 11);
+        let resp = send_hover_request(&client_connection, script.clone(), 2, 11);
         let result = result_from_message(resp);
         assert_json_eq!(
             result.pointer("/contents/value").unwrap(),
@@ -394,4 +394,39 @@ mod hover_tests {
             "\"# module doc\""
         );
     }
+
+    #[test]
+    fn hover_on_use_command() {
+        let mut script = fixtures();
+        script.push("lsp");
+        script.push("hover");
+        script.push("use.nu");
+        let script_uri = path_to_uri(&script);
+
+        let config = format!("use {}", script.to_str().unwrap());
+        let (client_connection, _recv) = initialize_language_server(Some(&config), None);
+
+        open_unchecked(&client_connection, script_uri.clone());
+        let resp = send_hover_request(&client_connection, script_uri.clone(), 0, 19);
+        let result = result_from_message(resp);
+
+        assert_eq!(
+            result
+                .pointer("/contents/value")
+                .unwrap()
+                .to_string()
+                .replace("\\r", ""),
+            "\"```\\nrecord<foo: list<any>>\\n``` \\n---\\nimmutable\""
+        );
+
+        let resp = send_hover_request(&client_connection, script_uri.clone(), 0, 22);
+        let result = result_from_message(resp);
+
+        assert!(result
+            .pointer("/contents/value")
+            .unwrap()
+            .to_string()
+            .replace("\\r", "")
+            .starts_with("\"\\n---\\n### Usage \\n```nu\\n  foo {flags}\\n```\\n\\n### Flags"));
+    }
 }

tests/fixtures/lsp/hover/cell_path.nu

@@ -1,5 +1,5 @@
-const r = {
+export const r = {
   foo: [1 {bar : 2}]
 }
 
-$r.foo.1.bar
+export def foo [] { }

tests/fixtures/lsp/hover/use.nu

@@ -0,0 +1,4 @@
+use cell_path.nu [ r foo ]
+def test [] {
+$r.foo.1.bar
+}