fix(lsp): goto definition on variables in match guard (#14818)

<!--
if this PR closes one or more issues, you can automatically link the PR
with
them by using one of the [*linking
keywords*](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword),
e.g.
- this PR should close #xxxx
- fixes #xxxx

you can also mention related issues, PRs or discussions!
-->

# Description
<!--
Thank you for improving Nushell. Please, check our [contributing
guide](../CONTRIBUTING.md) and talk to the core team before making major
changes.

Description of your pull request goes here. **Provide examples and/or
screenshots** if your changes affect the user experience.
-->

This PR fixes a corner case of goto definition in lsp server.

```nushell
let foo = 1
match $foo {
  _ if $foo == 1 => 1
        # |_______________ goto definition does not work here
  _ => 2
}
```

Since `match_pattern.guard` is not handled in this function (which could
be another issue).


23dc1b600a/crates/nu-parser/src/flatten.rs (L604-L658)

In this PR, however, finding leaf expression at the cursor is done with
the new AST traversing helper functions.
Theoretically, this is faster as the flattening and filtering are
combined in a single scan; the difference could be negligible though.
 

# User-Facing Changes
<!-- List of all changes that impact the user experience here. This
helps us keep track of breaking changes. -->

# Tests + Formatting
<!--
Don't forget to add tests that cover your changes.

Make sure you've run and fixed any issues with these commands:

- `cargo fmt --all -- --check` to check standard code formatting (`cargo
fmt --all` applies these changes)
- `cargo clippy --workspace -- -D warnings -D clippy::unwrap_used` to
check that you're using the standard code style
- `cargo test --workspace` to check that all tests pass (on Windows make
sure to [enable developer
mode](https://learn.microsoft.com/en-us/windows/apps/get-started/developer-mode-features-and-debugging))
- `cargo run -- -c "use toolkit.nu; toolkit test stdlib"` to run the
tests for the standard library

> **Note**
> from `nushell` you can also use the `toolkit` as follows
> ```bash
> use toolkit.nu # or use an `env_change` hook to activate it
automatically
> toolkit check pr
> ```
-->

3 new test cases added, will add more if new issues found.

# After Submitting
<!-- If your PR had any user-facing changes, update [the
documentation](https://github.com/nushell/nushell.github.io) after the
PR is merged, if necessary. This will help us keep the docs up to date.
-->

crates/nu-lsp/src/ast.rs

@@ -0,0 +1,210 @@
+use std::sync::Arc;
+
+use nu_protocol::{
+    ast::{
+        Block, Expr, Expression, ExternalArgument, ListItem, MatchPattern, Pattern,
+        PipelineRedirection, RecordItem,
+    },
+    engine::StateWorkingSet,
+};
+
+use crate::Id;
+
+/// similar to flatten_block, but allows extra map function
+pub fn ast_flat_map<T, E>(
+    ast: &Arc<Block>,
+    working_set: &StateWorkingSet,
+    extra_args: &E,
+    f_special: fn(&Expression, &StateWorkingSet, &E) -> Option<Vec<T>>,
+) -> Vec<T> {
+    ast.pipelines
+        .iter()
+        .flat_map(|pipeline| {
+            pipeline.elements.iter().flat_map(|element| {
+                expr_flat_map(&element.expr, working_set, extra_args, f_special)
+                    .into_iter()
+                    .chain(
+                        element
+                            .redirection
+                            .as_ref()
+                            .map(|redir| {
+                                redirect_flat_map(redir, working_set, extra_args, f_special)
+                            })
+                            .unwrap_or_default(),
+                    )
+            })
+        })
+        .collect()
+}
+
+/// generic function that do flat_map on an expression
+/// concats all recursive results on sub-expressions
+///
+/// # Arguments
+/// * `f_special` - function that overrides the default behavior
+pub fn expr_flat_map<T, E>(
+    expr: &Expression,
+    working_set: &StateWorkingSet,
+    extra_args: &E,
+    f_special: fn(&Expression, &StateWorkingSet, &E) -> Option<Vec<T>>,
+) -> Vec<T> {
+    // behavior overridden by f_special
+    if let Some(vec) = f_special(expr, working_set, extra_args) {
+        return vec;
+    }
+    let recur = |expr| expr_flat_map(expr, working_set, extra_args, f_special);
+    match &expr.expr {
+        Expr::RowCondition(block_id)
+        | Expr::Subexpression(block_id)
+        | Expr::Block(block_id)
+        | Expr::Closure(block_id) => {
+            let block = working_set.get_block(block_id.to_owned());
+            ast_flat_map(block, working_set, extra_args, f_special)
+        }
+        Expr::Range(range) => [&range.from, &range.next, &range.to]
+            .iter()
+            .filter_map(|e| e.as_ref())
+            .flat_map(recur)
+            .collect(),
+        Expr::Call(call) => call
+            .arguments
+            .iter()
+            .filter_map(|arg| arg.expr())
+            .flat_map(recur)
+            .collect(),
+        Expr::ExternalCall(head, args) => recur(head)
+            .into_iter()
+            .chain(args.iter().flat_map(|arg| match arg {
+                ExternalArgument::Regular(e) | ExternalArgument::Spread(e) => recur(e),
+            }))
+            .collect(),
+        Expr::UnaryNot(expr) | Expr::Collect(_, expr) => recur(expr),
+        Expr::BinaryOp(lhs, op, rhs) => recur(lhs)
+            .into_iter()
+            .chain(recur(op))
+            .chain(recur(rhs))
+            .collect(),
+        Expr::MatchBlock(matches) => matches
+            .iter()
+            .flat_map(|(pattern, expr)| {
+                match_pattern_flat_map(pattern, working_set, extra_args, f_special)
+                    .into_iter()
+                    .chain(recur(expr))
+            })
+            .collect(),
+        Expr::List(items) => items
+            .iter()
+            .flat_map(|item| match item {
+                ListItem::Item(expr) | ListItem::Spread(_, expr) => recur(expr),
+            })
+            .collect(),
+        Expr::Record(items) => items
+            .iter()
+            .flat_map(|item| match item {
+                RecordItem::Spread(_, expr) => recur(expr),
+                RecordItem::Pair(key, val) => [key, val].into_iter().flat_map(recur).collect(),
+            })
+            .collect(),
+        Expr::Table(table) => table
+            .columns
+            .iter()
+            .flat_map(recur)
+            .chain(table.rows.iter().flat_map(|row| row.iter().flat_map(recur)))
+            .collect(),
+        Expr::ValueWithUnit(vu) => recur(&vu.expr),
+        Expr::FullCellPath(fcp) => recur(&fcp.head),
+        Expr::Keyword(kw) => recur(&kw.expr),
+        Expr::StringInterpolation(vec) | Expr::GlobInterpolation(vec, _) => {
+            vec.iter().flat_map(recur).collect()
+        }
+
+        _ => Vec::new(),
+    }
+}
+
+/// flat_map on match patterns
+fn match_pattern_flat_map<T, E>(
+    pattern: &MatchPattern,
+    working_set: &StateWorkingSet,
+    extra_args: &E,
+    f_special: fn(&Expression, &StateWorkingSet, &E) -> Option<Vec<T>>,
+) -> Vec<T> {
+    let recur = |expr| expr_flat_map(expr, working_set, extra_args, f_special);
+    let recur_match = |p| match_pattern_flat_map(p, working_set, extra_args, f_special);
+    match &pattern.pattern {
+        Pattern::Expression(expr) => recur(expr),
+        Pattern::List(patterns) | Pattern::Or(patterns) => {
+            patterns.iter().flat_map(recur_match).collect()
+        }
+        Pattern::Record(entries) => entries.iter().flat_map(|(_, p)| recur_match(p)).collect(),
+        _ => Vec::new(),
+    }
+    .into_iter()
+    .chain(pattern.guard.as_ref().map(|g| recur(g)).unwrap_or_default())
+    .collect()
+}
+
+/// flat_map on redirections
+fn redirect_flat_map<T, E>(
+    redir: &PipelineRedirection,
+    working_set: &StateWorkingSet,
+    extra_args: &E,
+    f_special: fn(&Expression, &StateWorkingSet, &E) -> Option<Vec<T>>,
+) -> Vec<T> {
+    let recur = |expr| expr_flat_map(expr, working_set, extra_args, f_special);
+    match redir {
+        PipelineRedirection::Single { target, .. } => target.expr().map(recur).unwrap_or_default(),
+        PipelineRedirection::Separate { out, err } => [out, err]
+            .iter()
+            .filter_map(|t| t.expr())
+            .flat_map(recur)
+            .collect(),
+    }
+}
+
+fn find_id_in_expr(expr: &Expression, _: &StateWorkingSet, location: &usize) -> Option<Vec<Id>> {
+    // skip the entire expression if the location is not in it
+    if !expr.span.contains(*location) {
+        // TODO: the span of Keyword does not include its subsidiary expression
+        // resort to `expr_flat_map` if location found in its expr
+        if let Expr::Keyword(kw) = &expr.expr {
+            if kw.expr.span.contains(*location) {
+                return None;
+            }
+        }
+        return Some(Vec::new());
+    }
+    match &expr.expr {
+        Expr::Var(var_id) | Expr::VarDecl(var_id) => Some(vec![Id::Variable(*var_id)]),
+        Expr::Call(call) => {
+            if call.head.contains(*location) {
+                Some(vec![Id::Declaration(call.decl_id)])
+            } else {
+                None
+            }
+        }
+        Expr::Overlay(Some(module_id)) => Some(vec![Id::Module(*module_id)]),
+        // terminal value expressions
+        Expr::Bool(_)
+        | Expr::Binary(_)
+        | Expr::DateTime(_)
+        | Expr::Directory(_, _)
+        | Expr::Filepath(_, _)
+        | Expr::Float(_)
+        | Expr::Garbage
+        | Expr::GlobPattern(_, _)
+        | Expr::Int(_)
+        | Expr::Nothing
+        | Expr::RawString(_)
+        | Expr::Signature(_)
+        | Expr::String(_) => Some(vec![Id::Value(expr.ty.clone())]),
+        _ => None,
+    }
+}
+
+/// find the leaf node at the given location from ast
+pub fn find_id(ast: &Arc<Block>, working_set: &StateWorkingSet, location: &usize) -> Option<Id> {
+    ast_flat_map(ast, working_set, location, find_id_in_expr)
+        .first()
+        .cloned()
+}

crates/nu-lsp/src/goto.rs

@@ -0,0 +1,344 @@
+use crate::ast::find_id;
+use crate::{Id, LanguageServer};
+use lsp_types::{GotoDefinitionParams, GotoDefinitionResponse};
+
+impl LanguageServer {
+    pub fn goto_definition(
+        &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 (block, file_offset, working_set, file) =
+            self.parse_file(&mut engine_state, &path_uri, false)?;
+        let location =
+            file.offset_at(params.text_document_position_params.position) as usize + file_offset;
+        let id = find_id(&block, &working_set, &location)?;
+
+        let span = match id {
+            Id::Declaration(decl_id) => {
+                let block_id = working_set.get_decl(decl_id).block_id()?;
+                working_set.get_block(block_id).span
+            }
+            Id::Variable(var_id) => {
+                let var = working_set.get_variable(var_id);
+                Some(var.declaration_span)
+            }
+            Id::Module(module_id) => {
+                let module = working_set.get_module(module_id);
+                module.span
+            }
+            _ => None,
+        }?;
+        Some(GotoDefinitionResponse::Scalar(
+            self.get_location_by_span(working_set.files(), &span)?,
+        ))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::path_to_uri;
+    use crate::tests::{initialize_language_server, open_unchecked};
+    use assert_json_diff::assert_json_eq;
+    use lsp_server::{Connection, Message};
+    use lsp_types::request::{GotoDefinition, Request};
+    use lsp_types::{
+        GotoDefinitionParams, PartialResultParams, Position, TextDocumentIdentifier,
+        TextDocumentPositionParams, Uri, WorkDoneProgressParams,
+    };
+    use nu_test_support::fs::{fixtures, root};
+
+    fn send_goto_definition_request(
+        client_connection: &Connection,
+        uri: Uri,
+        line: u32,
+        character: u32,
+    ) -> Message {
+        client_connection
+            .sender
+            .send(Message::Request(lsp_server::Request {
+                id: 2.into(),
+                method: GotoDefinition::METHOD.to_string(),
+                params: serde_json::to_value(GotoDefinitionParams {
+                    text_document_position_params: TextDocumentPositionParams {
+                        text_document: TextDocumentIdentifier { uri },
+                        position: Position { line, character },
+                    },
+                    work_done_progress_params: WorkDoneProgressParams::default(),
+                    partial_result_params: PartialResultParams::default(),
+                })
+                .unwrap(),
+            }))
+            .unwrap();
+
+        client_connection
+            .receiver
+            .recv_timeout(std::time::Duration::from_secs(2))
+            .unwrap()
+    }
+
+    #[test]
+    fn goto_definition_for_none_existing_file() {
+        let (client_connection, _recv) = initialize_language_server();
+
+        let mut none_existent_path = root();
+        none_existent_path.push("none-existent.nu");
+
+        client_connection
+            .sender
+            .send(Message::Request(lsp_server::Request {
+                id: 2.into(),
+                method: GotoDefinition::METHOD.to_string(),
+                params: serde_json::to_value(GotoDefinitionParams {
+                    text_document_position_params: TextDocumentPositionParams {
+                        text_document: TextDocumentIdentifier {
+                            uri: path_to_uri(&none_existent_path),
+                        },
+                        position: Position {
+                            line: 0,
+                            character: 0,
+                        },
+                    },
+                    work_done_progress_params: WorkDoneProgressParams::default(),
+                    partial_result_params: PartialResultParams::default(),
+                })
+                .unwrap(),
+            }))
+            .unwrap();
+
+        let resp = client_connection
+            .receiver
+            .recv_timeout(std::time::Duration::from_secs(2))
+            .unwrap();
+        let result = if let Message::Response(response) = resp {
+            response.result
+        } else {
+            panic!()
+        };
+
+        assert_json_eq!(result, serde_json::json!(null));
+    }
+
+    #[test]
+    fn goto_definition_of_variable() {
+        let (client_connection, _recv) = initialize_language_server();
+
+        let mut script = fixtures();
+        script.push("lsp");
+        script.push("goto");
+        script.push("var.nu");
+        let script = path_to_uri(&script);
+
+        open_unchecked(&client_connection, script.clone());
+
+        let resp = send_goto_definition_request(&client_connection, script.clone(), 2, 12);
+        let result = if let Message::Response(response) = resp {
+            response.result
+        } else {
+            panic!()
+        };
+
+        assert_json_eq!(
+            result,
+            serde_json::json!({
+                "uri": script,
+                "range": {
+                    "start": { "line": 0, "character": 4 },
+                    "end": { "line": 0, "character": 12 }
+                }
+            })
+        );
+    }
+
+    #[test]
+    fn goto_definition_of_command() {
+        let (client_connection, _recv) = initialize_language_server();
+
+        let mut script = fixtures();
+        script.push("lsp");
+        script.push("goto");
+        script.push("command.nu");
+        let script = path_to_uri(&script);
+
+        open_unchecked(&client_connection, script.clone());
+
+        let resp = send_goto_definition_request(&client_connection, script.clone(), 4, 1);
+        let result = if let Message::Response(response) = resp {
+            response.result
+        } else {
+            panic!()
+        };
+
+        assert_json_eq!(
+            result,
+            serde_json::json!({
+                "uri": script,
+                "range": {
+                "start": { "line": 0, "character": 17 },
+                "end": { "line": 2, "character": 1 }
+            }
+            })
+        );
+    }
+
+    #[test]
+    fn goto_definition_of_command_unicode() {
+        let (client_connection, _recv) = initialize_language_server();
+
+        let mut script = fixtures();
+        script.push("lsp");
+        script.push("goto");
+        script.push("command_unicode.nu");
+        let script = path_to_uri(&script);
+
+        open_unchecked(&client_connection, script.clone());
+
+        let resp = send_goto_definition_request(&client_connection, script.clone(), 4, 2);
+        let result = if let Message::Response(response) = resp {
+            response.result
+        } else {
+            panic!()
+        };
+
+        assert_json_eq!(
+            result,
+            serde_json::json!({
+                "uri": script,
+                "range": {
+                "start": { "line": 0, "character": 19 },
+                "end": { "line": 2, "character": 1 }
+            }
+            })
+        );
+    }
+
+    #[test]
+    fn goto_definition_of_command_parameter() {
+        let (client_connection, _recv) = initialize_language_server();
+
+        let mut script = fixtures();
+        script.push("lsp");
+        script.push("goto");
+        script.push("command.nu");
+        let script = path_to_uri(&script);
+
+        open_unchecked(&client_connection, script.clone());
+
+        let resp = send_goto_definition_request(&client_connection, script.clone(), 1, 14);
+        let result = if let Message::Response(response) = resp {
+            response.result
+        } else {
+            panic!()
+        };
+
+        assert_json_eq!(
+            result,
+            serde_json::json!({
+                "uri": script,
+                "range": {
+                "start": { "line": 0, "character": 11 },
+                "end": { "line": 0, "character": 15 }
+            }
+            })
+        );
+    }
+
+    #[test]
+    fn goto_definition_of_variable_in_else_block() {
+        let (client_connection, _recv) = initialize_language_server();
+
+        let mut script = fixtures();
+        script.push("lsp");
+        script.push("goto");
+        script.push("else.nu");
+        let script = path_to_uri(&script);
+
+        open_unchecked(&client_connection, script.clone());
+
+        let resp = send_goto_definition_request(&client_connection, script.clone(), 1, 21);
+        let result = if let Message::Response(response) = resp {
+            response.result
+        } else {
+            panic!()
+        };
+
+        assert_json_eq!(
+            result,
+            serde_json::json!({
+                "uri": script,
+                "range": {
+                    "start": { "line": 0, "character": 4 },
+                    "end": { "line": 0, "character": 7 }
+                }
+            })
+        );
+    }
+
+    #[test]
+    fn goto_definition_of_variable_in_match_guard() {
+        let (client_connection, _recv) = initialize_language_server();
+
+        let mut script = fixtures();
+        script.push("lsp");
+        script.push("goto");
+        script.push("match.nu");
+        let script = path_to_uri(&script);
+
+        open_unchecked(&client_connection, script.clone());
+
+        let resp = send_goto_definition_request(&client_connection, script.clone(), 2, 9);
+        let result = if let Message::Response(response) = resp {
+            response.result
+        } else {
+            panic!()
+        };
+
+        assert_json_eq!(
+            result,
+            serde_json::json!({
+                "uri": script,
+                "range": {
+                    "start": { "line": 0, "character": 4 },
+                    "end": { "line": 0, "character": 7 }
+                }
+            })
+        );
+    }
+
+    #[test]
+    fn goto_definition_of_variable_in_each() {
+        let (client_connection, _recv) = initialize_language_server();
+
+        let mut script = fixtures();
+        script.push("lsp");
+        script.push("goto");
+        script.push("collect.nu");
+        let script = path_to_uri(&script);
+
+        open_unchecked(&client_connection, script.clone());
+
+        let resp = send_goto_definition_request(&client_connection, script.clone(), 1, 16);
+        let result = if let Message::Response(response) = resp {
+            response.result
+        } else {
+            panic!()
+        };
+
+        assert_json_eq!(
+            result,
+            serde_json::json!({
+                "uri": script,
+                "range": {
+                    "start": { "line": 0, "character": 4 },
+                    "end": { "line": 0, "character": 7 }
+                }
+            })
+        );
+    }
+}

crates/nu-lsp/src/hints.rs

@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use crate::ast::{ast_flat_map, expr_flat_map};
 use crate::{span_to_range, LanguageServer};
 use lsp_textdocument::FullTextDocument;
 use lsp_types::{
@@ -7,166 +8,11 @@ use lsp_types::{
     MarkupKind, Position, Range,
 };
 use nu_protocol::{
-    ast::{
-        Argument, Block, Expr, Expression, ExternalArgument, ListItem, MatchPattern, Operator,
-        Pattern, PipelineRedirection, RecordItem,
-    },
+    ast::{Argument, Block, Expr, Expression, Operator},
     engine::StateWorkingSet,
     Type,
 };
 
-/// similar to flatten_block, but allows extra map function
-fn ast_flat_map<T, E>(
-    ast: &Arc<Block>,
-    working_set: &StateWorkingSet,
-    extra_args: &E,
-    f_special: fn(&Expression, &StateWorkingSet, &E) -> Option<Vec<T>>,
-) -> Vec<T> {
-    ast.pipelines
-        .iter()
-        .flat_map(|pipeline| {
-            pipeline.elements.iter().flat_map(|element| {
-                expr_flat_map(&element.expr, working_set, extra_args, f_special)
-                    .into_iter()
-                    .chain(
-                        element
-                            .redirection
-                            .as_ref()
-                            .map(|redir| {
-                                redirect_flat_map(redir, working_set, extra_args, f_special)
-                            })
-                            .unwrap_or_default(),
-                    )
-            })
-        })
-        .collect()
-}
-
-/// generic function that do flat_map on an expression
-/// concats all recursive results on sub-expressions
-///
-/// # Arguments
-/// * `f_special` - function that overrides the default behavior
-fn expr_flat_map<T, E>(
-    expr: &Expression,
-    working_set: &StateWorkingSet,
-    extra_args: &E,
-    f_special: fn(&Expression, &StateWorkingSet, &E) -> Option<Vec<T>>,
-) -> Vec<T> {
-    // behavior overridden by f_special
-    if let Some(vec) = f_special(expr, working_set, extra_args) {
-        return vec;
-    }
-    let recur = |expr| expr_flat_map(expr, working_set, extra_args, f_special);
-    match &expr.expr {
-        Expr::RowCondition(block_id)
-        | Expr::Subexpression(block_id)
-        | Expr::Block(block_id)
-        | Expr::Closure(block_id) => {
-            let block = working_set.get_block(block_id.to_owned());
-            ast_flat_map(block, working_set, extra_args, f_special)
-        }
-        Expr::Range(range) => [&range.from, &range.next, &range.to]
-            .iter()
-            .filter_map(|e| e.as_ref())
-            .flat_map(recur)
-            .collect(),
-        Expr::Call(call) => call
-            .arguments
-            .iter()
-            .filter_map(|arg| arg.expr())
-            .flat_map(recur)
-            .collect(),
-        Expr::ExternalCall(head, args) => recur(head)
-            .into_iter()
-            .chain(args.iter().flat_map(|arg| match arg {
-                ExternalArgument::Regular(e) | ExternalArgument::Spread(e) => recur(e),
-            }))
-            .collect(),
-        Expr::UnaryNot(expr) | Expr::Collect(_, expr) => recur(expr),
-        Expr::BinaryOp(lhs, op, rhs) => recur(lhs)
-            .into_iter()
-            .chain(recur(op))
-            .chain(recur(rhs))
-            .collect(),
-        Expr::MatchBlock(matches) => matches
-            .iter()
-            .flat_map(|(pattern, expr)| {
-                match_pattern_flat_map(pattern, working_set, extra_args, f_special)
-                    .into_iter()
-                    .chain(recur(expr))
-            })
-            .collect(),
-        Expr::List(items) => items
-            .iter()
-            .flat_map(|item| match item {
-                ListItem::Item(expr) | ListItem::Spread(_, expr) => recur(expr),
-            })
-            .collect(),
-        Expr::Record(items) => items
-            .iter()
-            .flat_map(|item| match item {
-                RecordItem::Spread(_, expr) => recur(expr),
-                RecordItem::Pair(key, val) => [key, val].into_iter().flat_map(recur).collect(),
-            })
-            .collect(),
-        Expr::Table(table) => table
-            .columns
-            .iter()
-            .flat_map(recur)
-            .chain(table.rows.iter().flat_map(|row| row.iter().flat_map(recur)))
-            .collect(),
-        Expr::ValueWithUnit(vu) => recur(&vu.expr),
-        Expr::FullCellPath(fcp) => recur(&fcp.head),
-        Expr::Keyword(kw) => recur(&kw.expr),
-        Expr::StringInterpolation(vec) | Expr::GlobInterpolation(vec, _) => {
-            vec.iter().flat_map(recur).collect()
-        }
-
-        _ => Vec::new(),
-    }
-}
-
-/// flat_map on match patterns
-fn match_pattern_flat_map<T, E>(
-    pattern: &MatchPattern,
-    working_set: &StateWorkingSet,
-    extra_args: &E,
-    f_special: fn(&Expression, &StateWorkingSet, &E) -> Option<Vec<T>>,
-) -> Vec<T> {
-    let recur = |expr| expr_flat_map(expr, working_set, extra_args, f_special);
-    let recur_match = |p| match_pattern_flat_map(p, working_set, extra_args, f_special);
-    match &pattern.pattern {
-        Pattern::Expression(expr) => recur(expr),
-        Pattern::List(patterns) | Pattern::Or(patterns) => {
-            patterns.iter().flat_map(recur_match).collect()
-        }
-        Pattern::Record(entries) => entries.iter().flat_map(|(_, p)| recur_match(p)).collect(),
-        _ => Vec::new(),
-    }
-    .into_iter()
-    .chain(pattern.guard.as_ref().map(|g| recur(g)).unwrap_or_default())
-    .collect()
-}
-
-/// flat_map on redirections
-fn redirect_flat_map<T, E>(
-    redir: &PipelineRedirection,
-    working_set: &StateWorkingSet,
-    extra_args: &E,
-    f_special: fn(&Expression, &StateWorkingSet, &E) -> Option<Vec<T>>,
-) -> Vec<T> {
-    let recur = |expr| expr_flat_map(expr, working_set, extra_args, f_special);
-    match redir {
-        PipelineRedirection::Single { target, .. } => target.expr().map(recur).unwrap_or_default(),
-        PipelineRedirection::Separate { out, err } => [out, err]
-            .iter()
-            .filter_map(|t| t.expr())
-            .flat_map(recur)
-            .collect(),
-    }
-}
-
 fn type_short_name(t: &Type) -> String {
     match t {
         Type::Custom(_) => String::from("custom"),
@@ -182,7 +28,6 @@ fn extract_inlay_hints_from_expression(
     working_set: &StateWorkingSet,
     extra_args: &(usize, &FullTextDocument),
 ) -> Option<Vec<InlayHint>> {
-    let span = expr.span;
     let (offset, file) = extra_args;
     let recur = |expr| {
         expr_flat_map(
@@ -219,7 +64,7 @@ fn extract_inlay_hints_from_expression(
             Some(hints)
         }
         Expr::VarDecl(var_id) => {
-            let position = span_to_range(&span, file, *offset).end;
+            let position = span_to_range(&expr.span, file, *offset).end;
             // skip if the type is already specified in code
             if file
                 .get_content(Some(Range {

crates/nu-lsp/src/lib.rs

@@ -1,23 +1,20 @@
 #![doc = include_str!("../README.md")]
+use ast::find_id;
 use lsp_server::{Connection, IoThreads, Message, Response, ResponseError};
 use lsp_textdocument::{FullTextDocument, TextDocuments};
 use lsp_types::{
-    request::{
-        Completion, DocumentSymbolRequest, GotoDefinition, HoverRequest, InlayHintRequest, Request,
-        WorkspaceSymbolRequest,
-    },
-    CompletionItem, CompletionItemKind, CompletionParams, CompletionResponse, CompletionTextEdit,
-    GotoDefinitionParams, GotoDefinitionResponse, Hover, HoverContents, HoverParams, InlayHint,
-    Location, MarkupContent, MarkupKind, OneOf, Range, ServerCapabilities, TextDocumentSyncKind,
-    TextEdit, Uri,
+    request, request::Request, CompletionItem, CompletionItemKind, CompletionParams,
+    CompletionResponse, CompletionTextEdit, Hover, HoverContents, HoverParams, InlayHint, Location,
+    MarkupContent, MarkupKind, OneOf, Range, RenameOptions, ServerCapabilities,
+    TextDocumentSyncKind, TextEdit, Uri, WorkDoneProgressOptions,
 };
 use miette::{IntoDiagnostic, Result};
 use nu_cli::{NuCompleter, SuggestionKind};
-use nu_parser::{flatten_block, parse, FlatShape};
+use nu_parser::parse;
 use nu_protocol::{
     ast::Block,
     engine::{CachedFile, EngineState, Stack, StateWorkingSet},
-    DeclId, ModuleId, Span, Value, VarId,
+    DeclId, ModuleId, Span, Type, Value, VarId,
 };
 use std::collections::BTreeMap;
 use std::{
@@ -29,16 +26,18 @@ use std::{
 use symbols::SymbolCache;
 use url::Url;
 
+mod ast;
 mod diagnostics;
+mod goto;
 mod hints;
 mod notification;
 mod symbols;
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 enum Id {
     Variable(VarId),
     Declaration(DeclId),
-    Value(FlatShape),
+    Value(Type),
     Module(ModuleId),
 }
 
@@ -105,10 +104,16 @@ impl LanguageServer {
             document_symbol_provider: Some(OneOf::Left(true)),
             workspace_symbol_provider: Some(OneOf::Left(true)),
             inlay_hint_provider: Some(OneOf::Left(true)),
+            rename_provider: Some(OneOf::Right(RenameOptions {
+                prepare_provider: Some(true),
+                work_done_progress_options: WorkDoneProgressOptions {
+                    work_done_progress: Some(true),
+                },
+            })),
+            references_provider: Some(OneOf::Left(true)),
             ..Default::default()
         })
         .expect("Must be serializable");
-
         let _ = self
             .connection
             .initialize_while(server_capabilities, || {
@@ -140,24 +145,24 @@ impl LanguageServer {
                     }
 
                     let resp = match request.method.as_str() {
-                        GotoDefinition::METHOD => {
+                        request::GotoDefinition::METHOD => {
                             Self::handle_lsp_request(request, |params| self.goto_definition(params))
                         }
-                        HoverRequest::METHOD => {
+                        request::HoverRequest::METHOD => {
                             Self::handle_lsp_request(request, |params| self.hover(params))
                         }
-                        Completion::METHOD => {
+                        request::Completion::METHOD => {
                             Self::handle_lsp_request(request, |params| self.complete(params))
                         }
-                        DocumentSymbolRequest::METHOD => {
+                        request::DocumentSymbolRequest::METHOD => {
                             Self::handle_lsp_request(request, |params| self.document_symbol(params))
                         }
-                        WorkspaceSymbolRequest::METHOD => {
+                        request::WorkspaceSymbolRequest::METHOD => {
                             Self::handle_lsp_request(request, |params| {
                                 self.workspace_symbol(params)
                             })
                         }
-                        InlayHintRequest::METHOD => {
+                        request::InlayHintRequest::METHOD => {
                             Self::handle_lsp_request(request, |params| self.get_inlay_hints(params))
                         }
                         _ => {
@@ -281,66 +286,6 @@ impl LanguageServer {
         }
     }
 
-    fn find_id(
-        flattened: Vec<(Span, FlatShape)>,
-        location: usize,
-        offset: usize,
-    ) -> Option<(Id, usize, Span)> {
-        let location = location + offset;
-
-        for (span, shape) in flattened {
-            if location >= span.start && location < span.end {
-                match &shape {
-                    FlatShape::Variable(var_id) | FlatShape::VarDecl(var_id) => {
-                        return Some((Id::Variable(*var_id), offset, span));
-                    }
-                    FlatShape::InternalCall(decl_id) | FlatShape::Custom(decl_id) => {
-                        return Some((Id::Declaration(*decl_id), offset, span));
-                    }
-                    _ => return Some((Id::Value(shape), offset, span)),
-                }
-            }
-        }
-        None
-    }
-
-    fn goto_definition(&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 (block, file_offset, working_set, file) =
-            self.parse_file(&mut engine_state, &path_uri, false)?;
-        let flattened = flatten_block(&working_set, &block);
-        let (id, _, _) = Self::find_id(
-            flattened,
-            file.offset_at(params.text_document_position_params.position) as usize,
-            file_offset,
-        )?;
-
-        let span = match id {
-            Id::Declaration(decl_id) => {
-                let block_id = working_set.get_decl(decl_id).block_id()?;
-                working_set.get_block(block_id).span
-            }
-            Id::Variable(var_id) => {
-                let var = working_set.get_variable(var_id);
-                Some(var.declaration_span)
-            }
-            Id::Module(module_id) => {
-                let module = working_set.get_module(module_id);
-                module.span
-            }
-            _ => None,
-        }?;
-        Some(GotoDefinitionResponse::Scalar(
-            self.get_location_by_span(working_set.files(), &span)?,
-        ))
-    }
-
     fn hover(&mut self, params: &HoverParams) -> Option<Hover> {
         let mut engine_state = self.new_engine_state();
 
@@ -351,12 +296,9 @@ impl LanguageServer {
             .to_owned();
         let (block, file_offset, working_set, file) =
             self.parse_file(&mut engine_state, &path_uri, false)?;
-        let flattened = flatten_block(&working_set, &block);
-        let (id, _, _) = Self::find_id(
-            flattened,
-            file.offset_at(params.text_document_position_params.position) as usize,
-            file_offset,
-        )?;
+        let location =
+            file.offset_at(params.text_document_position_params.position) as usize + file_offset;
+        let id = find_id(&block, &working_set, &location)?;
 
         match id {
             Id::Variable(var_id) => {
@@ -518,37 +460,9 @@ impl LanguageServer {
                     range: None,
                 })
             }
-            Id::Value(shape) => {
-                let hover = String::from(match shape {
-                    FlatShape::Binary => "binary",
-                    FlatShape::Block => "block",
-                    FlatShape::Bool => "bool",
-                    FlatShape::Closure => "closure",
-                    FlatShape::DateTime => "datetime",
-                    FlatShape::Directory => "directory",
-                    FlatShape::External => "external",
-                    FlatShape::ExternalArg => "external arg",
-                    FlatShape::Filepath => "file path",
-                    FlatShape::Flag => "flag",
-                    FlatShape::Float => "float",
-                    FlatShape::GlobPattern => "glob pattern",
-                    FlatShape::Int => "int",
-                    FlatShape::Keyword => "keyword",
-                    FlatShape::List => "list",
-                    FlatShape::MatchPattern => "match-pattern",
-                    FlatShape::Nothing => "nothing",
-                    FlatShape::Range => "range",
-                    FlatShape::Record => "record",
-                    FlatShape::String => "string",
-                    FlatShape::StringInterpolation => "string interpolation",
-                    FlatShape::Table => "table",
-                    _ => {
-                        return None;
-                    }
-                });
-
+            Id::Value(t) => {
                 Some(Hover {
-                    contents: HoverContents::Scalar(lsp_types::MarkedString::String(hover)),
+                    contents: HoverContents::Scalar(lsp_types::MarkedString::String(t.to_string())),
                     // TODO
                     range: None,
                 })
@@ -620,13 +534,13 @@ mod tests {
         notification::{
             DidChangeTextDocument, DidOpenTextDocument, Exit, Initialized, Notification,
         },
-        request::{Completion, GotoDefinition, HoverRequest, Initialize, Request, Shutdown},
-        CompletionParams, DidChangeTextDocumentParams, DidOpenTextDocumentParams,
-        GotoDefinitionParams, InitializeParams, InitializedParams, PartialResultParams, Position,
-        TextDocumentContentChangeEvent, TextDocumentIdentifier, TextDocumentItem,
-        TextDocumentPositionParams, WorkDoneProgressParams,
+        request::{Completion, HoverRequest, Initialize, Request, Shutdown},
+        CompletionParams, DidChangeTextDocumentParams, DidOpenTextDocumentParams, InitializeParams,
+        InitializedParams, PartialResultParams, Position, TextDocumentContentChangeEvent,
+        TextDocumentIdentifier, TextDocumentItem, TextDocumentPositionParams,
+        WorkDoneProgressParams,
     };
-    use nu_test_support::fs::{fixtures, root};
+    use nu_test_support::fs::fixtures;
     use std::sync::mpsc::Receiver;
 
     pub fn initialize_language_server() -> (Connection, Receiver<Result<()>>) {
@@ -696,48 +610,6 @@ mod tests {
             .is_ok());
     }
 
-    #[test]
-    fn goto_definition_for_none_existing_file() {
-        let (client_connection, _recv) = initialize_language_server();
-
-        let mut none_existent_path = root();
-        none_existent_path.push("none-existent.nu");
-
-        client_connection
-            .sender
-            .send(Message::Request(lsp_server::Request {
-                id: 2.into(),
-                method: GotoDefinition::METHOD.to_string(),
-                params: serde_json::to_value(GotoDefinitionParams {
-                    text_document_position_params: TextDocumentPositionParams {
-                        text_document: TextDocumentIdentifier {
-                            uri: path_to_uri(&none_existent_path),
-                        },
-                        position: Position {
-                            line: 0,
-                            character: 0,
-                        },
-                    },
-                    work_done_progress_params: WorkDoneProgressParams::default(),
-                    partial_result_params: PartialResultParams::default(),
-                })
-                .unwrap(),
-            }))
-            .unwrap();
-
-        let resp = client_connection
-            .receiver
-            .recv_timeout(std::time::Duration::from_secs(2))
-            .unwrap();
-        let result = if let Message::Response(response) = resp {
-            response.result
-        } else {
-            panic!()
-        };
-
-        assert_json_eq!(result, serde_json::json!(null));
-    }
-
     pub fn open_unchecked(client_connection: &Connection, uri: Uri) -> lsp_server::Notification {
         open(client_connection, uri).unwrap()
     }
@@ -815,159 +687,6 @@ mod tests {
         }
     }
 
-    fn send_goto_definition_request(
-        client_connection: &Connection,
-        uri: Uri,
-        line: u32,
-        character: u32,
-    ) -> Message {
-        client_connection
-            .sender
-            .send(Message::Request(lsp_server::Request {
-                id: 2.into(),
-                method: GotoDefinition::METHOD.to_string(),
-                params: serde_json::to_value(GotoDefinitionParams {
-                    text_document_position_params: TextDocumentPositionParams {
-                        text_document: TextDocumentIdentifier { uri },
-                        position: Position { line, character },
-                    },
-                    work_done_progress_params: WorkDoneProgressParams::default(),
-                    partial_result_params: PartialResultParams::default(),
-                })
-                .unwrap(),
-            }))
-            .unwrap();
-
-        client_connection
-            .receiver
-            .recv_timeout(std::time::Duration::from_secs(2))
-            .unwrap()
-    }
-
-    #[test]
-    fn goto_definition_of_variable() {
-        let (client_connection, _recv) = initialize_language_server();
-
-        let mut script = fixtures();
-        script.push("lsp");
-        script.push("goto");
-        script.push("var.nu");
-        let script = path_to_uri(&script);
-
-        open_unchecked(&client_connection, script.clone());
-
-        let resp = send_goto_definition_request(&client_connection, script.clone(), 2, 12);
-        let result = if let Message::Response(response) = resp {
-            response.result
-        } else {
-            panic!()
-        };
-
-        assert_json_eq!(
-            result,
-            serde_json::json!({
-                "uri": script,
-                "range": {
-                    "start": { "line": 0, "character": 4 },
-                    "end": { "line": 0, "character": 12 }
-                }
-            })
-        );
-    }
-
-    #[test]
-    fn goto_definition_of_command() {
-        let (client_connection, _recv) = initialize_language_server();
-
-        let mut script = fixtures();
-        script.push("lsp");
-        script.push("goto");
-        script.push("command.nu");
-        let script = path_to_uri(&script);
-
-        open_unchecked(&client_connection, script.clone());
-
-        let resp = send_goto_definition_request(&client_connection, script.clone(), 4, 1);
-        let result = if let Message::Response(response) = resp {
-            response.result
-        } else {
-            panic!()
-        };
-
-        assert_json_eq!(
-            result,
-            serde_json::json!({
-                "uri": script,
-                "range": {
-                "start": { "line": 0, "character": 17 },
-                "end": { "line": 2, "character": 1 }
-            }
-            })
-        );
-    }
-
-    #[test]
-    fn goto_definition_of_command_unicode() {
-        let (client_connection, _recv) = initialize_language_server();
-
-        let mut script = fixtures();
-        script.push("lsp");
-        script.push("goto");
-        script.push("command_unicode.nu");
-        let script = path_to_uri(&script);
-
-        open_unchecked(&client_connection, script.clone());
-
-        let resp = send_goto_definition_request(&client_connection, script.clone(), 4, 2);
-        let result = if let Message::Response(response) = resp {
-            response.result
-        } else {
-            panic!()
-        };
-
-        assert_json_eq!(
-            result,
-            serde_json::json!({
-                "uri": script,
-                "range": {
-                "start": { "line": 0, "character": 19 },
-                "end": { "line": 2, "character": 1 }
-            }
-            })
-        );
-    }
-
-    #[test]
-    fn goto_definition_of_command_parameter() {
-        let (client_connection, _recv) = initialize_language_server();
-
-        let mut script = fixtures();
-        script.push("lsp");
-        script.push("goto");
-        script.push("command.nu");
-        let script = path_to_uri(&script);
-
-        open_unchecked(&client_connection, script.clone());
-
-        let resp = send_goto_definition_request(&client_connection, script.clone(), 1, 14);
-        let result = if let Message::Response(response) = resp {
-            response.result
-        } else {
-            panic!()
-        };
-
-        assert_json_eq!(
-            result,
-            serde_json::json!({
-                "uri": script,
-                "range": {
-                "start": { "line": 0, "character": 11 },
-                "end": { "line": 0, "character": 15 }
-            }
-            })
-        );
-    }
-
     pub fn send_hover_request(
         client_connection: &Connection,
         uri: Uri,

tests/fixtures/lsp/goto/collect.nu

@@ -0,0 +1,2 @@
+let foo = 1
+..2 | each { $foo }

tests/fixtures/lsp/goto/else.nu

@@ -0,0 +1,2 @@
+let foo = 1
+if true { } else { $foo }

tests/fixtures/lsp/goto/match.nu

@@ -0,0 +1,5 @@
+let foo = 1
+match $foo {
+  _ if $foo == 1 => 1
+  _ => 2
+}