WIP

crates/nu-command/src/core_commands/do_.rs

@@ -40,7 +40,7 @@ impl Command for Do {
 
         let block = engine_state.get_block(block_id);
 
-        let mut stack = stack.enter_scope();
+        let mut stack = stack.collect_captures(&block.captures);
 
         let params: Vec<_> = block
             .signature

crates/nu-command/src/core_commands/for_.rs

@@ -51,22 +51,22 @@ impl Command for For {
             .expect("internal error: missing keyword");
         let values = eval_expression(engine_state, stack, keyword_expr)?;
 
-        let block = call.positional[2]
+        let block_id = call.positional[2]
             .as_block()
             .expect("internal error: expected block");
 
         let engine_state = engine_state.clone();
-        let stack = stack.enter_scope();
+        let block = engine_state.get_block(block_id);
+        let mut stack = stack.collect_captures(&block.captures);
 
         match values {
             Value::List { vals, span } => Ok(vals
                 .into_iter()
                 .map(move |x| {
-                    let block = engine_state.get_block(block);
-
                     let mut stack = stack.clone();
                     stack.add_var(var_id, x);
 
+                    let block = engine_state.get_block(block_id);
                     match eval_block(&engine_state, &mut stack, block, PipelineData::new()) {
                         Ok(value) => Value::List {
                             vals: value.collect(),
@@ -79,12 +79,9 @@ impl Command for For {
             Value::Range { val, span } => Ok(val
                 .into_range_iter()?
                 .map(move |x| {
-                    let block = engine_state.get_block(block);
-
-                    let mut stack = stack.enter_scope();
-
                     stack.add_var(var_id, x);
 
+                    let block = engine_state.get_block(block_id);
                     match eval_block(&engine_state, &mut stack, block, PipelineData::new()) {
                         Ok(value) => Value::List {
                             vals: value.collect(),
@@ -95,10 +92,6 @@ impl Command for For {
                 })
                 .into_pipeline_data()),
             x => {
-                let block = engine_state.get_block(block);
-
-                let mut stack = stack.enter_scope();
-
                 stack.add_var(var_id, x);
 
                 eval_block(&engine_state, &mut stack, block, PipelineData::new())

crates/nu-command/src/core_commands/if_.rs

@@ -44,13 +44,13 @@ impl Command for If {
             Value::Bool { val, .. } => {
                 if val {
                     let block = engine_state.get_block(then_block);
-                    let mut stack = stack.enter_scope();
+                    let mut stack = stack.collect_captures(&block.captures);
                     eval_block(engine_state, &mut stack, block, input)
                 } else if let Some(else_case) = else_case {
                     if let Some(else_expr) = else_case.as_keyword() {
                         if let Some(block_id) = else_expr.as_block() {
                             let block = engine_state.get_block(block_id);
-                            let mut stack = stack.enter_scope();
+                            let mut stack = stack.collect_captures(&block.captures);
                             eval_block(engine_state, &mut stack, block, input)
                         } else {
                             eval_expression(engine_state, stack, else_expr)

crates/nu-command/src/filters/each.rs

@@ -64,7 +64,8 @@ impl Command for Each {
 
         let numbered = call.has_flag("numbered");
         let engine_state = engine_state.clone();
-        let stack = stack.clone();
+        let block = engine_state.get_block(block_id);
+        let mut stack = stack.collect_captures(&block.captures);
         let span = call.head;
 
         match input {
@@ -74,7 +75,7 @@ impl Command for Each {
                 .map(move |(idx, x)| {
                     let block = engine_state.get_block(block_id);
 
-                    let mut stack = stack.enter_scope();
+                    let mut stack = stack.clone();
 
                     if let Some(var) = block.signature.get_positional(0) {
                         if let Some(var_id) = &var.var_id {
@@ -112,7 +113,8 @@ impl Command for Each {
                 .map(move |(idx, x)| {
                     let block = engine_state.get_block(block_id);
 
-                    let mut stack = stack.enter_scope();
+                    let mut stack = stack.clone();
+
                     if let Some(var) = block.signature.get_positional(0) {
                         if let Some(var_id) = &var.var_id {
                             if numbered {
@@ -148,7 +150,8 @@ impl Command for Each {
                 .map(move |(idx, x)| {
                     let block = engine_state.get_block(block_id);
 
-                    let mut stack = stack.enter_scope();
+                    let mut stack = stack.clone();
+
                     if let Some(var) = block.signature.get_positional(0) {
                         if let Some(var_id) = &var.var_id {
                             if numbered {
@@ -186,7 +189,8 @@ impl Command for Each {
                 for (col, val) in cols.into_iter().zip(vals.into_iter()) {
                     let block = engine_state.get_block(block_id);
 
-                    let mut stack = stack.enter_scope();
+                    let mut stack = stack.clone();
+
                     if let Some(var) = block.signature.get_positional(0) {
                         if let Some(var_id) = &var.var_id {
                             stack.add_var(
@@ -231,7 +235,6 @@ impl Command for Each {
             PipelineData::Value(x) => {
                 let block = engine_state.get_block(block_id);
 
-                let mut stack = stack.enter_scope();
                 if let Some(var) = block.signature.get_positional(0) {
                     if let Some(var_id) = &var.var_id {
                         stack.add_var(*var_id, x);

crates/nu-command/src/filters/where_.rs

@@ -29,7 +29,9 @@ impl Command for Where {
         let cond = call.positional[0].clone();
 
         let engine_state = engine_state.clone();
-        let mut stack = stack.enter_scope();
+
+        // FIXME: very expensive
+        let mut stack = stack.clone();
 
         let (var_id, cond) = match cond {
             Expression {

crates/nu-command/src/system/benchmark.rs

@@ -37,7 +37,7 @@ impl Command for Benchmark {
             .expect("internal error: expected block");
         let block = engine_state.get_block(block);
 
-        let mut stack = stack.enter_scope();
+        let mut stack = stack.collect_captures(&block.captures);
         let start_time = Instant::now();
         eval_block(engine_state, &mut stack, block, PipelineData::new())?.into_value();
 

crates/nu-engine/src/eval.rs

@@ -33,7 +33,9 @@ fn eval_call(
         }
         .into_pipeline_data())
     } else if let Some(block_id) = decl.get_block_id() {
-        let mut stack = stack.enter_scope();
+        let block = engine_state.get_block(block_id);
+
+        let mut stack = stack.collect_captures(&block.captures);
         for (arg, param) in call.positional.iter().zip(
             decl.signature()
                 .required_positional
@@ -109,7 +111,6 @@ fn eval_call(
                 }
             }
         }
-        let block = engine_state.get_block(block_id);
         eval_block(engine_state, &mut stack, block, input)
     } else {
         decl.run(engine_state, stack, call, input)
@@ -212,6 +213,7 @@ pub fn eval_expression(
         Expr::Var(var_id) => stack
             .get_var(*var_id)
             .map_err(move |_| ShellError::VariableNotFoundAtRuntime(expr.span)),
+        Expr::VarDecl(_) => Ok(Value::Nothing { span: expr.span }),
         Expr::CellPath(cell_path) => Ok(Value::CellPath {
             val: cell_path.clone(),
             span: expr.span,
@@ -276,7 +278,7 @@ pub fn eval_expression(
         Expr::Subexpression(block_id) => {
             let block = engine_state.get_block(*block_id);
 
-            let mut stack = stack.enter_scope();
+            let mut stack = stack.collect_captures(&block.captures);
 
             // FIXME: protect this collect with ctrl-c
             Ok(Value::List {

crates/nu-parser/src/flatten.rs

@@ -194,7 +194,7 @@ pub fn flatten_expression(
             }
             output
         }
-        Expr::Var(_) => {
+        Expr::Var(_) | Expr::VarDecl(_) => {
             vec![(expr.span, FlatShape::Variable)]
         }
     }

crates/nu-parser/src/lib.rs

@@ -13,4 +13,4 @@ pub use lite_parse::{lite_parse, LiteBlock};
 pub use parse_keywords::{
     parse_alias, parse_def, parse_def_predecl, parse_let, parse_module, parse_use,
 };
-pub use parser::{parse, Import, VarDecl};
+pub use parser::{find_captures_in_expr, parse, Import, VarDecl};

crates/nu-parser/src/parser.rs

@@ -1806,7 +1806,7 @@ pub fn parse_var_with_opt_type(
 
             (
                 Expression {
-                    expr: Expr::Var(id),
+                    expr: Expr::VarDecl(id),
                     span: span(&spans[*spans_idx - 1..*spans_idx + 1]),
                     ty,
                     custom_completion: None,
@@ -1817,7 +1817,7 @@ pub fn parse_var_with_opt_type(
             let id = working_set.add_variable(bytes[0..(bytes.len() - 1)].to_vec(), Type::Unknown);
             (
                 Expression {
-                    expr: Expr::Var(id),
+                    expr: Expr::VarDecl(id),
                     span: spans[*spans_idx],
                     ty: Type::Unknown,
                     custom_completion: None,
@@ -1830,7 +1830,7 @@ pub fn parse_var_with_opt_type(
 
         (
             Expression {
-                expr: Expr::Var(id),
+                expr: Expr::VarDecl(id),
                 span: span(&spans[*spans_idx..*spans_idx + 1]),
                 ty: Type::Unknown,
                 custom_completion: None,
@@ -2574,6 +2574,11 @@ pub fn parse_block_expression(
         }
     }
 
+    let mut seen = vec![];
+    let captures = find_captures_in_block(working_set, &output, &mut seen);
+
+    output.captures = captures;
+
     working_set.exit_scope();
 
     let block_id = working_set.add_block(output);
@@ -3041,6 +3046,177 @@ pub fn parse_block(
     (block, error)
 }
 
+fn find_captures_in_block(
+    working_set: &StateWorkingSet,
+    block: &Block,
+    seen: &mut Vec<VarId>,
+) -> Vec<VarId> {
+    let mut output = vec![];
+
+    for flag in &block.signature.named {
+        if let Some(var_id) = flag.var_id {
+            seen.push(var_id);
+        }
+    }
+
+    for positional in &block.signature.required_positional {
+        if let Some(var_id) = positional.var_id {
+            seen.push(var_id);
+        }
+    }
+    for positional in &block.signature.optional_positional {
+        if let Some(var_id) = positional.var_id {
+            seen.push(var_id);
+        }
+    }
+    for positional in &block.signature.rest_positional {
+        if let Some(var_id) = positional.var_id {
+            seen.push(var_id);
+        }
+    }
+
+    for stmt in &block.stmts {
+        match stmt {
+            Statement::Pipeline(pipeline) => {
+                let result = find_captures_in_pipeline(working_set, pipeline, seen);
+                output.extend(&result);
+            }
+            Statement::Declaration(_) => {}
+        }
+    }
+
+    output
+}
+
+fn find_captures_in_pipeline(
+    working_set: &StateWorkingSet,
+    pipeline: &Pipeline,
+    seen: &mut Vec<VarId>,
+) -> Vec<VarId> {
+    let mut output = vec![];
+    for expr in &pipeline.expressions {
+        let result = find_captures_in_expr(working_set, expr, seen);
+        output.extend(&result);
+    }
+
+    output
+}
+
+pub fn find_captures_in_expr(
+    working_set: &StateWorkingSet,
+    expr: &Expression,
+    seen: &mut Vec<VarId>,
+) -> Vec<VarId> {
+    let mut output = vec![];
+    match &expr.expr {
+        Expr::BinaryOp(lhs, _, rhs) => {
+            let lhs_result = find_captures_in_expr(working_set, lhs, seen);
+            let rhs_result = find_captures_in_expr(working_set, rhs, seen);
+
+            output.extend(&lhs_result);
+            output.extend(&rhs_result);
+        }
+        Expr::Block(block_id) => {
+            let block = working_set.get_block(*block_id);
+            let result = find_captures_in_block(working_set, block, seen);
+            output.extend(&result);
+        }
+        Expr::Bool(_) => {}
+        Expr::Call(call) => {
+            for named in &call.named {
+                if let Some(arg) = &named.1 {
+                    let result = find_captures_in_expr(working_set, arg, seen);
+                    output.extend(&result);
+                }
+            }
+
+            for positional in &call.positional {
+                let result = find_captures_in_expr(working_set, positional, seen);
+                output.extend(&result);
+            }
+        }
+        Expr::CellPath(_) => {}
+        Expr::ExternalCall(_, _, exprs) => {
+            for expr in exprs {
+                let result = find_captures_in_expr(working_set, expr, seen);
+                output.extend(&result);
+            }
+        }
+        Expr::Filepath(_) => {}
+        Expr::Float(_) => {}
+        Expr::FullCellPath(cell_path) => {
+            let result = find_captures_in_expr(working_set, &cell_path.head, seen);
+            output.extend(&result);
+        }
+        Expr::Garbage => {}
+        Expr::GlobPattern(_) => {}
+        Expr::Int(_) => {}
+        Expr::Keyword(_, _, expr) => {
+            let result = find_captures_in_expr(working_set, expr, seen);
+            output.extend(&result);
+        }
+        Expr::List(exprs) => {
+            for expr in exprs {
+                let result = find_captures_in_expr(working_set, expr, seen);
+                output.extend(&result);
+            }
+        }
+        Expr::Operator(_) => {}
+        Expr::Range(expr1, expr2, expr3, _) => {
+            if let Some(expr) = expr1 {
+                let result = find_captures_in_expr(working_set, expr, seen);
+                output.extend(&result);
+            }
+            if let Some(expr) = expr2 {
+                let result = find_captures_in_expr(working_set, expr, seen);
+                output.extend(&result);
+            }
+            if let Some(expr) = expr3 {
+                let result = find_captures_in_expr(working_set, expr, seen);
+                output.extend(&result);
+            }
+        }
+        Expr::RowCondition(var_id, expr) => {
+            seen.push(*var_id);
+
+            let result = find_captures_in_expr(working_set, expr, seen);
+            output.extend(&result);
+        }
+        Expr::Signature(_) => {}
+        Expr::String(_) => {}
+        Expr::Subexpression(block_id) => {
+            let block = working_set.get_block(*block_id);
+            let result = find_captures_in_block(working_set, block, seen);
+            output.extend(&result);
+        }
+        Expr::Table(headers, values) => {
+            for header in headers {
+                let result = find_captures_in_expr(working_set, header, seen);
+                output.extend(&result);
+            }
+            for row in values {
+                for cell in row {
+                    let result = find_captures_in_expr(working_set, cell, seen);
+                    output.extend(&result);
+                }
+            }
+        }
+        Expr::ValueWithUnit(expr, _) => {
+            let result = find_captures_in_expr(working_set, expr, seen);
+            output.extend(&result);
+        }
+        Expr::Var(var_id) => {
+            if !seen.contains(var_id) {
+                output.push(*var_id);
+            }
+        }
+        Expr::VarDecl(var_id) => {
+            seen.push(*var_id);
+        }
+    }
+    output
+}
+
 // Parses a vector of u8 to create an AST Block. If a file name is given, then
 // the name is stored in the working set. When parsing a source without a file
 // name, the source of bytes is stored as "source"

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

@@ -1,6 +1,6 @@
 use std::ops::{Index, IndexMut};
 
-use crate::{DeclId, Signature};
+use crate::{DeclId, Signature, VarId};
 
 use super::Statement;
 
@@ -9,6 +9,7 @@ pub struct Block {
     pub signature: Box<Signature>,
     pub stmts: Vec<Statement>,
     pub exports: Vec<(Vec<u8>, DeclId)>, // Assuming just defs for now
+    pub captures: Vec<VarId>,
 }
 
 impl Block {
@@ -47,6 +48,7 @@ impl Block {
             signature: Box::new(Signature::new("")),
             stmts: vec![],
             exports: vec![],
+            captures: vec![],
         }
     }
 
@@ -55,6 +57,7 @@ impl Block {
             signature: self.signature,
             stmts: self.stmts,
             exports,
+            captures: self.captures,
         }
     }
 }
@@ -68,6 +71,7 @@ where
             signature: Box::new(Signature::new("")),
             stmts: stmts.collect(),
             exports: vec![],
+            captures: vec![],
         }
     }
 }

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

@@ -13,6 +13,7 @@ pub enum Expr {
         RangeOperator,
     ),
     Var(VarId),
+    VarDecl(VarId),
     Call(Box<Call>),
     ExternalCall(String, Span, Vec<Expression>),
     Operator(Operator),

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

@@ -77,6 +77,7 @@ impl Expression {
     pub fn as_var(&self) -> Option<VarId> {
         match self.expr {
             Expr::Var(var_id) => Some(var_id),
+            Expr::VarDecl(var_id) => Some(var_id),
             _ => None,
         }
     }

crates/nu-protocol/src/engine/stack.rs

@@ -3,14 +3,11 @@ use std::collections::HashMap;
 use crate::{ShellError, Value, VarId};
 
 #[derive(Debug, Clone)]
-pub struct StackFrame {
+pub struct Stack {
     pub vars: HashMap<VarId, Value>,
     pub env_vars: HashMap<String, String>,
 }
 
-#[derive(Clone, Debug)]
-pub struct Stack(Vec<StackFrame>);
-
 impl Default for Stack {
     fn default() -> Self {
         Self::new()
@@ -19,77 +16,82 @@ impl Default for Stack {
 
 impl Stack {
     pub fn new() -> Stack {
-        Stack(vec![StackFrame {
+        Stack {
             vars: HashMap::new(),
             env_vars: HashMap::new(),
-        }])
+        }
     }
     pub fn get_var(&self, var_id: VarId) -> Result<Value, ShellError> {
-        for frame in self.0.iter().rev() {
+        if let Some(v) = self.vars.get(&var_id) {
-            if let Some(v) = frame.vars.get(&var_id) {
+            return Ok(v.clone());
-                return Ok(v.clone());
-            }
         }
         Err(ShellError::InternalError("variable not found".into()))
     }
 
     pub fn add_var(&mut self, var_id: VarId, value: Value) {
-        let frame = self
+        self.vars.insert(var_id, value);
-            .0
-            .last_mut()
-            .expect("internal error: can't access stack frame");
-        frame.vars.insert(var_id, value);
     }
 
     pub fn add_env_var(&mut self, var: String, value: String) {
-        let frame = self
+        self.env_vars.insert(var, value);
-            .0
-            .last_mut()
-            .expect("internal error: can't access stack frame");
-        frame.env_vars.insert(var, value);
     }
 
-    pub fn enter_scope(&self) -> Stack {
+    pub fn collect_captures(&self, captures: &[VarId]) -> Stack {
-        // FIXME: VERY EXPENSIVE to clone entire stack
+        let mut output = Stack::new();
-        let mut output = self.clone();
-        output.0.push(StackFrame {
-            vars: HashMap::new(),
-            env_vars: HashMap::new(),
-        });
 
-        output
+        for capture in captures {
-    }
+            output.vars.insert(
-
+                *capture,
-    pub fn get_env_vars(&self) -> HashMap<String, String> {
+                self.get_var(*capture)
-        let mut output = HashMap::new();
+                    .expect("internal error: capture of missing variable"),
-
+            );
-        for frame in &self.0 {
-            output.extend(frame.env_vars.clone().into_iter());
         }
 
         output
     }
 
+    // pub fn enter_scope(&self) -> Stack {
+    //     // FIXME: VERY EXPENSIVE to clone entire stack
+    //     let mut output = self.clone();
+    //     output.0.push(StackFrame {
+    //         vars: HashMap::new(),
+    //         env_vars: HashMap::new(),
+    //     });
+
+    //     output
+    // }
+
+    pub fn get_env_vars(&self) -> HashMap<String, String> {
+        // let mut output = HashMap::new();
+
+        // for frame in &self.0 {
+        //     output.extend(frame.env_vars.clone().into_iter());
+        // }
+
+        // output
+        self.env_vars.clone()
+    }
+
     pub fn get_env_var(&self, name: &str) -> Option<String> {
-        for frame in self.0.iter().rev() {
+        // for frame in self.0.iter().rev() {
-            if let Some(v) = frame.env_vars.get(name) {
+        if let Some(v) = self.env_vars.get(name) {
-                return Some(v.to_string());
+            return Some(v.to_string());
-            }
         }
+        // }
         None
     }
 
     pub fn print_stack(&self) {
-        for frame in self.0.iter().rev() {
+        // for frame in self.0.iter().rev() {
-            println!("===frame===");
+        // println!("===frame===");
-            println!("vars:");
+        println!("vars:");
-            for (var, val) in &frame.vars {
+        for (var, val) in &self.vars {
-                println!("  {}: {:?}", var, val);
+            println!("  {}: {:?}", var, val);
-            }
-            println!("env vars:");
-            for (var, val) in &frame.env_vars {
-                println!("  {}: {:?}", var, val);
-            }
         }
+        println!("env vars:");
+        for (var, val) in &self.env_vars {
+            println!("  {}: {:?}", var, val);
+        }
+        // }
     }
 }