Repository: nushell/nushell (language: rust)

Issue:
External completer needs more precise argument (spans) for completion task
### Describe the bug

1. `ls | grep --<tab>` will call the external completer with `["|" "grep" "--"]`, better to remove "|" here.
2. `(ls; git --<tab>)` will call the external completer with `["(", "ls", "git", "--"]`, failed to parse the correct command, also missing ";" here, internal completion fails either.
3. other complicated cases where command is surrounded by `{} ()`

### How to reproduce

1. type `ls | grep --`
2. press tab

### Expected behavior

Should be able to parse the whole command line into AST, and find the command node where current cursor is in.
Possibly bugs lie in current parser implementation.

### Configuration

| key                | value                                          |
| ------------------ | ---------------------------------------------- |
| version            | 0.98.0                                         |
| major              | 0                                              |
| minor              | 98                                             |
| patch              | 0                                              |
| branch             |                                                |
| commit_hash        |                                                |
| build_os           | macos-x86_64                                   |
| build_target       | x86_64-apple-darwin                            |
| rust_version       | rustc 1.81.0 (eeb90cda1 2024-09-04) (Homebrew) |
| cargo_version      | cargo 1.81.0                                   |
| build_time         | 2024-09-18 07:48:46 +00:00                     |
| build_rust_channel | release                                        |
| allocator          | mimalloc                                       |
| features           | default, sqlite, trash                         |
| installed_plugins  |                                                |


Discussion / hints:

