fix(lsp): goto/hover on module name in some commands (nushell#14924)

# Description
This PR enables basic goto def/hover features on module name in
commands:
1. hide
2. overlay use
3. overlay hide

## Some pending issues

1. Somewhat related to nushell#14816
```nushell
overlay use foo as bar
                  # |_______ cursor here
```
fails to work, since the position of the cursor is outside of the whole
span of this call expression.
I'll try to fix nushell#14816 later instead of implementing new weird
workarounds.

2. references/renaming however is still buggy on overlays with
`as`/`--prefix` features enabled.

# User-Facing Changes
# Tests + Formatting
3 more tests
# After Submitting

crates/nu-lsp/src/ast.rs

@@ -266,8 +266,8 @@ fn try_find_id_in_mod(
     })
 }
 
-/// Find id in use command
-/// `use foo.nu bar` or `use foo.nu [bar baz]`
+/// 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
@@ -282,9 +282,11 @@ fn try_find_id_in_use(
     id: Option<&Id>,
 ) -> Option<(Id, Span)> {
     let call_name = working_set.get_span_contents(call.head);
-    if call_name != b"use" {
+    if call_name != b"use" && call_name != b"hide" {
         return None;
     }
+    // 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)
@@ -303,28 +305,21 @@ fn try_find_id_in_use(
         _ => None,
     };
     let check_location = |span: &Span| location.map_or(true, |pos| span.contains(*pos));
-    let get_module_id = |span: 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.map_or(true, |id_r| found_id == *id_r)
-            .then_some((found_id, span))
-    };
 
     // Get module id if required
     let module_name = call.arguments.first()?;
     let span = module_name.span();
     if let Some(Id::Module(_)) = id {
-        // still need to check the second argument
-        if let Some(res) = get_module_id(span) {
+        // 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 let Some(pos) = location {
-        if span.contains(*pos) {
-            return get_module_id(span);
+        // first argument of `use` should always be module name
+        // while it is optional in `hide`
+        if span.contains(*pos) && call_name == b"use" {
+            return get_matched_module_id(working_set, span, id);
         }
     }
 
@@ -343,26 +338,103 @@ fn try_find_id_in_use(
         })
     };
 
-    // the imported name is always at the second argument
-    let Argument::Positional(expr) = call.arguments.get(1)? else {
-        return None;
+    let arguments = if call_name == b"use" {
+        call.arguments.get(1..)?
+    } else {
+        call.arguments.as_slice()
     };
-    if !check_location(&expr.span) {
+
+    for arg in arguments {
+        let Argument::Positional(expr) = arg else {
+            continue;
+        };
+        if !check_location(&expr.span) {
+            continue;
+        }
+        let matched = match &expr.expr {
+            Expr::String(name) => {
+                find_by_name(name.as_bytes()).map(|id| (id, strip_quotes(expr.span, working_set)))
+            }
+            Expr::List(items) => search_in_list_items(items),
+            Expr::FullCellPath(fcp) => {
+                let Expr::List(items) = &fcp.head.expr else {
+                    return None;
+                };
+                search_in_list_items(items)
+            }
+            _ => None,
+        };
+        if matched.is_some() || location.is_some() {
+            return matched;
+        }
+    }
+    None
+}
+
+/// Find id in use/hide command
+///
+/// TODO: rename of `overlay use as new_name`, `overlay use --prefix`
+///
+/// # Arguments
+/// - `location`: None if no `contains` check required
+/// - `id`: None if no id equal check required
+fn try_find_id_in_overlay(
+    call: &Call,
+    working_set: &StateWorkingSet,
+    location: Option<&usize>,
+    id: Option<&Id>,
+) -> Option<(Id, Span)> {
+    let call_name = working_set.get_span_contents(call.head);
+    if call_name != b"overlay use" && call_name != b"overlay hide" {
         return None;
     }
-    match &expr.expr {
-        Expr::String(name) => {
-            find_by_name(name.as_bytes()).map(|id| (id, strip_quotes(expr.span, working_set)))
-        }
-        Expr::List(items) => search_in_list_items(items),
-        Expr::FullCellPath(fcp) => {
-            let Expr::List(items) = &fcp.head.expr else {
-                return None;
-            };
-            search_in_list_items(items)
+    let check_location = |span: &Span| location.map_or(true, |pos| span.contains(*pos));
+    for arg in call.arguments.iter() {
+        let Argument::Positional(expr) = arg else {
+            continue;
+        };
+        if !check_location(&expr.span) {
+            continue;
+        };
+        let matched = match &expr.expr {
+            Expr::String(name) => {
+                let name = name.as_bytes();
+                get_matched_module_id(working_set, expr.span, id).or_else(|| {
+                    let found_id = Id::Module(working_set.find_overlay(name)?.origin);
+                    id.map_or(true, |id_r| found_id == *id_r)
+                        .then_some((found_id, strip_quotes(expr.span, working_set)))
+                })
+            }
+            // keyword 'as'
+            Expr::Keyword(kwd) => match &kwd.expr.expr {
+                Expr::String(name) => {
+                    let found_id = Id::Module(working_set.find_overlay(name.as_bytes())?.origin);
+                    id.map_or(true, |id_r| found_id == *id_r)
+                        .then_some((found_id, strip_quotes(kwd.expr.span, working_set)))
+                }
+                _ => None,
+            },
+            _ => None,
+        };
+        if matched.is_some() || location.is_some() {
+            return matched;
         }
-        _ => None,
     }
+    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.map_or(true, |id_r| found_id == *id_r)
+        .then_some((found_id, span))
 }
 
 fn find_id_in_expr(
@@ -396,6 +468,12 @@ fn find_id_in_expr(
                 try_find_id_in_def(call, working_set, Some(location), None)
                     .or(try_find_id_in_mod(call, working_set, Some(location), None))
                     .or(try_find_id_in_use(call, working_set, Some(location), None))
+                    .or(try_find_id_in_overlay(
+                        call,
+                        working_set,
+                        Some(location),
+                        None,
+                    ))
                     .map(|p| vec![p])
             }
         }
@@ -458,6 +536,7 @@ fn find_reference_by_id_in_expr(
             if let Some((_, span_found)) = try_find_id_in_def(call, working_set, None, Some(id))
                 .or(try_find_id_in_mod(call, working_set, None, Some(id)))
                 .or(try_find_id_in_use(call, working_set, None, Some(id)))
+                .or(try_find_id_in_overlay(call, working_set, None, Some(id)))
             {
                 occurs.push(span_found);
             }

crates/nu-lsp/src/goto.rs

@@ -55,7 +55,7 @@ impl LanguageServer {
 mod tests {
     use crate::path_to_uri;
     use crate::tests::{initialize_language_server, open_unchecked};
-    use assert_json_diff::assert_json_eq;
+    use assert_json_diff::{assert_json_eq, assert_json_include};
     use lsp_server::{Connection, Message};
     use lsp_types::{
         request::{GotoDefinition, Request},
@@ -413,4 +413,84 @@ mod tests {
             })
         );
     }
+
+    #[test]
+    fn goto_definition_of_module_in_hide() {
+        let (client_connection, _recv) = initialize_language_server(None);
+
+        let mut script = fixtures();
+        script.push("lsp");
+        script.push("goto");
+        script.push("use_module.nu");
+        let script = path_to_uri(&script);
+
+        open_unchecked(&client_connection, script.clone());
+
+        let resp = send_goto_definition_request(&client_connection, script.clone(), 3, 6);
+        let result = if let Message::Response(response) = resp {
+            response.result
+        } else {
+            panic!()
+        };
+
+        assert_json_eq!(
+            result,
+            serde_json::json!({
+                "uri": script.to_string().replace("use_module", "module"),
+                "range": {
+                    "start": { "line": 1, "character": 29 },
+                    "end": { "line": 1, "character": 30 }
+                }
+            })
+        );
+    }
+
+    #[test]
+    fn goto_definition_of_module_in_overlay() {
+        let (client_connection, _recv) = initialize_language_server(None);
+
+        let mut script = fixtures();
+        script.push("lsp");
+        script.push("goto");
+        script.push("use_module.nu");
+        let script = path_to_uri(&script);
+
+        open_unchecked(&client_connection, script.clone());
+
+        let resp = send_goto_definition_request(&client_connection, script.clone(), 1, 20);
+        let result = if let Message::Response(response) = resp {
+            response.result
+        } else {
+            panic!()
+        };
+
+        assert_json_include!(
+            actual: result,
+            expected: serde_json::json!({
+                "uri": script.to_string().replace("use_module", "module"),
+                "range": {
+                    "start": { "line": 0, "character": 0 },
+                    "end": { "line": 3 }
+                }
+            })
+        );
+
+        let resp = send_goto_definition_request(&client_connection, script.clone(), 2, 30);
+        let result = if let Message::Response(response) = resp {
+            response.result
+        } else {
+            panic!()
+        };
+
+        assert_json_include!(
+            actual: result,
+            expected: serde_json::json!({
+                "uri": script.to_string().replace("use_module", "module"),
+                "range": {
+                    "start": { "line": 0, "character": 0 },
+                    "end": { "line": 3 }
+                }
+            })
+        );
+    }
 }

tests/fixtures/lsp/goto/use_module.nu

@@ -1 +1,4 @@
 use module.nu "module name"
+overlay use module.nu as new_name
+overlay hide --keep-env [PWD] new_name
+hide "module name"