Allocate less when doing a capture discovery (#8903)

# Description

This should be a little more efficient when running the algorithm to
find the captured variables.

# 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 -A
clippy::needless_collect` to check that you're using the standard code
style
- `cargo test --workspace` to check that all tests pass
- `cargo run -- crates/nu-std/tests/run.nu` 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
> ```
-->

# 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-parser/src/parser.rs

@@ -5357,9 +5357,8 @@ pub fn discover_captures_in_closure(
     block: &Block,
     seen: &mut Vec<VarId>,
     seen_blocks: &mut HashMap<BlockId, Vec<(VarId, Span)>>,
-) -> Result<Vec<(VarId, Span)>, ParseError> {
+    output: &mut Vec<(VarId, Span)>,
-    let mut output = vec![];
+) -> Result<(), ParseError> {
-
     for flag in &block.signature.named {
         if let Some(var_id) = flag.var_id {
             seen.push(var_id);
@@ -5383,11 +5382,10 @@ pub fn discover_captures_in_closure(
     }
 
     for pipeline in &block.pipelines {
-        let result = discover_captures_in_pipeline(working_set, pipeline, seen, seen_blocks)?;
+        discover_captures_in_pipeline(working_set, pipeline, seen, seen_blocks, output)?;
-        output.extend(&result);
     }
 
-    Ok(output)
+    Ok(())
 }
 
 fn discover_captures_in_pipeline(
@@ -5395,15 +5393,13 @@ fn discover_captures_in_pipeline(
     pipeline: &Pipeline,
     seen: &mut Vec<VarId>,
     seen_blocks: &mut HashMap<BlockId, Vec<(VarId, Span)>>,
-) -> Result<Vec<(VarId, Span)>, ParseError> {
+    output: &mut Vec<(VarId, Span)>,
-    let mut output = vec![];
+) -> Result<(), ParseError> {
     for element in &pipeline.elements {
-        let result =
+        discover_captures_in_pipeline_element(working_set, element, seen, seen_blocks, output)?;
-            discover_captures_in_pipeline_element(working_set, element, seen, seen_blocks)?;
-        output.extend(&result);
     }
 
-    Ok(output)
+    Ok(())
 }
 
 // Closes over captured variables
@@ -5412,26 +5408,22 @@ pub fn discover_captures_in_pipeline_element(
     element: &PipelineElement,
     seen: &mut Vec<VarId>,
     seen_blocks: &mut HashMap<BlockId, Vec<(VarId, Span)>>,
-) -> Result<Vec<(VarId, Span)>, ParseError> {
+    output: &mut Vec<(VarId, Span)>,
+) -> Result<(), ParseError> {
     match element {
         PipelineElement::Expression(_, expression)
         | PipelineElement::Redirection(_, _, expression)
         | PipelineElement::And(_, expression)
         | PipelineElement::Or(_, expression) => {
-            discover_captures_in_expr(working_set, expression, seen, seen_blocks)
+            discover_captures_in_expr(working_set, expression, seen, seen_blocks, output)
         }
         PipelineElement::SeparateRedirection {
             out: (_, out_expr),
             err: (_, err_expr),
         } => {
-            let mut result = discover_captures_in_expr(working_set, out_expr, seen, seen_blocks)?;
+            discover_captures_in_expr(working_set, out_expr, seen, seen_blocks, output)?;
-            result.append(&mut discover_captures_in_expr(
+            discover_captures_in_expr(working_set, err_expr, seen, seen_blocks, output)?;
-                working_set,
+            Ok(())
-                err_expr,
-                seen,
-                seen_blocks,
-            )?);
-            Ok(result)
         }
     }
 }
@@ -5465,26 +5457,29 @@ pub fn discover_captures_in_expr(
     expr: &Expression,
     seen: &mut Vec<VarId>,
     seen_blocks: &mut HashMap<BlockId, Vec<(VarId, Span)>>,
-) -> Result<Vec<(VarId, Span)>, ParseError> {
+    output: &mut Vec<(VarId, Span)>,
-    let mut output: Vec<(VarId, Span)> = vec![];
+) -> Result<(), ParseError> {
     match &expr.expr {
         Expr::BinaryOp(lhs, _, rhs) => {
-            let lhs_result = discover_captures_in_expr(working_set, lhs, seen, seen_blocks)?;
+            discover_captures_in_expr(working_set, lhs, seen, seen_blocks, output)?;
-            let rhs_result = discover_captures_in_expr(working_set, rhs, seen, seen_blocks)?;
+            discover_captures_in_expr(working_set, rhs, seen, seen_blocks, output)?;
-
-            output.extend(&lhs_result);
-            output.extend(&rhs_result);
         }
         Expr::UnaryNot(expr) => {
-            let result = discover_captures_in_expr(working_set, expr, seen, seen_blocks)?;
+            discover_captures_in_expr(working_set, expr, seen, seen_blocks, output)?;
-            output.extend(&result);
         }
         Expr::Closure(block_id) => {
             let block = working_set.get_block(*block_id);
             let results = {
                 let mut seen = vec![];
-                let results =
+                let mut results = vec![];
-                    discover_captures_in_closure(working_set, block, &mut seen, seen_blocks)?;
+
+                discover_captures_in_closure(
+                    working_set,
+                    block,
+                    &mut seen,
+                    seen_blocks,
+                    &mut results,
+                )?;
 
                 for (var_id, span) in results.iter() {
                     if !seen.contains(var_id) {
@@ -5510,8 +5505,17 @@ pub fn discover_captures_in_expr(
             // FIXME: is this correct?
             let results = {
                 let mut seen = vec![];
-                discover_captures_in_closure(working_set, block, &mut seen, seen_blocks)?
+                let mut results = vec![];
+                discover_captures_in_closure(
+                    working_set,
+                    block,
+                    &mut seen,
+                    seen_blocks,
+                    &mut results,
+                )?;
+                results
             };
+
             seen_blocks.insert(*block_id, results.clone());
             for (var_id, span) in results.into_iter() {
                 if !seen.contains(&var_id) {
@@ -5536,11 +5540,13 @@ pub fn discover_captures_in_expr(
                             let mut seen = vec![];
                             seen_blocks.insert(block_id, output.clone());
 
-                            let result = discover_captures_in_closure(
+                            let mut result = vec![];
+                            discover_captures_in_closure(
                                 working_set,
                                 block,
                                 &mut seen,
                                 seen_blocks,
+                                &mut result,
                             )?;
                             output.extend(&result);
                             seen_blocks.insert(block_id, result);
@@ -5551,34 +5557,28 @@ pub fn discover_captures_in_expr(
 
             for named in call.named_iter() {
                 if let Some(arg) = &named.2 {
-                    let result = discover_captures_in_expr(working_set, arg, seen, seen_blocks)?;
+                    discover_captures_in_expr(working_set, arg, seen, seen_blocks, output)?;
-                    output.extend(&result);
                 }
             }
 
             for positional in call.positional_iter() {
-                let result = discover_captures_in_expr(working_set, positional, seen, seen_blocks)?;
+                discover_captures_in_expr(working_set, positional, seen, seen_blocks, output)?;
-                output.extend(&result);
             }
         }
         Expr::CellPath(_) => {}
         Expr::DateTime(_) => {}
         Expr::ExternalCall(head, exprs, _) => {
-            let result = discover_captures_in_expr(working_set, head, seen, seen_blocks)?;
+            discover_captures_in_expr(working_set, head, seen, seen_blocks, output)?;
-            output.extend(&result);
 
             for expr in exprs {
-                let result = discover_captures_in_expr(working_set, expr, seen, seen_blocks)?;
+                discover_captures_in_expr(working_set, expr, seen, seen_blocks, output)?;
-                output.extend(&result);
             }
         }
         Expr::Filepath(_) => {}
         Expr::Directory(_) => {}
         Expr::Float(_) => {}
         Expr::FullCellPath(cell_path) => {
-            let result =
+            discover_captures_in_expr(working_set, &cell_path.head, seen, seen_blocks, output)?;
-                discover_captures_in_expr(working_set, &cell_path.head, seen, seen_blocks)?;
-            output.extend(&result);
         }
         Expr::ImportPattern(_) => {}
         Expr::Overlay(_) => {}
@@ -5587,44 +5587,29 @@ pub fn discover_captures_in_expr(
         Expr::GlobPattern(_) => {}
         Expr::Int(_) => {}
         Expr::Keyword(_, _, expr) => {
-            let result = discover_captures_in_expr(working_set, expr, seen, seen_blocks)?;
+            discover_captures_in_expr(working_set, expr, seen, seen_blocks, output)?;
-            output.extend(&result);
         }
         Expr::List(exprs) => {
             for expr in exprs {
-                let result = discover_captures_in_expr(working_set, expr, seen, seen_blocks)?;
+                discover_captures_in_expr(working_set, expr, seen, seen_blocks, output)?;
-                output.extend(&result);
             }
         }
         Expr::Operator(_) => {}
         Expr::Range(expr1, expr2, expr3, _) => {
             if let Some(expr) = expr1 {
-                let result = discover_captures_in_expr(working_set, expr, seen, seen_blocks)?;
+                discover_captures_in_expr(working_set, expr, seen, seen_blocks, output)?;
-                output.extend(&result);
             }
             if let Some(expr) = expr2 {
-                let result = discover_captures_in_expr(working_set, expr, seen, seen_blocks)?;
+                discover_captures_in_expr(working_set, expr, seen, seen_blocks, output)?;
-                output.extend(&result);
             }
             if let Some(expr) = expr3 {
-                let result = discover_captures_in_expr(working_set, expr, seen, seen_blocks)?;
+                discover_captures_in_expr(working_set, expr, seen, seen_blocks, output)?;
-                output.extend(&result);
             }
         }
         Expr::Record(fields) => {
             for (field_name, field_value) in fields {
-                output.extend(&discover_captures_in_expr(
+                discover_captures_in_expr(working_set, field_name, seen, seen_blocks, output)?;
-                    working_set,
+                discover_captures_in_expr(working_set, field_value, seen, seen_blocks, output)?;
-                    field_name,
-                    seen,
-                    seen_blocks,
-                )?);
-                output.extend(&discover_captures_in_expr(
-                    working_set,
-                    field_value,
-                    seen,
-                    seen_blocks,
-                )?);
             }
         }
         Expr::Signature(sig) => {
@@ -5653,23 +5638,30 @@ pub fn discover_captures_in_expr(
         Expr::String(_) => {}
         Expr::StringInterpolation(exprs) => {
             for expr in exprs {
-                let result = discover_captures_in_expr(working_set, expr, seen, seen_blocks)?;
+                discover_captures_in_expr(working_set, expr, seen, seen_blocks, output)?;
-                output.extend(&result);
             }
         }
         Expr::MatchPattern(_) => {}
         Expr::MatchBlock(match_block) => {
             for match_ in match_block {
                 discover_captures_in_pattern(&match_.0, seen);
-                let result = discover_captures_in_expr(working_set, &match_.1, seen, seen_blocks)?;
+                discover_captures_in_expr(working_set, &match_.1, seen, seen_blocks, output)?;
-                output.extend(&result);
             }
         }
         Expr::RowCondition(block_id) | Expr::Subexpression(block_id) => {
             let block = working_set.get_block(*block_id);
+
             let results = {
+                let mut results = vec![];
                 let mut seen = vec![];
-                discover_captures_in_closure(working_set, block, &mut seen, seen_blocks)?
+                discover_captures_in_closure(
+                    working_set,
+                    block,
+                    &mut seen,
+                    seen_blocks,
+                    &mut results,
+                )?;
+                results
             };
             seen_blocks.insert(*block_id, results.clone());
             for (var_id, span) in results.into_iter() {
@@ -5680,19 +5672,16 @@ pub fn discover_captures_in_expr(
         }
         Expr::Table(headers, values) => {
             for header in headers {
-                let result = discover_captures_in_expr(working_set, header, seen, seen_blocks)?;
+                discover_captures_in_expr(working_set, header, seen, seen_blocks, output)?;
-                output.extend(&result);
             }
             for row in values {
                 for cell in row {
-                    let result = discover_captures_in_expr(working_set, cell, seen, seen_blocks)?;
+                    discover_captures_in_expr(working_set, cell, seen, seen_blocks, output)?;
-                    output.extend(&result);
                 }
             }
         }
         Expr::ValueWithUnit(expr, _) => {
-            let result = discover_captures_in_expr(working_set, expr, seen, seen_blocks)?;
+            discover_captures_in_expr(working_set, expr, seen, seen_blocks, output)?;
-            output.extend(&result);
         }
         Expr::Var(var_id) => {
             if (*var_id > ENV_VARIABLE_ID || *var_id == IN_VARIABLE_ID) && !seen.contains(var_id) {
@@ -5703,7 +5692,7 @@ pub fn discover_captures_in_expr(
             seen.push(*var_id);
         }
     }
-    Ok(output)
+    Ok(())
 }
 
 fn wrap_element_with_collect(
@@ -5833,9 +5822,15 @@ pub fn parse(
     let mut seen = vec![];
     let mut seen_blocks = HashMap::new();
 
-    let captures = discover_captures_in_closure(working_set, &output, &mut seen, &mut seen_blocks);
+    let mut captures = vec![];
-    match captures {
+    match discover_captures_in_closure(
-        Ok(captures) => output.captures = captures.into_iter().map(|(var_id, _)| var_id).collect(),
+        working_set,
+        &output,
+        &mut seen,
+        &mut seen_blocks,
+        &mut captures,
+    ) {
+        Ok(_) => output.captures = captures.into_iter().map(|(var_id, _)| var_id).collect(),
         Err(err) => working_set.error(err),
     }
 
@@ -5845,10 +5840,16 @@ pub fn parse(
         let block_id = block_idx + working_set.permanent_state.num_blocks();
 
         if !seen_blocks.contains_key(&block_id) {
-            let captures =
+            let mut captures = vec![];
-                discover_captures_in_closure(working_set, block, &mut seen, &mut seen_blocks);
+
-            match captures {
+            match discover_captures_in_closure(
-                Ok(captures) => {
+                working_set,
+                block,
+                &mut seen,
+                &mut seen_blocks,
+                &mut captures,
+            ) {
+                Ok(_) => {
                     seen_blocks.insert(block_id, captures);
                 }
                 Err(err) => {

src/tests.rs

@@ -48,6 +48,24 @@ pub fn run_test(input: &str, expected: &str) -> TestResult {
     let mut file = NamedTempFile::new()?;
     let name = file.path();
 
+    let mut cmd = Command::cargo_bin("nu")?;
+    cmd.arg("--no-std-lib");
+    cmd.arg(name);
+    cmd.env(
+        "PWD",
+        std::env::current_dir().expect("Can't get current dir"),
+    );
+
+    writeln!(file, "{input}")?;
+
+    run_cmd_and_assert(cmd, expected)
+}
+
+#[cfg(test)]
+pub fn run_test_std(input: &str, expected: &str) -> TestResult {
+    let mut file = NamedTempFile::new()?;
+    let name = file.path();
+
     let mut cmd = Command::cargo_bin("nu")?;
     cmd.arg(name);
     cmd.env(
@@ -83,6 +101,7 @@ pub fn run_test_contains(input: &str, expected: &str) -> TestResult {
     let name = file.path();
 
     let mut cmd = Command::cargo_bin("nu")?;
+    cmd.arg("--no-std-lib");
     cmd.arg(name);
 
     writeln!(file, "{input}")?;
@@ -108,6 +127,7 @@ pub fn test_ide_contains(input: &str, ide_commands: &[&str], expected: &str) ->
     let name = file.path();
 
     let mut cmd = Command::cargo_bin("nu")?;
+    cmd.arg("--no-std-lib");
     for ide_command in ide_commands {
         cmd.arg(ide_command);
     }
@@ -136,6 +156,7 @@ pub fn fail_test(input: &str, expected: &str) -> TestResult {
     let name = file.path();
 
     let mut cmd = Command::cargo_bin("nu")?;
+    cmd.arg("--no-std-lib");
     cmd.arg(name);
     cmd.env(
         "PWD",

src/tests/test_stdlib.rs

@@ -1,8 +1,8 @@
-use crate::tests::{fail_test, run_test, TestResult};
+use crate::tests::{fail_test, run_test_std, TestResult};
 
 #[test]
 fn library_loaded() -> TestResult {
-    run_test(
+    run_test_std(
         "help std | lines | first 1 | to text",
         "std.nu, `used` to load all standard library components",
     )
@@ -10,7 +10,7 @@ fn library_loaded() -> TestResult {
 
 #[test]
 fn prelude_loaded() -> TestResult {
-    run_test("std help commands | where name == open | length", "1")
+    run_test_std("std help commands | where name == open | length", "1")
 }
 
 #[test]
@@ -20,5 +20,5 @@ fn not_loaded() -> TestResult {
 
 #[test]
 fn use_command() -> TestResult {
-    run_test("use std assert; assert true; print 'it works'", "it works")
+    run_test_std("use std assert; assert true; print 'it works'", "it works")
 }