From 886ed5ab2d847ec446e959126c980f56384c28df Mon Sep 17 00:00:00 2001
From: JT <547158+jntrnr@users.noreply.github.com>
Date: Fri, 11 Feb 2022 07:37:10 -0500
Subject: [PATCH] Fix captures (#4421)

* Fix rowcondition and import captures

* Only check extra blocks if not yet seen
---
 crates/nu-parser/src/parser.rs                | 23 +++++++++++++++++--
 crates/nu-protocol/src/engine/engine_state.rs |  2 +-
 src/tests/test_parser.rs                      |  8 +++++++
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/crates/nu-parser/src/parser.rs b/crates/nu-parser/src/parser.rs
index 4fc6df69b..4825c686e 100644
--- a/crates/nu-parser/src/parser.rs
+++ b/crates/nu-parser/src/parser.rs
@@ -3944,8 +3944,16 @@ pub fn discover_captures_in_expr(
         }
         Expr::RowCondition(block_id) | Expr::Subexpression(block_id) => {
             let block = working_set.get_block(*block_id);
-            let result = discover_captures_in_block(working_set, block, seen, seen_blocks);
-            output.extend(&result);
+            let results = {
+                let mut seen = vec![];
+                discover_captures_in_block(working_set, block, &mut seen, seen_blocks)
+            };
+            seen_blocks.insert(*block_id, results.clone());
+            for var_id in results.into_iter() {
+                if !seen.contains(&var_id) {
+                    output.push(var_id)
+                }
+            }
         }
         Expr::Table(headers, values) => {
             for header in headers {
@@ -4066,6 +4074,17 @@ pub fn parse(
     let captures = discover_captures_in_block(working_set, &output, &mut seen, &mut seen_blocks);
     output.captures = captures;
 
+    // Also check other blocks that might have been imported
+    for (block_idx, block) in working_set.delta.blocks.iter().enumerate() {
+        let block_id = block_idx + working_set.permanent_state.num_blocks();
+
+        if !seen_blocks.contains_key(&block_id) {
+            let captures =
+                discover_captures_in_block(working_set, block, &mut seen, &mut seen_blocks);
+            seen_blocks.insert(block_id, captures);
+        }
+    }
+
     for (block_id, captures) in seen_blocks.into_iter() {
         // In theory, we should only be updating captures where we have new information
         // the only place where this is possible would be blocks that are newly created
diff --git a/crates/nu-protocol/src/engine/engine_state.rs b/crates/nu-protocol/src/engine/engine_state.rs
index 7377feaa4..fdc9ab8de 100644
--- a/crates/nu-protocol/src/engine/engine_state.rs
+++ b/crates/nu-protocol/src/engine/engine_state.rs
@@ -607,7 +607,7 @@ pub struct StateDelta {
     pub(crate) file_contents: Vec<(Vec<u8>, usize, usize)>,
     vars: Vec<Type>,              // indexed by VarId
     decls: Vec<Box<dyn Command>>, // indexed by DeclId
-    blocks: Vec<Block>,           // indexed by BlockId
+    pub blocks: Vec<Block>,       // indexed by BlockId
     overlays: Vec<Overlay>,       // indexed by OverlayId
     pub scope: Vec<ScopeFrame>,
     #[cfg(feature = "plugin")]
diff --git a/src/tests/test_parser.rs b/src/tests/test_parser.rs
index 4ecca20b3..d97526328 100644
--- a/src/tests/test_parser.rs
+++ b/src/tests/test_parser.rs
@@ -296,3 +296,11 @@ def 'say-hi' [] {
         "HelloHello",
     )
 }
+
+#[test]
+fn capture_row_condition() -> TestResult {
+    run_test(
+        r#"let name = "foo"; [foo] | where $'($name)' =~ $it | str collect"#,
+        "foo",
+    )
+}