Fix patch:
diff --git a/crates/nu-cli/src/completions/completer.rs b/crates/nu-cli/src/completions/completer.rs
index 71125e04e9950..9ab0cba5668de 100644
--- a/crates/nu-cli/src/completions/completer.rs
+++ b/crates/nu-cli/src/completions/completer.rs
@@ -5,8 +5,9 @@ use crate::completions::{
 use log::debug;
 use nu_color_config::{color_record_to_nustyle, lookup_ansi_color_style};
 use nu_engine::eval_block;
-use nu_parser::{flatten_pipeline_element, parse, FlatShape};
+use nu_parser::{flatten_expression, parse, FlatShape};
 use nu_protocol::{
+    ast::{Expr, Expression, FindMapResult, Traverse},
     debugger::WithoutDebug,
     engine::{Closure, EngineState, Stack, StateWorkingSet},
     PipelineData, Span, Value,
@@ -16,6 +17,51 @@ use std::{str, sync::Arc};
 
 use super::base::{SemanticSuggestion, SuggestionKind};
 
+fn find_pipeline_element_by_position<'a>(
+    expr: &'a Expression,
+    working_set: &'a StateWorkingSet,
+    pos: usize,
+) -> FindMapResult<&'a Expression> {
+    // skip the entire expression if the position is not in it
+    if !expr.span.contains(pos) {
+        return FindMapResult::Stop;
+    }
+    let closure = |expr: &'a Expression| find_pipeline_element_by_position(expr, working_set, pos);
+    match &expr.expr {
+        Expr::Call(call) => call
+            .arguments
+            .iter()
+            .find_map(|arg| arg.expr().and_then(|e| e.find_map(working_set, &closure)))
+            // if no inner call/external_call found, then this is the inner-most one
+            .or(Some(expr))
+            .map(FindMapResult::Found)
+            .unwrap_or_default(),
+        // TODO: clear separation of internal/external completion logic
+        Expr::ExternalCall(head, arguments) => arguments
+            .iter()
+            .find_map(|arg| arg.expr().find_map(working_set, &closure))
+            .or(head.as_ref().find_map(working_set, &closure))
+            .or(Some(expr))
+            .map(FindMapResult::Found)
+            .unwrap_or_default(),
+        // complete the operator
+        Expr::BinaryOp(lhs, _, rhs) => lhs
+            .find_map(working_set, &closure)
+            .or(rhs.find_map(working_set, &closure))
+            .or(Some(expr))
+            .map(FindMapResult::Found)
+            .unwrap_or_default(),
+        Expr::FullCellPath(fcp) => fcp
+            .head
+            .find_map(working_set, &closure)
+            .or(Some(expr))
+            .map(FindMapResult::Found)
+            .unwrap_or_default(),
+        Expr::Var(_) => FindMapResult::Found(expr),
+        _ => FindMapResult::Continue,
+    }
+}
+
 #[derive(Clone)]
 pub struct NuCompleter {
     engine_state: Arc<EngineState>,
@@ -140,106 +186,175 @@ impl NuCompleter {
 
         let config = self.engine_state.get_config();
 
-        let outermost_block = parse(&mut working_set, Some("completer"), line.as_bytes(), false);
+        let block = parse(&mut working_set, Some("completer"), line.as_bytes(), false);
+        let Some(element_expression) = block.find_map(&working_set, &|expr: &Expression| {
+            find_pipeline_element_by_position(expr, &working_set, pos)
+        }) else {
+            return vec![];
+        };
 
-        // Try to get the innermost block parsed (by span) so that we consider the correct context/scope.
-        let target_block = working_set
-            .delta
-            .blocks
-            .iter()
-            .filter_map(|block| match block.span {
-                Some(span) if span.contains(pos) => Some((block, span)),
-                _ => None,
-            })
-            .reduce(|prev, cur| {
-                // |(block, span), (block, span)|
-                match cur.1.start.cmp(&prev.1.start) {
-                    core::cmp::Ordering::Greater => cur,
-                    core::cmp::Ordering::Equal if cur.1.end < prev.1.end => cur,
-                    _ => prev,
+        let flattened = flatten_expression(&working_set, element_expression);
+        let mut spans: Vec<String> = vec![];
+
+        for (flat_idx, (span, shape)) in flattened.iter().enumerate() {
+            let is_passthrough_command = spans
+                .first()
+                .filter(|content| content.as_str() == "sudo" || content.as_str() == "doas")
+                .is_some();
+
+            // Read the current span to string
+            let current_span = working_set.get_span_contents(*span);
+            let current_span_str = String::from_utf8_lossy(current_span);
+            let is_last_span = span.contains(pos);
+
+            // Skip the last 'a' as span item
+            if is_last_span {
+                let offset = pos - span.start;
+                if offset == 0 {
+                    spans.push(String::new())
+                } else {
+                    let mut current_span_str = current_span_str.to_string();
+                    current_span_str.remove(offset);
+                    spans.push(current_span_str);
                 }
-            })
-            .map(|(block, _)| block)
-            .unwrap_or(&outermost_block);
-
-        for pipeline in &target_block.pipelines {
-            for pipeline_element in &pipeline.elements {
-                let flattened = flatten_pipeline_element(&working_set, pipeline_element);
-                let mut spans: Vec<String> = vec![];
-
-                for (flat_idx, flat) in flattened.iter().enumerate() {
-                    let is_passthrough_command = spans
-                        .first()
-                        .filter(|content| content.as_str() == "sudo" || content.as_str() == "doas")
-                        .is_some();
-
-                    // Read the current span to string
-                    let current_span = working_set.get_span_contents(flat.0);
-                    let current_span_str = String::from_utf8_lossy(current_span);
-                    let is_last_span = pos >= flat.0.start && pos < flat.0.end;
-
-                    // Skip the last 'a' as span item
-                    if is_last_span {
-                        let offset = pos - flat.0.start;
-                        if offset == 0 {
-                            spans.push(String::new())
-                        } else {
-                            let mut current_span_str = current_span_str.to_string();
-                            current_span_str.remove(offset);
-                            spans.push(current_span_str);
-                        }
-                    } else {
-                        spans.push(current_span_str.to_string());
+            } else {
+                spans.push(current_span_str.to_string());
+            }
+
+            // Complete based on the last span
+            if is_last_span {
+                // Context variables
+                let most_left_var = most_left_variable(flat_idx, &working_set, flattened.clone());
+
+                // Create a new span
+                let new_span = Span::new(span.start, span.end - 1);
+
+                // Parses the prefix. Completion should look up to the cursor position, not after.
+                let index = pos - span.start;
+                let prefix = &current_span[..index];
+
+                // Variables completion
+                if prefix.starts_with(b"$") || most_left_var.is_some() {
+                    let mut variable_names_completer =
+                        VariableCompletion::new(most_left_var.unwrap_or((vec![], vec![])));
+
+                    let mut variable_completions = self.process_completion(
+                        &mut variable_names_completer,
+                        &working_set,
+                        prefix,
+                        new_span,
+                        fake_offset,
+                        pos,
+                    );
+
+                    let mut variable_operations_completer =
+                        OperatorCompletion::new(element_expression.clone());
+
+                    let mut variable_operations_completions = self.process_completion(
+                        &mut variable_operations_completer,
+                        &working_set,
+                        prefix,
+                        new_span,
+                        fake_offset,
+                        pos,
+                    );
+
+                    variable_completions.append(&mut variable_operations_completions);
+                    return variable_completions;
+                }
+
+                // Flags completion
+                if prefix.starts_with(b"-") {
+                    // Try to complete flag internally
+                    let mut completer = FlagCompletion::new(element_expression.clone());
+                    let result = self.process_completion(
+                        &mut completer,
+                        &working_set,
+                        prefix,
+                        new_span,
+                        fake_offset,
+                        pos,
+                    );
+
+                    if !result.is_empty() {
+                        return result;
                     }
 
-                    // Complete based on the last span
-                    if is_last_span {
-                        // Context variables
-                        let most_left_var =
-                            most_left_variable(flat_idx, &working_set, flattened.clone());
+                    // We got no results for internal completion
+                    // now we can check if external completer is set and use it
+                    if let Some(closure) = config.completions.external.completer.as_ref() {
+                        if let Some(external_result) =
+                            self.external_completion(closure, &spans, fake_offset, new_span)
+                        {
+                            return external_result;
+                        }
+                    }
+                }
 
-                        // Create a new span
-                        let new_span = Span::new(flat.0.start, flat.0.end - 1);
+                // specially check if it is currently empty - always complete commands
+                if (is_passthrough_command && flat_idx == 1)
+                    || (flat_idx == 0 && working_set.get_span_contents(new_span).is_empty())
+                {
+                    let mut completer = CommandCompletion::new(
+                        flattened.clone(),
+                        // flat_idx,
+                        FlatShape::String,
+                        true,
+                    );
+                    return self.process_completion(
+                        &mut completer,
+                        &working_set,
+                        prefix,
+                        new_span,
+                        fake_offset,
+                        pos,
+                    );
+                }
 
-                        // Parses the prefix. Completion should look up to the cursor position, not after.
-                        let index = pos - flat.0.start;
-                        let prefix = &current_span[..index];
+                // Completions that depends on the previous expression (e.g: use, source-env)
+                if (is_passthrough_command && flat_idx > 1) || flat_idx > 0 {
+                    if let Some(previous_expr) = flattened.get(flat_idx - 1) {
+                        // Read the content for the previous expression
+                        let prev_expr_str = working_set.get_span_contents(previous_expr.0).to_vec();
 
-                        // Variables completion
-                        if prefix.starts_with(b"$") || most_left_var.is_some() {
-                            let mut variable_names_completer =
-                                VariableCompletion::new(most_left_var.unwrap_or((vec![], vec![])));
+                        // Completion for .nu files
+                        if prev_expr_str == b"use"
+                            || prev_expr_str == b"overlay use"
+                            || prev_expr_str == b"source-env"
+                        {
+                            let mut completer = DotNuCompletion::new();
 
-                            let mut variable_completions = self.process_completion(
-                                &mut variable_names_completer,
+                            return self.process_completion(
+                                &mut completer,
                                 &working_set,
                                 prefix,
                                 new_span,
                                 fake_offset,
                                 pos,
                             );
+                        } else if prev_expr_str == b"ls" {
+                            let mut completer = FileCompletion::new();
 
-                            let mut variable_operations_completer =
-                                OperatorCompletion::new(pipeline_element.expr.clone());
-
-                            let mut variable_operations_completions = self.process_completion(
-                                &mut variable_operations_completer,
+                            return self.process_completion(
+                                &mut completer,
                                 &working_set,
                                 prefix,
                                 new_span,
                                 fake_offset,
                                 pos,
                             );
-
-                            variable_completions.append(&mut variable_operations_completions);
-                            return variable_completions;
-                        }
-
-                        // Flags completion
-                        if prefix.starts_with(b"-") {
-                            // Try to complete flag internally
-                            let mut completer = FlagCompletion::new(pipeline_element.expr.clone());
-                            let result = self.process_completion(
+                        } else if matches!(
+                            previous_expr.1,
+                            FlatShape::Float
+                                | FlatShape::Int
+                                | FlatShape::String
+                                | FlatShape::List
+                                | FlatShape::Bool
+                                | FlatShape::Variable(_)
+                        ) {
+                            let mut completer = OperatorCompletion::new(element_expression.clone());
+
+                            let operator_suggestion = self.process_completion(
                                 &mut completer,
                                 &working_set,
                                 prefix,
@@ -247,198 +362,102 @@ impl NuCompleter {
                                 fake_offset,
                                 pos,
                             );
-
-                            if !result.is_empty() {
-                                return result;
-                            }
-
-                            // We got no results for internal completion
-                            // now we can check if external completer is set and use it
-                            if let Some(closure) = config.completions.external.completer.as_ref() {
-                                if let Some(external_result) =
-                                    self.external_completion(closure, &spans, fake_offset, new_span)
-                                {
-                                    return external_result;
-                                }
+                            if !operator_suggestion.is_empty() {
+                                return operator_suggestion;
                             }
                         }
+                    }
+                }
 
-                        // specially check if it is currently empty - always complete commands
-                        if (is_passthrough_command && flat_idx == 1)
-                            || (flat_idx == 0 && working_set.get_span_contents(new_span).is_empty())
-                        {
-                            let mut completer = CommandCompletion::new(
-                                flattened.clone(),
-                                // flat_idx,
-                                FlatShape::String,
-                                true,
-                            );
-                            return self.process_completion(
-                                &mut completer,
-                                &working_set,
-                                prefix,
-                                new_span,
-                                fake_offset,
-                                pos,
-                            );
+                // Match other types
+                match shape {
+                    FlatShape::Custom(decl_id) => {
+                        let mut completer = CustomCompletion::new(
+                            self.stack.clone(),
+                            *decl_id,
+                            initial_line,
+                            FileCompletion::new(),
+                        );
+
+                        return self.process_completion(
+                            &mut completer,
+                            &working_set,
+                            prefix,
+                            new_span,
+                            fake_offset,
+                            pos,
+                        );
+                    }
+                    FlatShape::Directory => {
+                        let mut completer = DirectoryCompletion::new();
+
+                        return self.process_completion(
+                            &mut completer,
+                            &working_set,
+                            prefix,
+                            new_span,
+                            fake_offset,
+                            pos,
+                        );
+                    }
+                    FlatShape::Filepath | FlatShape::GlobPattern => {
+                        let mut completer = FileCompletion::new();
+
+                        return self.process_completion(
+                            &mut completer,
+                            &working_set,
+                            prefix,
+                            new_span,
+                            fake_offset,
+                            pos,
+                        );
+                    }
+                    flat_shape => {
+                        let mut completer = CommandCompletion::new(
+                            flattened.clone(),
+                            // flat_idx,
+                            flat_shape.clone(),
+                            false,
+                        );
+
+                        let mut out: Vec<_> = self.process_completion(
+                            &mut completer,
+                            &working_set,
+                            prefix,
+                            new_span,
+                            fake_offset,
+                            pos,
+                        );
+
+                        if !out.is_empty() {
+                            return out;
                         }
 
-                        // Completions that depends on the previous expression (e.g: use, source-env)
-                        if (is_passthrough_command && flat_idx > 1) || flat_idx > 0 {
-                            if let Some(previous_expr) = flattened.get(flat_idx - 1) {
-                                // Read the content for the previous expression
-                                let prev_expr_str =
-                                    working_set.get_span_contents(previous_expr.0).to_vec();
-
-                                // Completion for .nu files
-                                if prev_expr_str == b"use"
-                                    || prev_expr_str == b"overlay use"
-                                    || prev_expr_str == b"source-env"
-                                {
-                                    let mut completer = DotNuCompletion::new();
-
-                                    return self.process_completion(
-                                        &mut completer,
-                                        &working_set,
-                                        prefix,
-                                        new_span,
-                                        fake_offset,
-                                        pos,
-                                    );
-                                } else if prev_expr_str == b"ls" {
-                                    let mut completer = FileCompletion::new();
-
-                                    return self.process_completion(
-                                        &mut completer,
-                                        &working_set,
-                                        prefix,
-                                        new_span,
-                                        fake_offset,
-                                        pos,
-                                    );
-                                } else if matches!(
-                                    previous_expr.1,
-                                    FlatShape::Float
-                                        | FlatShape::Int
-                                        | FlatShape::String
-                                        | FlatShape::List
-                                        | FlatShape::Bool
-                                        | FlatShape::Variable(_)
-                                ) {
-                                    let mut completer =
-                                        OperatorCompletion::new(pipeline_element.expr.clone());
-
-                                    let operator_suggestion = self.process_completion(
-                                        &mut completer,
-                                        &working_set,
-                                        prefix,
-                                        new_span,
-                                        fake_offset,
-                                        pos,
-                                    );
-                                    if !operator_suggestion.is_empty() {
-                                        return operator_suggestion;
-                                    }
-                                }
+                        // Try to complete using an external completer (if set)
+                        if let Some(closure) = config.completions.external.completer.as_ref() {
+                            if let Some(external_result) =
+                                self.external_completion(closure, &spans, fake_offset, new_span)
+                            {
+                                return external_result;
                             }
                         }
 
-                        // Match other types
-                        match &flat.1 {
-                            FlatShape::Custom(decl_id) => {
-                                let mut completer = CustomCompletion::new(
-                                    self.stack.clone(),
-                                    *decl_id,
-                                    initial_line,
-                                    FileCompletion::new(),
-                                );
-
-                                return self.process_completion(
-                                    &mut completer,
-                                    &working_set,
-                                    prefix,
-                                    new_span,
-                                    fake_offset,
-                                    pos,
-                                );
-                            }
-                            FlatShape::Directory => {
-                                let mut completer = DirectoryCompletion::new();
-
-                                return self.process_completion(
-                                    &mut completer,
-                                    &working_set,
-                                    prefix,
-                                    new_span,
-                                    fake_offset,
-                                    pos,
-                                );
-                            }
-                            FlatShape::Filepath | FlatShape::GlobPattern => {
-                                let mut completer = FileCompletion::new();
-
-                                return self.process_completion(
-                                    &mut completer,
-                                    &working_set,
-                                    prefix,
-                                    new_span,
-                                    fake_offset,
-                                    pos,
-                                );
-                            }
-                            flat_shape => {
-                                let mut completer = CommandCompletion::new(
-                                    flattened.clone(),
-                                    // flat_idx,
-                                    flat_shape.clone(),
-                                    false,
-                                );
-
-                                let mut out: Vec<_> = self.process_completion(
-                                    &mut completer,
-                                    &working_set,
-                                    prefix,
-                                    new_span,
-                                    fake_offset,
-                                    pos,
-                                );
-
-                                if !out.is_empty() {
-                                    return out;
-                                }
-
-                                // Try to complete using an external completer (if set)
-                                if let Some(closure) =
-                                    config.completions.external.completer.as_ref()
-                                {
-                                    if let Some(external_result) = self.external_completion(
-                                        closure,
-                                        &spans,
-                                        fake_offset,
-                                        new_span,
-                                    ) {
-                                        return external_result;
-                                    }
-                                }
-
-                                // Check for file completion
-                                let mut completer = FileCompletion::new();
-                                out = self.process_completion(
-                                    &mut completer,
-                                    &working_set,
-                                    prefix,
-                                    new_span,
-                                    fake_offset,
-                                    pos,
-                                );
-
-                                if !out.is_empty() {
-                                    return out;
-                                }
-                            }
-                        };
+                        // Check for file completion
+                        let mut completer = FileCompletion::new();
+                        out = self.process_completion(
+                            &mut completer,
+                            &working_set,
+                            prefix,
+                            new_span,
+                            fake_offset,
+                            pos,
+                        );
+
+                        if !out.is_empty() {
+                            return out;
+                        }
                     }
-                }
+                };
             }
         }
 
diff --git a/crates/nu-lsp/src/ast.rs b/crates/nu-lsp/src/ast.rs
index 511a25ca6b117..775093c1334d5 100644
--- a/crates/nu-lsp/src/ast.rs
+++ b/crates/nu-lsp/src/ast.rs
@@ -1,172 +1,11 @@
 use crate::Id;
 use nu_protocol::{
-    ast::{
-        Argument, Block, Call, Expr, Expression, ExternalArgument, ListItem, MatchPattern,
-        PathMember, Pattern, PipelineRedirection, RecordItem,
-    },
+    ast::{Argument, Block, Call, Expr, Expression, FindMapResult, ListItem, PathMember, Traverse},
     engine::StateWorkingSet,
     Span,
 };
 use std::sync::Arc;
 
-/// similar to flatten_block, but allows extra map function
-pub fn ast_flat_map<'a, T, F>(
-    ast: &'a Arc<Block>,
-    working_set: &'a StateWorkingSet,
-    f_special: &F,
-) -> Vec<T>
-where
-    F: Fn(&'a Expression) -> Option<Vec<T>>,
-{
-    ast.pipelines
-        .iter()
-        .flat_map(|pipeline| {
-            pipeline.elements.iter().flat_map(|element| {
-                expr_flat_map(&element.expr, working_set, f_special)
-                    .into_iter()
-                    .chain(
-                        element
-                            .redirection
-                            .as_ref()
-                            .map(|redir| redirect_flat_map(redir, working_set, 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<'a, T, F>(
-    expr: &'a Expression,
-    working_set: &'a StateWorkingSet,
-    f_special: &F,
-) -> Vec<T>
-where
-    F: Fn(&'a Expression) -> Option<Vec<T>>,
-{
-    // behavior overridden by f_special
-    if let Some(vec) = f_special(expr) {
-        return vec;
-    }
-    let recur = |expr| expr_flat_map(expr, working_set, 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, 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, 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<'a, T, F>(
-    pattern: &'a MatchPattern,
-    working_set: &'a StateWorkingSet,
-    f_special: &F,
-) -> Vec<T>
-where
-    F: Fn(&'a Expression) -> Option<Vec<T>>,
-{
-    let recur = |expr| expr_flat_map(expr, working_set, f_special);
-    let recur_match = |p| match_pattern_flat_map(p, working_set, 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<'a, T, F>(
-    redir: &'a PipelineRedirection,
-    working_set: &'a StateWorkingSet,
-    f_special: &F,
-) -> Vec<T>
-where
-    F: Fn(&'a Expression) -> Option<Vec<T>>,
-{
-    let recur = |expr| expr_flat_map(expr, working_set, 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(),
-    }
-}
-
 /// Adjust span if quoted
 fn strip_quotes(span: Span, working_set: &StateWorkingSet) -> Span {
     let text = String::from_utf8_lossy(working_set.get_span_contents(span));
@@ -436,22 +275,22 @@ fn find_id_in_expr(
     expr: &Expression,
     working_set: &StateWorkingSet,
     location: &usize,
-) -> Option<Vec<(Id, Span)>> {
+) -> FindMapResult<(Id, Span)> {
     // skip the entire expression if the location is not in it
     if !expr.span.contains(*location) {
-        return Some(Vec::new());
+        return FindMapResult::Stop;
     }
     let span = expr.span;
     match &expr.expr {
-        Expr::VarDecl(var_id) => Some(vec![(Id::Variable(*var_id), span)]),
+        Expr::VarDecl(var_id) => FindMapResult::Found((Id::Variable(*var_id), span)),
         // trim leading `$` sign
-        Expr::Var(var_id) => Some(vec![(
+        Expr::Var(var_id) => FindMapResult::Found((
             Id::Variable(*var_id),
             Span::new(span.start.saturating_add(1), span.end),
-        )]),
+        )),
         Expr::Call(call) => {
             if call.head.contains(*location) {
-                Some(vec![(Id::Declaration(call.decl_id), call.head)])
+                FindMapResult::Found((Id::Declaration(call.decl_id), call.head))
             } else {
                 try_find_id_in_def(call, working_set, Some(location), None)
                     .or(try_find_id_in_mod(call, working_set, Some(location), None))
@@ -462,19 +301,20 @@ fn find_id_in_expr(
                         Some(location),
                         None,
                     ))
-                    .map(|p| vec![p])
+                    .map(FindMapResult::Found)
+                    .unwrap_or_default()
             }
         }
         Expr::FullCellPath(fcp) => {
             if fcp.head.span.contains(*location) {
-                None
+                FindMapResult::Continue
             } else {
                 let Expression {
                     expr: Expr::Var(var_id),
                     ..
                 } = fcp.head
                 else {
-                    return None;
+                    return FindMapResult::Continue;
                 };
                 let tail: Vec<PathMember> = fcp
                     .tail
@@ -482,11 +322,13 @@ fn find_id_in_expr(
                     .into_iter()
                     .take_while(|pm| pm.span().start <= *location)
                     .collect();
-                let span = tail.last()?.span();
-                Some(vec![(Id::CellPath(var_id, tail), span)])
+                let Some(span) = tail.last().map(|pm| pm.span()) else {
+                    return FindMapResult::Stop;
+                };
+                FindMapResult::Found((Id::CellPath(var_id, tail), span))
             }
         }
-        Expr::Overlay(Some(module_id)) => Some(vec![(Id::Module(*module_id), span)]),
+        Expr::Overlay(Some(module_id)) => FindMapResult::Found((Id::Module(*module_id), span)),
         // terminal value expressions
         Expr::Bool(_)
         | Expr::Binary(_)
@@ -500,8 +342,8 @@ fn find_id_in_expr(
         | Expr::Nothing
         | Expr::RawString(_)
         | Expr::Signature(_)
-        | Expr::String(_) => Some(vec![(Id::Value(expr.ty.clone()), span)]),
-        _ => None,
+        | Expr::String(_) => FindMapResult::Found((Id::Value(expr.ty.clone()), span)),
+        _ => FindMapResult::Continue,
     }
 }
 
@@ -512,7 +354,7 @@ pub(crate) fn find_id(
     location: &usize,
 ) -> Option<(Id, Span)> {
     let closure = |e| find_id_in_expr(e, working_set, location);
-    ast_flat_map(ast, working_set, &closure).first().cloned()
+    ast.find_map(working_set, &closure)
 }
 
 fn find_reference_by_id_in_expr(
@@ -521,7 +363,6 @@ fn find_reference_by_id_in_expr(
     id: &Id,
 ) -> Option<Vec<Span>> {
     let closure = |e| find_reference_by_id_in_expr(e, working_set, id);
-    let recur = |expr| expr_flat_map(expr, working_set, &closure);
     match (&expr.expr, id) {
         (Expr::Var(vid1), Id::Variable(vid2)) if *vid1 == *vid2 => Some(vec![Span::new(
             // we want to exclude the `$` sign for renaming
@@ -536,7 +377,7 @@ fn find_reference_by_id_in_expr(
                 .arguments
                 .iter()
                 .filter_map(|arg| arg.expr())
-                .flat_map(recur)
+                .flat_map(|e| e.flat_map(working_set, &closure))
                 .collect();
             if matches!(id, Id::Declaration(decl_id) if call.decl_id == *decl_id) {
                 occurs.push(call.head);
@@ -560,7 +401,7 @@ pub(crate) fn find_reference_by_id(
     working_set: &StateWorkingSet,
     id: &Id,
 ) -> Vec<Span> {
-    ast_flat_map(ast, working_set, &|e| {
+    ast.flat_map(working_set, &|e| {
         find_reference_by_id_in_expr(e, working_set, id)
     })
 }
diff --git a/crates/nu-lsp/src/hints.rs b/crates/nu-lsp/src/hints.rs
index 96c191c7db80e..581a2e70b0931 100644
--- a/crates/nu-lsp/src/hints.rs
+++ b/crates/nu-lsp/src/hints.rs
@@ -1,4 +1,3 @@
-use crate::ast::{ast_flat_map, expr_flat_map};
 use crate::{span_to_range, LanguageServer};
 use lsp_textdocument::FullTextDocument;
 use lsp_types::{
@@ -6,7 +5,7 @@ use lsp_types::{
     MarkupKind, Position, Range,
 };
 use nu_protocol::{
-    ast::{Argument, Block, Expr, Expression, Operator},
+    ast::{Argument, Block, Expr, Expression, Operator, Traverse},
     engine::StateWorkingSet,
     Type,
 };
@@ -29,11 +28,12 @@ fn extract_inlay_hints_from_expression(
     file: &FullTextDocument,
 ) -> Option<Vec<InlayHint>> {
     let closure = |e| extract_inlay_hints_from_expression(e, working_set, offset, file);
-    let recur = |expr| expr_flat_map(expr, working_set, &closure);
     match &expr.expr {
         Expr::BinaryOp(lhs, op, rhs) => {
-            let mut hints: Vec<InlayHint> =
-                [lhs, op, rhs].into_iter().flat_map(|e| recur(e)).collect();
+            let mut hints: Vec<InlayHint> = [lhs, op, rhs]
+                .into_iter()
+                .flat_map(|e| e.flat_map(working_set, &closure))
+                .collect();
             if let Expr::Operator(Operator::Assignment(_)) = op.expr {
                 let position = span_to_range(&lhs.span, file, *offset).end;
                 let type_rhs = type_short_name(&rhs.ty);
@@ -103,13 +103,13 @@ fn extract_inlay_hints_from_expression(
                 match arg {
                     // skip the rest when spread/unknown arguments encountered
                     Argument::Spread(expr) | Argument::Unknown(expr) => {
-                        hints.extend(recur(expr));
+                        hints.extend(expr.flat_map(working_set, &closure));
                         sig_idx = signatures.len();
                         continue;
                     }
                     // skip current for flags
                     Argument::Named((_, _, Some(expr))) => {
-                        hints.extend(recur(expr));
+                        hints.extend(expr.flat_map(working_set, &closure));
                         continue;
                     }
                     Argument::Positional(expr) => {
@@ -130,7 +130,7 @@ fn extract_inlay_hints_from_expression(
                                 padding_right: None,
                             });
                         }
-                        hints.extend(recur(expr));
+                        hints.extend(expr.flat_map(working_set, &closure));
                     }
                     _ => {
                         continue;
@@ -154,7 +154,7 @@ impl LanguageServer {
         offset: usize,
         file: &FullTextDocument,
     ) -> Vec<InlayHint> {
-        ast_flat_map(block, working_set, &|e| {
+        block.flat_map(working_set, &|e| {
             extract_inlay_hints_from_expression(e, working_set, &offset, file)
         })
     }
diff --git a/crates/nu-protocol/src/ast/call.rs b/crates/nu-protocol/src/ast/call.rs
index 0c42583c5cb67..7a0d2b0cb4ca9 100644
--- a/crates/nu-protocol/src/ast/call.rs
+++ b/crates/nu-protocol/src/ast/call.rs
@@ -76,6 +76,15 @@ pub enum ExternalArgument {
     Spread(Expression),
 }
 
+impl ExternalArgument {
+    pub fn expr(&self) -> &Expression {
+        match self {
+            ExternalArgument::Regular(expr) => expr,
+            ExternalArgument::Spread(expr) => expr,
+        }
+    }
+}
+
 /// Parsed call of a `Command`
 ///
 /// As we also implement some internal keywords in terms of the `Command` trait, this type stores the passed arguments as [`Expression`].
diff --git a/crates/nu-protocol/src/ast/mod.rs b/crates/nu-protocol/src/ast/mod.rs
index cb0916880597e..0138431ec6077 100644
--- a/crates/nu-protocol/src/ast/mod.rs
+++ b/crates/nu-protocol/src/ast/mod.rs
@@ -11,6 +11,7 @@ mod operator;
 mod pipeline;
 mod range;
 mod table;
+mod traverse;
 mod unit;
 mod value_with_unit;
 
@@ -26,5 +27,6 @@ pub use operator::*;
 pub use pipeline::*;
 pub use range::*;
 pub use table::Table;
+pub use traverse::*;
 pub use unit::*;
 pub use value_with_unit::*;
diff --git a/crates/nu-protocol/src/ast/traverse.rs b/crates/nu-protocol/src/ast/traverse.rs
new file mode 100644
index 0000000000000..2c9e6dbaa1304
--- /dev/null
+++ b/crates/nu-protocol/src/ast/traverse.rs
@@ -0,0 +1,282 @@
+use crate::engine::StateWorkingSet;
+
+use super::{
+    Block, Expr, Expression, ListItem, MatchPattern, Pattern, PipelineRedirection, RecordItem,
+};
+
+/// Result of find_map closure
+#[derive(Default)]
+pub enum FindMapResult<T> {
+    Found(T),
+    #[default]
+    Continue,
+    Stop,
+}
+
+/// Trait for traversing the AST
+pub trait Traverse {
+    /// Generic function that do flat_map on an AST node
+    /// concatenates all recursive results on sub-expressions
+    ///
+    /// # Arguments
+    /// * `f` - function that overrides the default behavior
+    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
+    where
+        F: Fn(&'a Expression) -> Option<Vec<T>>;
+
+    /// Generic function that do find_map on an AST node
+    /// return the first Some
+    ///
+    /// # Arguments
+    /// * `f` - function that overrides the default behavior
+    fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
+    where
+        F: Fn(&'a Expression) -> FindMapResult<T>;
+}
+
+impl Traverse for Block {
+    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
+    where
+        F: Fn(&'a Expression) -> Option<Vec<T>>,
+    {
+        self.pipelines
+            .iter()
+            .flat_map(|pipeline| {
+                pipeline.elements.iter().flat_map(|element| {
+                    element.expr.flat_map(working_set, f).into_iter().chain(
+                        element
+                            .redirection
+                            .as_ref()
+                            .map(|redir| redir.flat_map(working_set, f))
+                            .unwrap_or_default(),
+                    )
+                })
+            })
+            .collect()
+    }
+
+    fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
+    where
+        F: Fn(&'a Expression) -> FindMapResult<T>,
+    {
+        self.pipelines.iter().find_map(|pipeline| {
+            pipeline.elements.iter().find_map(|element| {
+                element.expr.find_map(working_set, f).or(element
+                    .redirection
+                    .as_ref()
+                    .and_then(|redir| redir.find_map(working_set, f)))
+            })
+        })
+    }
+}
+
+impl Traverse for PipelineRedirection {
+    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
+    where
+        F: Fn(&'a Expression) -> Option<Vec<T>>,
+    {
+        let recur = |expr: &'a Expression| expr.flat_map(working_set, f);
+        match self {
+            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_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
+    where
+        F: Fn(&'a Expression) -> FindMapResult<T>,
+    {
+        let recur = |expr: &'a Expression| expr.find_map(working_set, f);
+        match self {
+            PipelineRedirection::Single { target, .. } => {
+                target.expr().map(recur).unwrap_or_default()
+            }
+            PipelineRedirection::Separate { out, err } => {
+                [out, err].iter().filter_map(|t| t.expr()).find_map(recur)
+            }
+        }
+    }
+}
+
+impl Traverse for Expression {
+    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
+    where
+        F: Fn(&'a Expression) -> Option<Vec<T>>,
+    {
+        // behavior overridden by f
+        if let Some(vec) = f(self) {
+            return vec;
+        }
+        let recur = |expr: &'a Expression| expr.flat_map(working_set, f);
+        match &self.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());
+                block.flat_map(working_set, f)
+            }
+            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.as_ref())
+                .into_iter()
+                .chain(args.iter().flat_map(|arg| recur(arg.expr())))
+                .collect(),
+            Expr::UnaryNot(expr) | Expr::Collect(_, expr) => recur(expr.as_ref()),
+            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)| {
+                    pattern
+                        .flat_map(working_set, f)
+                        .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(),
+        }
+    }
+
+    fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
+    where
+        F: Fn(&'a Expression) -> FindMapResult<T>,
+    {
+        // behavior overridden by f
+        match f(self) {
+            FindMapResult::Found(t) => Some(t),
+            FindMapResult::Stop => None,
+            FindMapResult::Continue => {
+                let recur = |expr: &'a Expression| expr.find_map(working_set, f);
+                match &self.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());
+                        block.find_map(working_set, f)
+                    }
+                    Expr::Range(range) => [&range.from, &range.next, &range.to]
+                        .iter()
+                        .find_map(|e| e.as_ref().and_then(recur)),
+                    Expr::Call(call) => call
+                        .arguments
+                        .iter()
+                        .find_map(|arg| arg.expr().and_then(recur)),
+                    Expr::ExternalCall(head, args) => {
+                        recur(head.as_ref()).or(args.iter().find_map(|arg| recur(arg.expr())))
+                    }
+                    Expr::UnaryNot(expr) | Expr::Collect(_, expr) => recur(expr.as_ref()),
+                    Expr::BinaryOp(lhs, op, rhs) => recur(lhs).or(recur(op)).or(recur(rhs)),
+                    Expr::MatchBlock(matches) => matches.iter().find_map(|(pattern, expr)| {
+                        pattern.find_map(working_set, f).or(recur(expr))
+                    }),
+                    Expr::List(items) => items.iter().find_map(|item| match item {
+                        ListItem::Item(expr) | ListItem::Spread(_, expr) => recur(expr),
+                    }),
+                    Expr::Record(items) => items.iter().find_map(|item| match item {
+                        RecordItem::Spread(_, expr) => recur(expr),
+                        RecordItem::Pair(key, val) => [key, val].into_iter().find_map(recur),
+                    }),
+                    Expr::Table(table) => table
+                        .columns
+                        .iter()
+                        .find_map(recur)
+                        .or(table.rows.iter().find_map(|row| row.iter().find_map(recur))),
+                    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().find_map(recur)
+                    }
+
+                    _ => None,
+                }
+            }
+        }
+    }
+}
+
+impl Traverse for MatchPattern {
+    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
+    where
+        F: Fn(&'a Expression) -> Option<Vec<T>>,
+    {
+        let recur = |expr: &'a Expression| expr.flat_map(working_set, f);
+        let recur_pattern = |pattern: &'a MatchPattern| pattern.flat_map(working_set, f);
+        match &self.pattern {
+            Pattern::Expression(expr) => recur(expr),
+            Pattern::List(patterns) | Pattern::Or(patterns) => {
+                patterns.iter().flat_map(recur_pattern).collect()
+            }
+            Pattern::Record(entries) => {
+                entries.iter().flat_map(|(_, p)| recur_pattern(p)).collect()
+            }
+            _ => Vec::new(),
+        }
+        .into_iter()
+        .chain(self.guard.as_ref().map(|g| recur(g)).unwrap_or_default())
+        .collect()
+    }
+
+    fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
+    where
+        F: Fn(&'a Expression) -> FindMapResult<T>,
+    {
+        let recur = |expr: &'a Expression| expr.find_map(working_set, f);
+        let recur_pattern = |pattern: &'a MatchPattern| pattern.find_map(working_set, f);
+        match &self.pattern {
+            Pattern::Expression(expr) => recur(expr),
+            Pattern::List(patterns) | Pattern::Or(patterns) => {
+                patterns.iter().find_map(recur_pattern)
+            }
+            Pattern::Record(entries) => entries.iter().find_map(|(_, p)| recur_pattern(p)),
+            _ => None,
+        }
+        .or(self.guard.as_ref().and_then(|g| recur(g)))
+    }
+}


Test patch: