refactor(lsp): flat_map with mutable accumulator (#15567)

# Description

Mainly performance improvement of lsp operations involving flat_map on
AST nodes.
Previous flat_map traversing is functional, which is a nice property to
have, but the heavy cost of vector collection on each tree node makes it
undesirable.

This PR mitigates the problem with a mutable accumulator.

# User-Facing Changes

Should be none.

# Tests + Formatting

# After Submitting

crates/nu-protocol/src/ast/traverse.rs

@@ -15,17 +15,19 @@ pub enum FindMapResult<T> {
 
 /// 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
+    /// Generic function that do flat_map on an AST node.
+    /// Concatenates all recursive results on sub-expressions
+    /// into the `results` accumulator.
     ///
     /// # Arguments
-    /// * `f` - function that overrides the default behavior
-    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
+    /// * `f` - function that generates leaf elements
+    /// * `results` - accumulator
+    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
     where
-        F: Fn(&'a Expression) -> Option<Vec<T>>;
+        F: Fn(&'a Expression) -> Vec<T>;
 
-    /// Generic function that do find_map on an AST node
-    /// return the first Some
+    /// Generic function that do find_map on an AST node.
+    /// Return the first result found by applying `f` on sub-expressions.
     ///
     /// # Arguments
     /// * `f` - function that overrides the default behavior
@@ -35,24 +37,18 @@ pub trait Traverse {
 }
 
 impl Traverse for Block {
-    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
+    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
     where
-        F: Fn(&'a Expression) -> Option<Vec<T>>,
+        F: Fn(&'a Expression) -> 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()
+        for pipeline in self.pipelines.iter() {
+            for element in pipeline.elements.iter() {
+                element.expr.flat_map(working_set, f, results);
+                if let Some(redir) = &element.redirection {
+                    redir.flat_map(working_set, f, results);
+                };
+            }
+        }
     }
 
     fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
@@ -71,21 +67,19 @@ impl Traverse for Block {
 }
 
 impl Traverse for PipelineRedirection {
-    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
+    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
     where
-        F: Fn(&'a Expression) -> Option<Vec<T>>,
+        F: Fn(&'a Expression) -> Vec<T>,
     {
-        let recur = |expr: &'a Expression| expr.flat_map(working_set, f);
+        let mut recur = |expr: &'a Expression| expr.flat_map(working_set, f, results);
+
         match self {
-            PipelineRedirection::Single { target, .. } => {
-                target.expr().map(recur).unwrap_or_default()
+            PipelineRedirection::Single { target, .. } => target.expr().map(recur),
+            PipelineRedirection::Separate { out, err } => {
+                out.expr().map(&mut recur);
+                err.expr().map(&mut recur)
             }
-            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>
@@ -94,9 +88,7 @@ impl Traverse for PipelineRedirection {
     {
         let recur = |expr: &'a Expression| expr.find_map(working_set, f);
         match self {
-            PipelineRedirection::Single { target, .. } => {
-                target.expr().map(recur).unwrap_or_default()
-            }
+            PipelineRedirection::Single { target, .. } => target.expr().and_then(recur),
             PipelineRedirection::Separate { out, err } => {
                 [out, err].iter().filter_map(|t| t.expr()).find_map(recur)
             }
@@ -105,87 +97,97 @@ impl Traverse for PipelineRedirection {
 }
 
 impl Traverse for Expression {
-    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
+    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
     where
-        F: Fn(&'a Expression) -> Option<Vec<T>>,
+        F: Fn(&'a Expression) -> 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);
+        // leaf elements generated by `f` for this expression
+        results.extend(f(self));
+        let mut recur = |expr: &'a Expression| expr.flat_map(working_set, f, results);
+
         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)
+                block.flat_map(working_set, f, results)
+            }
+            Expr::Range(range) => {
+                for sub_expr in [&range.from, &range.next, &range.to].into_iter().flatten() {
+                    recur(sub_expr);
+                }
+            }
+            Expr::Call(call) => {
+                for arg in &call.arguments {
+                    if let Some(sub_expr) = arg.expr() {
+                        recur(sub_expr);
+                    }
+                }
+            }
+            Expr::ExternalCall(head, args) => {
+                recur(head.as_ref());
+                for arg in args {
+                    recur(arg.expr());
+                }
             }
-            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::BinaryOp(lhs, op, rhs) => {
+                recur(lhs);
+                recur(op);
+                recur(rhs);
+            }
+            Expr::MatchBlock(matches) => {
+                for (pattern, expr) in matches {
+                    pattern.flat_map(working_set, f, results);
+                    expr.flat_map(working_set, f, results);
+                }
+            }
+            Expr::List(items) => {
+                for item in items {
+                    match item {
+                        ListItem::Item(expr) | ListItem::Spread(_, expr) => recur(expr),
+                    }
+                }
+            }
+            Expr::Record(items) => {
+                for item in items {
+                    match item {
+                        RecordItem::Spread(_, expr) => recur(expr),
+                        RecordItem::Pair(key, val) => {
+                            recur(key);
+                            recur(val);
+                        }
+                    }
+                }
+            }
+            Expr::Table(table) => {
+                for column in &table.columns {
+                    recur(column);
+                }
+                for row in &table.rows {
+                    for item in row {
+                        recur(item);
+                    }
+                }
+            }
             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()
+                for item in vec {
+                    recur(item);
+                }
+            }
+            Expr::AttributeBlock(ab) => {
+                for attr in &ab.attributes {
+                    recur(&attr.expr);
+                }
+                recur(&ab.item);
             }
-            Expr::AttributeBlock(ab) => ab
-                .attributes
-                .iter()
-                .flat_map(|attr| recur(&attr.expr))
-                .chain(recur(&ab.item))
-                .collect(),
 
-            _ => Vec::new(),
-        }
+            _ => (),
+        };
     }
 
     fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
@@ -203,7 +205,9 @@ impl Traverse for Expression {
                     | Expr::Subexpression(block_id)
                     | Expr::Block(block_id)
                     | Expr::Closure(block_id) => {
-                        let block = working_set.get_block(block_id.to_owned());
+                        // Clone the block_id to create an owned value
+                        let block_id = block_id.to_owned();
+                        let block = working_set.get_block(block_id);
                         block.find_map(working_set, f)
                     }
                     Expr::Range(range) => [&range.from, &range.next, &range.to]
@@ -253,25 +257,31 @@ impl Traverse for Expression {
 }
 
 impl Traverse for MatchPattern {
-    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
+    fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
     where
-        F: Fn(&'a Expression) -> Option<Vec<T>>,
+        F: Fn(&'a Expression) -> 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);
+        let mut recur_pattern =
+            |pattern: &'a MatchPattern| pattern.flat_map(working_set, f, results);
+
         match &self.pattern {
-            Pattern::Expression(expr) => recur(expr),
+            Pattern::Expression(expr) => expr.flat_map(working_set, f, results),
             Pattern::List(patterns) | Pattern::Or(patterns) => {
-                patterns.iter().flat_map(recur_pattern).collect()
+                for pattern in patterns {
+                    recur_pattern(pattern);
+                }
             }
             Pattern::Record(entries) => {
-                entries.iter().flat_map(|(_, p)| recur_pattern(p)).collect()
+                for (_, p) in entries {
+                    recur_pattern(p);
+                }
             }
-            _ => Vec::new(),
+            _ => (),
+        };
+
+        if let Some(g) = self.guard.as_ref() {
+            g.flat_map(working_set, f, results);
         }
-        .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>