From 5c419a39e84f7d7bfbd685d87b6c6e19c8788b5f Mon Sep 17 00:00:00 2001 From: Devyn Cairns Date: Wed, 3 Jul 2024 20:15:00 -0700 Subject: [PATCH] custom commands, fix let, fix bare $env var --- crates/nu-engine/src/compile/builder.rs | 4 + crates/nu-engine/src/compile/call.rs | 209 +++++++++++++++++++++++- crates/nu-engine/src/compile/keyword.rs | 2 +- crates/nu-engine/src/eval_ir.rs | 167 ++++++++++++++----- crates/nu-protocol/src/ir/display.rs | 15 ++ crates/nu-protocol/src/ir/mod.rs | 20 ++- 6 files changed, 368 insertions(+), 49 deletions(-) diff --git a/crates/nu-engine/src/compile/builder.rs b/crates/nu-engine/src/compile/builder.rs index e0f1e3acfc..beff2bba2d 100644 --- a/crates/nu-engine/src/compile/builder.rs +++ b/crates/nu-engine/src/compile/builder.rs @@ -136,6 +136,7 @@ impl BlockBuilder { | Literal::Nothing => (), } } + Instruction::LoadValue { dst, val: _ } => self.mark_register(*dst)?, Instruction::Move { dst, src } => { self.free_register(*src)?; self.mark_register(*dst)?; @@ -149,6 +150,9 @@ impl BlockBuilder { Instruction::LoadEnv { dst, key: _ } => self.mark_register(*dst)?, Instruction::LoadEnvOpt { dst, key: _ } => self.mark_register(*dst)?, Instruction::StoreEnv { key: _, src } => self.free_register(*src)?, + Instruction::NewCalleeStack => (), + Instruction::CaptureVariable { var_id: _ } => (), + Instruction::PushVariable { var_id: _, src } => self.free_register(*src)?, Instruction::PushPositional { src } => self.free_register(*src)?, Instruction::AppendRest { src } => self.free_register(*src)?, Instruction::PushFlag { name: _ } => (), diff --git a/crates/nu-engine/src/compile/call.rs b/crates/nu-engine/src/compile/call.rs index 4a4568a2dc..5967aebf05 100644 --- a/crates/nu-engine/src/compile/call.rs +++ b/crates/nu-engine/src/compile/call.rs @@ -1,7 +1,9 @@ +use std::iter::repeat; + use nu_protocol::{ - ast::{Argument, Call, Expression, ExternalArgument}, + ast::{Argument, Block, Call, Expression, ExternalArgument}, engine::StateWorkingSet, - ir::Instruction, + ir::{Instruction, Literal}, IntoSpanned, RegId, Span, }; @@ -53,6 +55,20 @@ pub(crate) fn compile_call( } } + // Also, if this is a custom command (block) call, we should handle this completely differently, + // setting args in variables on a callee stack. + if let Some(block_id) = decl.block_id() { + let block = working_set.get_block(block_id); + return compile_custom_command_call( + working_set, + builder, + block, + call, + redirect_modes, + io_reg, + ); + } + // It's important that we evaluate the args first before trying to set up the argument // state for the call. // @@ -149,6 +165,195 @@ pub(crate) fn compile_call( Ok(()) } +pub(crate) fn compile_custom_command_call( + working_set: &StateWorkingSet<'_>, + builder: &mut BlockBuilder, + decl_block: &Block, + call: &Call, + redirect_modes: RedirectModes, + io_reg: RegId, +) -> Result<(), CompileError> { + // Custom command calls take place on a new stack + builder.push(Instruction::NewCalleeStack.into_spanned(call.head))?; + + // Add captures + for &var_id in &decl_block.captures { + builder.push(Instruction::CaptureVariable { var_id }.into_spanned(call.head))?; + } + + // Add positional arguments + let positional_iter = decl_block + .signature + .required_positional + .iter() + .chain(&decl_block.signature.optional_positional) + .zip(call.positional_iter().map(Some).chain(repeat(None))); + + for (declared_positional, provided_positional) in positional_iter { + let var_id = declared_positional + .var_id + .expect("custom command positional parameter missing var_id"); + + let var_reg = builder.next_register()?; + let span = provided_positional.map(|b| b.span).unwrap_or(call.head); + if let Some(expr) = provided_positional { + // If the arg was provided, compile it + compile_expression( + working_set, + builder, + expr, + redirect_modes.with_capture_out(span), + None, + var_reg, + )?; + } else { + if let Some(default) = declared_positional.default_value.clone() { + // Put the default as a Value if it's present + builder.push( + Instruction::LoadValue { + dst: var_reg, + val: Box::new(default), + } + .into_spanned(span), + )?; + } else { + // It doesn't really matter if this is a required argument, because we're supposed + // to check that in the parser + builder.load_literal(var_reg, Literal::Nothing.into_spanned(span))?; + } + } + builder.push( + Instruction::PushVariable { + var_id, + src: var_reg, + } + .into_spanned(span), + )?; + } + + // Add rest arguments + if let Some(declared_rest) = &decl_block.signature.rest_positional { + let var_id = declared_rest + .var_id + .expect("custom command rest parameter missing var_id"); + + let rest_index = decl_block.signature.required_positional.len() + + decl_block.signature.optional_positional.len(); + + // Use the first span of rest if possible + let span = call + .rest_iter(rest_index) + .next() + .map(|(e, _)| e.span) + .unwrap_or(call.head); + + // Build the rest list from the remaining arguments + let var_reg = builder.literal(Literal::List { capacity: 0 }.into_spanned(span))?; + + for (expr, spread) in call.rest_iter(rest_index) { + let expr_reg = builder.next_register()?; + compile_expression( + working_set, + builder, + expr, + redirect_modes.with_capture_out(expr.span), + None, + expr_reg, + )?; + builder.push( + // If the original argument was a spread argument, spread it into the list + if spread { + Instruction::ListSpread { + src_dst: var_reg, + items: var_reg, + } + } else { + Instruction::ListPush { + src_dst: var_reg, + item: expr_reg, + } + } + .into_spanned(expr.span), + )?; + } + + builder.push( + Instruction::PushVariable { + var_id, + src: var_reg, + } + .into_spanned(span), + )?; + } + + // Add named arguments + for flag in &decl_block.signature.named { + if let Some(var_id) = flag.var_id { + let var_reg = builder.next_register()?; + + let provided_arg = call + .named_iter() + .find(|(name, _, _)| name.item == flag.long); + + let span = provided_arg + .map(|(name, _, _)| name.span) + .unwrap_or(call.head); + + if let Some(expr) = provided_arg.and_then(|(_, _, expr)| expr.as_ref()) { + // It was provided - compile it + compile_expression( + working_set, + builder, + expr, + redirect_modes.with_capture_out(expr.span), + None, + var_reg, + )?; + } else { + // It wasn't provided - use the default or `true` + if let Some(default) = flag.default_value.clone() { + builder.push( + Instruction::LoadValue { + dst: var_reg, + val: Box::new(default), + } + .into_spanned(span), + )?; + } else { + builder.load_literal(var_reg, Literal::Bool(true).into_spanned(span))?; + } + } + + builder.push( + Instruction::PushVariable { + var_id, + src: var_reg, + } + .into_spanned(span), + )?; + } + } + + if let Some(mode) = redirect_modes.out { + builder.push(mode.map(|mode| Instruction::RedirectOut { mode }))?; + } + + if let Some(mode) = redirect_modes.err { + builder.push(mode.map(|mode| Instruction::RedirectErr { mode }))?; + } + + // The state is set up, so we can do the call into io_reg + builder.push( + Instruction::Call { + decl_id: call.decl_id, + src_dst: io_reg, + } + .into_spanned(call.head), + )?; + + Ok(()) +} + pub(crate) fn compile_external_call( working_set: &StateWorkingSet, builder: &mut BlockBuilder, diff --git a/crates/nu-engine/src/compile/keyword.rs b/crates/nu-engine/src/compile/keyword.rs index ce3d79caf1..c12dbd9ca6 100644 --- a/crates/nu-engine/src/compile/keyword.rs +++ b/crates/nu-engine/src/compile/keyword.rs @@ -307,7 +307,7 @@ pub(crate) fn compile_let( builder, block, redirect_modes.with_capture_out(call.head), - None, + Some(io_reg), io_reg, )?; diff --git a/crates/nu-engine/src/eval_ir.rs b/crates/nu-engine/src/eval_ir.rs index dfc381a0b5..932099b560 100644 --- a/crates/nu-engine/src/eval_ir.rs +++ b/crates/nu-engine/src/eval_ir.rs @@ -7,10 +7,10 @@ use nu_protocol::{ engine::{Argument, Closure, EngineState, ErrorHandler, Matcher, Redirection, Stack}, ir::{Call, DataSlice, Instruction, IrBlock, Literal, RedirectMode}, record, DeclId, IntoPipelineData, IntoSpanned, ListStream, OutDest, PipelineData, Range, - Record, RegId, ShellError, Span, Spanned, Value, VarId, + Record, RegId, ShellError, Span, Spanned, Value, VarId, ENV_VARIABLE_ID, }; -use crate::eval::is_automatic_env_var; +use crate::{eval::is_automatic_env_var, eval_block_with_early_return}; /// Evaluate the compiled representation of a [`Block`]. pub fn eval_ir_block( @@ -35,6 +35,7 @@ pub fn eval_ir_block( data: &ir_block.data, args_base, error_handler_base, + callee_stack: None, redirect_out: None, redirect_err: None, matches: vec![], @@ -73,6 +74,8 @@ struct EvalContext<'a> { args_base: usize, /// Base index on the error handler stack to reset to after a call error_handler_base: usize, + /// Stack to use for callee + callee_stack: Option>, /// State set by redirect-out redirect_out: Option, /// State set by redirect-err @@ -108,6 +111,11 @@ impl<'a> EvalContext<'a> { span: Some(error_span), }) } + + /// Get the current stack to be used for the callee - either `callee_stack` if set, or `stack` + fn callee_stack(&mut self) -> &mut Stack { + self.callee_stack.as_deref_mut().unwrap_or(self.stack) + } } /// Eval an IR block on the provided slice of registers. @@ -131,7 +139,7 @@ fn eval_ir_block_impl( "{pc:-4}: {}", instruction.display(ctx.engine_state, ctx.data) ); - match eval_instruction(ctx, instruction, span) { + match eval_instruction::(ctx, instruction, span) { Ok(InstructionResult::Continue) => { pc += 1; } @@ -193,7 +201,7 @@ enum InstructionResult { } /// Perform an instruction -fn eval_instruction( +fn eval_instruction( ctx: &mut EvalContext<'_>, instruction: &Instruction, span: &Span, @@ -202,6 +210,10 @@ fn eval_instruction( match instruction { Instruction::LoadLiteral { dst, lit } => load_literal(ctx, *dst, lit, *span), + Instruction::LoadValue { dst, val } => { + ctx.put_reg(*dst, Value::clone(val).into_pipeline_data()); + Ok(Continue) + } Instruction::Move { dst, src } => { let val = ctx.take_reg(*src); ctx.put_reg(*dst, val); @@ -290,23 +302,57 @@ fn eval_instruction( }) } } + Instruction::NewCalleeStack => { + let new_stack = ctx.stack.gather_captures(ctx.engine_state, &[]); + ctx.callee_stack = Some(Box::new(new_stack)); + Ok(Continue) + } + Instruction::CaptureVariable { var_id } => { + if let Some(callee_stack) = &mut ctx.callee_stack { + let value = ctx.stack.get_var_with_origin(*var_id, *span)?; + callee_stack.add_var(*var_id, value); + Ok(Continue) + } else { + Err(ShellError::IrEvalError { + msg: "capture-variable instruction without prior new-callee-stack \ + to initialize it" + .into(), + span: Some(*span), + }) + } + } + Instruction::PushVariable { var_id, src } => { + let value = ctx.collect_reg(*src, *span)?; + if let Some(callee_stack) = &mut ctx.callee_stack { + callee_stack.add_var(*var_id, value); + Ok(Continue) + } else { + Err(ShellError::IrEvalError { + msg: + "push-variable instruction without prior new-callee-stack to initialize it" + .into(), + span: Some(*span), + }) + } + } Instruction::PushPositional { src } => { let val = ctx.collect_reg(*src, *span)?; - ctx.stack + ctx.callee_stack() .arguments .push(Argument::Positional { span: *span, val }); Ok(Continue) } Instruction::AppendRest { src } => { let vals = ctx.collect_reg(*src, *span)?; - ctx.stack + ctx.callee_stack() .arguments .push(Argument::Spread { span: *span, vals }); Ok(Continue) } Instruction::PushFlag { name } => { - ctx.stack.arguments.push(Argument::Flag { - data: ctx.data.clone(), + let data = ctx.data.clone(); + ctx.callee_stack().arguments.push(Argument::Flag { + data, name: *name, span: *span, }); @@ -314,8 +360,9 @@ fn eval_instruction( } Instruction::PushNamed { name, src } => { let val = ctx.collect_reg(*src, *span)?; - ctx.stack.arguments.push(Argument::Named { - data: ctx.data.clone(), + let data = ctx.data.clone(); + ctx.callee_stack().arguments.push(Argument::Named { + data, name: *name, span: *span, val, @@ -323,8 +370,9 @@ fn eval_instruction( Ok(Continue) } Instruction::PushParserInfo { name, info } => { - ctx.stack.arguments.push(Argument::ParserInfo { - data: ctx.data.clone(), + let data = ctx.data.clone(); + ctx.callee_stack().arguments.push(Argument::ParserInfo { + data, name: *name, info: info.clone(), }); @@ -342,7 +390,7 @@ fn eval_instruction( } Instruction::Call { decl_id, src_dst } => { let input = ctx.take_reg(*src_dst); - let result = eval_call(ctx, *decl_id, *span, input)?; + let result = eval_call::(ctx, *decl_id, *span, input)?; ctx.put_reg(*src_dst, result); Ok(Continue) } @@ -726,7 +774,7 @@ fn binary_op( } /// Evaluate a call -fn eval_call( +fn eval_call( ctx: &mut EvalContext<'_>, decl_id: DeclId, head: Span, @@ -736,52 +784,91 @@ fn eval_call( engine_state, stack, args_base, + callee_stack, redirect_out, redirect_err, .. } = ctx; - // TODO: handle block eval + let stack = callee_stack.as_deref_mut().unwrap_or(stack); + let args_len = stack.arguments.get_len(*args_base); let decl = engine_state.get_decl(decl_id); // Set up redirect modes let mut stack = stack.push_redirection(redirect_out.take(), redirect_err.take()); - // should this be precalculated? ideally we just use the call builder... - let span = Span::merge_many( - std::iter::once(head).chain( - stack - .arguments - .get_args(*args_base, args_len) - .into_iter() - .flat_map(|arg| arg.span()), - ), - ); - let call = Call { - decl_id, - head, - span, - args_base: *args_base, - args_len, + let result = if let Some(block_id) = decl.block_id() { + // If the decl is a custom command, we assume that we have set up the arguments using + // new-callee-stack and push-variable instead of stack.arguments + // + // This saves us from having to parse through the declaration at eval time to figure out + // what to put where. + let block = engine_state.get_block(block_id); + + eval_block_with_early_return::(engine_state, &mut stack, block, input) + } else { + // should this be precalculated? ideally we just use the call builder... + let span = Span::merge_many( + std::iter::once(head).chain( + stack + .arguments + .get_args(*args_base, args_len) + .into_iter() + .flat_map(|arg| arg.span()), + ), + ); + + let call = Call { + decl_id, + head, + span, + args_base: *args_base, + args_len, + }; + + // Run the call + decl.run(engine_state, &mut stack, &(&call).into(), input) }; - // Run the call - let result = decl.run(engine_state, &mut stack, &(&call).into(), input); - // Important that this runs: + // Important that this runs, to reset state post-call: stack.arguments.leave_frame(ctx.args_base); + *redirect_out = None; + *redirect_err = None; + + drop(stack); + *callee_stack = None; + result } /// Get variable from [`Stack`] or [`EngineState`] fn get_var(ctx: &EvalContext<'_>, var_id: VarId, span: Span) -> Result { - ctx.stack.get_var(var_id, span).or_else(|err| { - if let Some(const_val) = ctx.engine_state.get_var(var_id).const_val.clone() { - Ok(const_val.with_span(span)) - } else { - Err(err) + match var_id { + // $env + ENV_VARIABLE_ID => { + let env_vars = ctx.stack.get_env_vars(ctx.engine_state); + let env_columns = env_vars.keys(); + let env_values = env_vars.values(); + + let mut pairs = env_columns + .map(|x| x.to_string()) + .zip(env_values.cloned()) + .collect::>(); + + pairs.sort_by(|a, b| a.0.cmp(&b.0)); + + Ok(Value::record(pairs.into_iter().collect(), span)) } - }) + _ => ctx.stack.get_var(var_id, span).or_else(|err| { + // $nu is handled by getting constant + if let Some(const_val) = ctx.engine_state.get_constant(var_id).cloned() { + Ok(const_val.with_span(span)) + } else { + Err(err) + } + }), + } } /// Helper to collect values into [`PipelineData`], preserving original span and metadata diff --git a/crates/nu-protocol/src/ir/display.rs b/crates/nu-protocol/src/ir/display.rs index 5d75b3a74d..8aa7e4ffaf 100644 --- a/crates/nu-protocol/src/ir/display.rs +++ b/crates/nu-protocol/src/ir/display.rs @@ -56,6 +56,10 @@ impl<'a> fmt::Display for FmtInstruction<'a> { }; write!(f, "{:WIDTH$} {dst}, {lit}", "load-literal") } + Instruction::LoadValue { dst, val } => { + let val = val.to_debug_string(); + write!(f, "{:WIDTH$} {dst}, {val}", "load-value") + } Instruction::Move { dst, src } => { write!(f, "{:WIDTH$} {dst}, {src}", "move") } @@ -91,6 +95,17 @@ impl<'a> fmt::Display for FmtInstruction<'a> { let key = FmtData(self.data, *key); write!(f, "{:WIDTH$} {key}, {src}", "store-env") } + Instruction::NewCalleeStack => { + write!(f, "{:WIDTH$}", "new-callee-stack") + } + Instruction::CaptureVariable { var_id } => { + let var = FmtVar::new(self.engine_state, *var_id); + write!(f, "{:WIDTH$} {var}", "capture-variable") + } + Instruction::PushVariable { var_id, src } => { + let var = FmtVar::new(self.engine_state, *var_id); + write!(f, "{:WIDTH$} {var}, {src}", "push-variable") + } Instruction::PushPositional { src } => { write!(f, "{:WIDTH$} {src}", "push-positional") } diff --git a/crates/nu-protocol/src/ir/mod.rs b/crates/nu-protocol/src/ir/mod.rs index a0b3687f72..0571142e73 100644 --- a/crates/nu-protocol/src/ir/mod.rs +++ b/crates/nu-protocol/src/ir/mod.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use crate::{ ast::{CellPath, Expression, Operator, Pattern, RangeInclusion}, engine::EngineState, - BlockId, DeclId, RegId, Span, VarId, + BlockId, DeclId, RegId, Span, Value, VarId, }; use serde::{Deserialize, Serialize}; @@ -61,6 +61,8 @@ impl std::ops::Index for [u8] { pub enum Instruction { /// Load a literal value into the `dst` register LoadLiteral { dst: RegId, lit: Literal }, + /// Load a clone of a boxed value into the `dst` register (e.g. from const evaluation) + LoadValue { dst: RegId, val: Box }, /// Move a register. Value is taken from `src` (used by this instruction). Move { dst: RegId, src: RegId }, /// Copy a register (must be a collected value). Value is still in `src` after this instruction. @@ -82,15 +84,21 @@ pub enum Instruction { LoadEnvOpt { dst: RegId, key: DataSlice }, /// Store the value of an environment variable from the `src` register StoreEnv { key: DataSlice, src: RegId }, - /// Add a positional arg to the next call + /// Create a stack for the next call + NewCalleeStack, + /// Capture a variable from the caller stack to the callee stack, preserving the original span. + CaptureVariable { var_id: VarId }, + /// Add a variable onto the callee's stack, for arguments. + PushVariable { var_id: VarId, src: RegId }, + /// Add a positional arg to the next (internal) call. PushPositional { src: RegId }, - /// Add a list of args to the next call (spread/rest) + /// Add a list of args to the next (internal) call (spread/rest). AppendRest { src: RegId }, - /// Add a named arg with no value to the next call. + /// Add a named arg with no value to the next (internal) call. PushFlag { name: DataSlice }, - /// Add a named arg with a value to the next call. + /// Add a named arg with a value to the next (internal) call. PushNamed { name: DataSlice, src: RegId }, - /// Add parser info to the next call. + /// Add parser info to the next (internal) call. PushParserInfo { name: DataSlice, info: Box,