Add magic $in variable (#309)

* Add magic in variable

* Oops, missing file

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

@@ -1,5 +1,5 @@
-use super::{Expr, Operator};
-use crate::{BlockId, Signature, Span, Type, VarId};
+use super::{Expr, Operator, Statement};
+use crate::{engine::StateWorkingSet, BlockId, Signature, Span, Type, VarId, IN_VARIABLE_ID};
 
 #[derive(Debug, Clone)]
 pub struct Expression {
@@ -88,4 +88,219 @@ impl Expression {
             _ => None,
         }
     }
+
+    pub fn has_in_variable(&self, working_set: &StateWorkingSet) -> bool {
+        match &self.expr {
+            Expr::BinaryOp(left, _, right) => {
+                left.has_in_variable(working_set) || right.has_in_variable(working_set)
+            }
+            Expr::Block(block_id) => {
+                let block = working_set.get_block(*block_id);
+
+                if let Some(Statement::Pipeline(pipeline)) = block.stmts.get(0) {
+                    match pipeline.expressions.get(0) {
+                        Some(expr) => expr.has_in_variable(working_set),
+                        None => false,
+                    }
+                } else {
+                    false
+                }
+            }
+            Expr::Bool(_) => false,
+            Expr::Call(call) => {
+                for positional in &call.positional {
+                    if positional.has_in_variable(working_set) {
+                        return true;
+                    }
+                }
+                for named in &call.named {
+                    if let Some(expr) = &named.1 {
+                        if expr.has_in_variable(working_set) {
+                            return true;
+                        }
+                    }
+                }
+                false
+            }
+            Expr::CellPath(_) => false,
+            Expr::ExternalCall(_, _, args) => {
+                for arg in args {
+                    if arg.has_in_variable(working_set) {
+                        return true;
+                    }
+                }
+                false
+            }
+            Expr::Filepath(_) => false,
+            Expr::Float(_) => false,
+            Expr::FullCellPath(full_cell_path) => {
+                if full_cell_path.head.has_in_variable(working_set) {
+                    return true;
+                }
+                false
+            }
+            Expr::Garbage => false,
+            Expr::GlobPattern(_) => false,
+            Expr::Int(_) => false,
+            Expr::Keyword(_, _, expr) => expr.has_in_variable(working_set),
+            Expr::List(list) => {
+                for l in list {
+                    if l.has_in_variable(working_set) {
+                        return true;
+                    }
+                }
+                false
+            }
+            Expr::Operator(_) => false,
+            Expr::Range(left, middle, right, ..) => {
+                if let Some(left) = &left {
+                    if left.has_in_variable(working_set) {
+                        return true;
+                    }
+                }
+                if let Some(middle) = &middle {
+                    if middle.has_in_variable(working_set) {
+                        return true;
+                    }
+                }
+                if let Some(right) = &right {
+                    if right.has_in_variable(working_set) {
+                        return true;
+                    }
+                }
+                false
+            }
+            Expr::RowCondition(_, expr) => expr.has_in_variable(working_set),
+            Expr::Signature(_) => false,
+            Expr::String(_) => false,
+            Expr::Subexpression(block_id) => {
+                let block = working_set.get_block(*block_id);
+
+                if let Some(Statement::Pipeline(pipeline)) = block.stmts.get(0) {
+                    if let Some(expr) = pipeline.expressions.get(0) {
+                        expr.has_in_variable(working_set)
+                    } else {
+                        false
+                    }
+                } else {
+                    false
+                }
+            }
+            Expr::Table(headers, cells) => {
+                for header in headers {
+                    if header.has_in_variable(working_set) {
+                        return true;
+                    }
+                }
+
+                for row in cells {
+                    for cell in row.iter() {
+                        if cell.has_in_variable(working_set) {
+                            return true;
+                        }
+                    }
+                }
+
+                false
+            }
+
+            Expr::ValueWithUnit(expr, _) => expr.has_in_variable(working_set),
+            Expr::Var(var_id) => *var_id == IN_VARIABLE_ID,
+            Expr::VarDecl(_) => false,
+        }
+    }
+
+    pub fn replace_in_variable(&mut self, working_set: &mut StateWorkingSet, new_var_id: VarId) {
+        match &mut self.expr {
+            Expr::BinaryOp(left, _, right) => {
+                left.replace_in_variable(working_set, new_var_id);
+                right.replace_in_variable(working_set, new_var_id);
+            }
+            Expr::Block(block_id) => {
+                let block = working_set.get_block_mut(*block_id);
+
+                if let Some(Statement::Pipeline(pipeline)) = block.stmts.get_mut(0) {
+                    if let Some(expr) = pipeline.expressions.get_mut(0) {
+                        expr.clone().replace_in_variable(working_set, new_var_id)
+                    }
+                }
+            }
+            Expr::Bool(_) => {}
+            Expr::Call(call) => {
+                for positional in &mut call.positional {
+                    positional.replace_in_variable(working_set, new_var_id);
+                }
+                for named in &mut call.named {
+                    if let Some(expr) = &mut named.1 {
+                        expr.replace_in_variable(working_set, new_var_id)
+                    }
+                }
+            }
+            Expr::CellPath(_) => {}
+            Expr::ExternalCall(_, _, args) => {
+                for arg in args {
+                    arg.replace_in_variable(working_set, new_var_id)
+                }
+            }
+            Expr::Filepath(_) => {}
+            Expr::Float(_) => {}
+            Expr::FullCellPath(full_cell_path) => {
+                full_cell_path
+                    .head
+                    .replace_in_variable(working_set, new_var_id);
+            }
+            Expr::Garbage => {}
+            Expr::GlobPattern(_) => {}
+            Expr::Int(_) => {}
+            Expr::Keyword(_, _, expr) => expr.replace_in_variable(working_set, new_var_id),
+            Expr::List(list) => {
+                for l in list {
+                    l.replace_in_variable(working_set, new_var_id)
+                }
+            }
+            Expr::Operator(_) => {}
+            Expr::Range(left, middle, right, ..) => {
+                if let Some(left) = left {
+                    left.replace_in_variable(working_set, new_var_id)
+                }
+                if let Some(middle) = middle {
+                    middle.replace_in_variable(working_set, new_var_id)
+                }
+                if let Some(right) = right {
+                    right.replace_in_variable(working_set, new_var_id)
+                }
+            }
+            Expr::RowCondition(_, expr) => expr.replace_in_variable(working_set, new_var_id),
+            Expr::Signature(_) => {}
+            Expr::String(_) => {}
+            Expr::Subexpression(block_id) => {
+                let block = working_set.get_block_mut(*block_id);
+
+                if let Some(Statement::Pipeline(pipeline)) = block.stmts.get_mut(0) {
+                    if let Some(expr) = pipeline.expressions.get_mut(0) {
+                        expr.clone().replace_in_variable(working_set, new_var_id)
+                    }
+                }
+            }
+            Expr::Table(headers, cells) => {
+                for header in headers {
+                    header.replace_in_variable(working_set, new_var_id)
+                }
+
+                for row in cells {
+                    for cell in row.iter_mut() {
+                        cell.replace_in_variable(working_set, new_var_id)
+                    }
+                }
+            }
+
+            Expr::ValueWithUnit(expr, _) => expr.replace_in_variable(working_set, new_var_id),
+            Expr::Var(x) => {
+                if *x == IN_VARIABLE_ID {
+                    *x = new_var_id
+                }
+            }
+            Expr::VarDecl(_) => {}
+        }
+    }
 }

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

