Fixing captures (#723)

* WIP fixing captures

* small fix

* WIP

* Rewrite to proof-of-concept better parse_def

* Add missing file

* Finish capture refactor

* Fix tests

* Add more tests

crates/nu-parser/src/flatten.rs

@@ -95,7 +95,7 @@ pub fn flatten_expression(
             output.extend(flatten_expression(working_set, rhs));
             output
         }
-        Expr::Block(block_id) => {
+        Expr::Block(block_id) | Expr::RowCondition(block_id) | Expr::Subexpression(block_id) => {
             let outer_span = expr.span;
 
             let mut output = vec![];
@@ -385,9 +385,6 @@ pub fn flatten_expression(
         Expr::String(_) => {
             vec![(expr.span, FlatShape::String)]
         }
-        Expr::RowCondition(block_id) | Expr::Subexpression(block_id) => {
-            flatten_block(working_set, working_set.get_block(*block_id))
-        }
         Expr::Table(headers, cells) => {
             let outer_span = expr.span;
             let mut last_end = outer_span.start;

crates/nu-parser/src/parse_keywords.rs

@@ -5,16 +5,16 @@ use nu_protocol::{
         Pipeline, Statement,
     },
     engine::StateWorkingSet,
-    span, Exportable, Overlay, Span, SyntaxShape, Type, CONFIG_VARIABLE_ID,
+    span, Exportable, Overlay, PositionalArg, Span, SyntaxShape, Type, CONFIG_VARIABLE_ID,
 };
 use std::collections::{HashMap, HashSet};
 
 use crate::{
     lex, lite_parse,
     parser::{
-        check_call, check_name, garbage, garbage_statement, parse, parse_block_expression,
-        parse_internal_call, parse_multispan_value, parse_signature, parse_string,
-        parse_var_with_opt_type, trim_quotes,
+        check_call, check_name, find_captures_in_block, garbage, garbage_statement, parse,
+        parse_block_expression, parse_internal_call, parse_multispan_value, parse_signature,
+        parse_string, parse_var_with_opt_type, trim_quotes,
     },
     ParseError,
 };
@@ -57,6 +57,121 @@ pub fn parse_def_predecl(working_set: &mut StateWorkingSet, spans: &[Span]) -> O
     None
 }
 
+pub fn parse_for(
+    working_set: &mut StateWorkingSet,
+    spans: &[Span],
+) -> (Statement, Option<ParseError>) {
+    // Checking that the function is used with the correct name
+    // Maybe this is not necessary but it is a sanity check
+    if working_set.get_span_contents(spans[0]) != b"for" {
+        return (
+            garbage_statement(spans),
+            Some(ParseError::UnknownState(
+                "internal error: Wrong call name for 'for' function".into(),
+                span(spans),
+            )),
+        );
+    }
+
+    // Parsing the spans and checking that they match the register signature
+    // Using a parsed call makes more sense than checking for how many spans are in the call
+    // Also, by creating a call, it can be checked if it matches the declaration signature
+    let (call, call_span) = match working_set.find_decl(b"for") {
+        None => {
+            return (
+                garbage_statement(spans),
+                Some(ParseError::UnknownState(
+                    "internal error: def declaration not found".into(),
+                    span(spans),
+                )),
+            )
+        }
+        Some(decl_id) => {
+            working_set.enter_scope();
+            let (call, mut err) = parse_internal_call(working_set, spans[0], &spans[1..], decl_id);
+            working_set.exit_scope();
+
+            let call_span = span(spans);
+            let decl = working_set.get_decl(decl_id);
+            let sig = decl.signature();
+
+            // Let's get our block and make sure it has the right signature
+            if let Some(arg) = call.positional.get(2) {
+                match arg {
+                    Expression {
+                        expr: Expr::Block(block_id),
+                        ..
+                    }
+                    | Expression {
+                        expr: Expr::RowCondition(block_id),
+                        ..
+                    } => {
+                        let block = working_set.get_block_mut(*block_id);
+
+                        block.signature = Box::new(sig.clone());
+                    }
+                    _ => {}
+                }
+            }
+
+            err = check_call(call_span, &sig, &call).or(err);
+            if err.is_some() || call.has_flag("help") {
+                return (
+                    Statement::Pipeline(Pipeline::from_vec(vec![Expression {
+                        expr: Expr::Call(call),
+                        span: call_span,
+                        ty: Type::Unknown,
+                        custom_completion: None,
+                    }])),
+                    err,
+                );
+            }
+
+            (call, call_span)
+        }
+    };
+
+    // All positional arguments must be in the call positional vector by this point
+    let var_decl = call.positional.get(0).expect("for call already checked");
+    let block = call.positional.get(2).expect("for call already checked");
+
+    let error = None;
+    if let (Some(var_id), Some(block_id)) = (&var_decl.as_var(), block.as_block()) {
+        let block = working_set.get_block_mut(block_id);
+
+        block.signature.required_positional.insert(
+            0,
+            PositionalArg {
+                name: String::new(),
+                desc: String::new(),
+                shape: SyntaxShape::Any,
+                var_id: Some(*var_id),
+            },
+        );
+
+        let block = working_set.get_block(block_id);
+
+        // Now that we have a signature for the block, we know more about what variables
+        // will come into scope as params. Because of this, we need to recalculated what
+        // variables this block will capture from the outside.
+        let mut seen = vec![];
+        let captures = find_captures_in_block(working_set, block, &mut seen);
+
+        let mut block = working_set.get_block_mut(block_id);
+        block.captures = captures;
+    }
+
+    (
+        Statement::Pipeline(Pipeline::from_vec(vec![Expression {
+            expr: Expr::Call(call),
+            span: call_span,
+            ty: Type::Unknown,
+            custom_completion: None,
+        }])),
+        error,
+    )
+}
+
 pub fn parse_def(
     working_set: &mut StateWorkingSet,
     spans: &[Span],
@@ -93,8 +208,28 @@ pub fn parse_def(
 
             let call_span = span(spans);
             let decl = working_set.get_decl(decl_id);
+            let sig = decl.signature();
 
-            err = check_call(call_span, &decl.signature(), &call).or(err);
+            // Let's get our block and make sure it has the right signature
+            if let Some(arg) = call.positional.get(2) {
+                match arg {
+                    Expression {
+                        expr: Expr::Block(block_id),
+                        ..
+                    }
+                    | Expression {
+                        expr: Expr::RowCondition(block_id),
+                        ..
+                    } => {
+                        let block = working_set.get_block_mut(*block_id);
+
+                        block.signature = Box::new(sig.clone());
+                    }
+                    _ => {}
+                }
+            }
+
+            err = check_call(call_span, &sig, &call).or(err);
             if err.is_some() || call.has_flag("help") {
                 return (
                     Statement::Pipeline(Pipeline::from_vec(vec![Expression {
@@ -124,7 +259,22 @@ pub fn parse_def(
             let declaration = working_set.get_decl_mut(decl_id);
 
             signature.name = name.clone();
-            *declaration = signature.into_block_command(block_id);
+
+            *declaration = signature.clone().into_block_command(block_id);
+
+            let mut block = working_set.get_block_mut(block_id);
+            block.signature = signature;
+
+            let block = working_set.get_block(block_id);
+
+            // Now that we have a signature for the block, we know more about what variables
+            // will come into scope as params. Because of this, we need to recalculated what
+            // variables this block will capture from the outside.
+            let mut seen = vec![];
+            let captures = find_captures_in_block(working_set, block, &mut seen);
+
+            let mut block = working_set.get_block_mut(block_id);
+            block.captures = captures;
         } else {
             error = error.or_else(|| {
                 Some(ParseError::InternalError(

crates/nu-parser/src/parser.rs

@@ -1,6 +1,6 @@
 use crate::{
     lex, lite_parse,
-    parse_keywords::parse_source,
+    parse_keywords::{parse_for, parse_source},
     type_check::{math_result_type, type_compatible},
     LiteBlock, ParseError, Token, TokenContents,
 };
@@ -13,7 +13,7 @@ use nu_protocol::{
     },
     engine::StateWorkingSet,
     span, Flag, PositionalArg, Signature, Span, Spanned, SyntaxShape, Type, Unit, VarId,
-    CONFIG_VARIABLE_ID,
+    CONFIG_VARIABLE_ID, ENV_VARIABLE_ID, IN_VARIABLE_ID,
 };
 
 use crate::parse_keywords::{
@@ -3402,6 +3402,7 @@ pub fn parse_statement(
     match name {
         b"def" => parse_def(working_set, spans),
         b"let" => parse_let(working_set, spans),
+        b"for" => parse_for(working_set, spans),
         b"alias" => parse_alias(working_set, spans),
         b"module" => parse_module(working_set, spans),
         b"use" => parse_use(working_set, spans),
@@ -3577,13 +3578,15 @@ pub fn parse_block(
     (block, error)
 }
 
-fn find_captures_in_block(
+pub fn find_captures_in_block(
     working_set: &StateWorkingSet,
     block: &Block,
     seen: &mut Vec<VarId>,
 ) -> Vec<VarId> {
     let mut output = vec![];
 
+    // println!("sig: {:#?}", block.signature);
+
     for flag in &block.signature.named {
         if let Some(var_id) = flag.var_id {
             seen.push(var_id);
@@ -3654,6 +3657,13 @@ pub fn find_captures_in_expr(
         }
         Expr::Bool(_) => {}
         Expr::Call(call) => {
+            let decl = working_set.get_decl(call.decl_id);
+            if let Some(block_id) = decl.get_block_id() {
+                let block = working_set.get_block(block_id);
+                let result = find_captures_in_block(working_set, block, seen);
+                output.extend(&result);
+            }
+
             for named in &call.named {
                 if let Some(arg) = &named.1 {
                     let result = find_captures_in_expr(working_set, arg, seen);
@@ -3715,7 +3725,30 @@ pub fn find_captures_in_expr(
                 output.extend(&find_captures_in_expr(working_set, field_value, seen));
             }
         }
-        Expr::Signature(_) => {}
+        Expr::Signature(sig) => {
+            // println!("Signature found! Adding var_ids");
+            // Something with a declaration, similar to a var decl, will introduce more VarIds into the stack at eval
+            for pos in &sig.required_positional {
+                if let Some(var_id) = pos.var_id {
+                    seen.push(var_id);
+                }
+            }
+            for pos in &sig.optional_positional {
+                if let Some(var_id) = pos.var_id {
+                    seen.push(var_id);
+                }
+            }
+            if let Some(rest) = &sig.rest_positional {
+                if let Some(var_id) = rest.var_id {
+                    seen.push(var_id);
+                }
+            }
+            for named in &sig.named {
+                if let Some(var_id) = named.var_id {
+                    seen.push(var_id);
+                }
+            }
+        }
         Expr::String(_) => {}
         Expr::StringInterpolation(exprs) => {
             for expr in exprs {
@@ -3745,7 +3778,7 @@ pub fn find_captures_in_expr(
             output.extend(&result);
         }
         Expr::Var(var_id) => {
-            if !seen.contains(var_id) {
+            if (*var_id > ENV_VARIABLE_ID || *var_id == IN_VARIABLE_ID) && !seen.contains(var_id) {
                 output.push(*var_id);
             }
         }