Repository: nushell/nushell (language: rust)

Issue:
The span of `Keyword` does not include its subsidiary expression
### Question

https://github.com/nushell/nushell/blob/23dc1b600a0e09279cef75f0dbf47505c37edeca/crates/nu-protocol/src/ast/keyword.rs#L6-L10

It's kind of weird that the Keyword expression contains at subsidiary expression (like an else block) in it, however the the span of the whole `Keyword` expression only covers the actual keyword syntax string.

https://github.com/nushell/nushell/blob/23dc1b600a0e09279cef75f0dbf47505c37edeca/crates/nu-parser/src/parser.rs#L913-L917

I don't know if this is designed on purpose or not, but it costs some extra attention dealing with those spans. 



Discussion / hints:
Doesn't `expr` have an `expr.span`? You maybe able to use `Span::concat()` to combine them. Or maybe I'm missing your point?
I mean the `Keyword` should be considered as a parent node? If so, a parent node's span should always cover all of its children, right?
I'm sorry, the phrasing should be, for an expression

```rust
Expression {expr: Expr::Keyword(Keyword {keyword: xxx, span: span_a, expr: expr_a} ), span: span_b}
```
where the `span_b` should cover `span_a` as well as the span of `expr_a` imo, however it's not the case,
`span_b` is the same as `span_a` which only covers the keyword string.

I'm referencing the wrong code snippet here, plz ignore that.

Fix patch:
diff --git a/crates/nu-lsp/src/ast.rs b/crates/nu-lsp/src/ast.rs
index 56d67aa5f3096..21c6ddc71d724 100644
--- a/crates/nu-lsp/src/ast.rs
+++ b/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);
             }
diff --git a/crates/nu-lsp/src/goto.rs b/crates/nu-lsp/src/goto.rs
index 15e607477e8af..ef1bcb3101536 100644
--- a/crates/nu-lsp/src/goto.rs
+++ b/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 }
+                }
+            })
+        );
+    }
 }


Test patch:
diff --git a/tests/fixtures/lsp/goto/use_module.nu b/tests/fixtures/lsp/goto/use_module.nu
index 0bbc2e19a5258..058e631fd73e5 100644
--- a/tests/fixtures/lsp/goto/use_module.nu
+++ b/tests/fixtures/lsp/goto/use_module.nu
@@ -1,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"