@@ -135,13 +135,14 @@ pub struct EngineState {
 
 pub const NU_VARIABLE_ID: usize = 0;
 pub const SCOPE_VARIABLE_ID: usize = 1;
+pub const IN_VARIABLE_ID: usize = 2;
 
 impl EngineState {
     pub fn new() -> Self {
         Self {
             files: im::vector![],
             file_contents: im::vector![],
-            vars: im::vector![Type::Unknown, Type::Unknown],
+            vars: im::vector![Type::Unknown, Type::Unknown, Type::Unknown],
             decls: im::vector![],
             blocks: im::vector![],
             scope: im::vector![ScopeFrame::new()],
@@ -857,6 +858,18 @@ impl<'a> StateWorkingSet<'a> {
         }
     }
 
+    pub fn get_block_mut(&mut self, block_id: BlockId) -> &mut Block {
+        let num_permanent_blocks = self.permanent_state.num_blocks();
+        if block_id < num_permanent_blocks {
+            panic!("Attempt to mutate a block that is in the permanent (immutable) state")
+        } else {
+            self.delta
+                .blocks
+                .get_mut(block_id - num_permanent_blocks)
+                .expect("internal error: missing block")
+        }
+    }
+
     pub fn render(self) -> StateDelta {
         self.delta
     }

crates/nu-protocol/src/lib.rs

@@ -11,7 +11,7 @@ mod ty;
 mod value;
 pub use value::Value;
 
-pub use engine::{NU_VARIABLE_ID, SCOPE_VARIABLE_ID};
+pub use engine::{IN_VARIABLE_ID, NU_VARIABLE_ID, SCOPE_VARIABLE_ID};
 pub use example::*;
 pub use id::*;
 pub use pipeline_data::*;