Move custom command call stack handling back to the eval side

This commit is contained in:
Devyn Cairns 2024-07-06 00:55:57 -07:00
parent c249922e03
commit 9498c06631
No known key found for this signature in database
6 changed files with 169 additions and 313 deletions

View File

@ -163,9 +163,6 @@ impl BlockBuilder {
Instruction::LoadEnv { dst, key: _ } => allocate(&[], &[*dst]),
Instruction::LoadEnvOpt { dst, key: _ } => allocate(&[], &[*dst]),
Instruction::StoreEnv { key: _, src } => allocate(&[*src], &[]),
Instruction::NewCalleeStack => Ok(()),
Instruction::CaptureVariable { var_id: _ } => Ok(()),
Instruction::PushVariable { var_id: _, src } => allocate(&[*src], &[]),
Instruction::PushPositional { src } => allocate(&[*src], &[]),
Instruction::AppendRest { src } => allocate(&[*src], &[]),
Instruction::PushFlag { name: _ } => Ok(()),

View File

@ -1,7 +1,7 @@
use std::{iter::repeat, sync::Arc};
use std::sync::Arc;
use nu_protocol::{
ast::{Argument, Block, Call, Expression, ExternalArgument},
ast::{Argument, Call, Expression, ExternalArgument},
engine::StateWorkingSet,
ir::{Instruction, IrAstRef, Literal},
IntoSpanned, RegId, Span, Spanned,
@ -69,20 +69,6 @@ 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,
);
}
// Keep AST if the decl needs it.
let requires_ast = decl.requires_ast_for_arguments();
@ -227,195 +213,6 @@ pub(crate) fn compile_help(
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,

View File

@ -6,8 +6,9 @@ use nu_protocol::{
debugger::DebugContext,
engine::{Argument, Closure, EngineState, ErrorHandler, Matcher, Redirection, Stack},
ir::{Call, DataSlice, Instruction, IrAstRef, IrBlock, Literal, RedirectMode},
record, DeclId, IntoPipelineData, IntoSpanned, ListStream, OutDest, PipelineData, Range,
Record, RegId, ShellError, Span, Spanned, Value, VarId, ENV_VARIABLE_ID,
record, DeclId, Flag, IntoPipelineData, IntoSpanned, ListStream, OutDest, PipelineData,
PositionalArg, Range, Record, RegId, ShellError, Signature, Span, Spanned, Value, VarId,
ENV_VARIABLE_ID,
};
use crate::{eval::is_automatic_env_var, eval_block_with_early_return};
@ -51,7 +52,6 @@ pub fn eval_ir_block<D: DebugContext>(
block_span: &block.span,
args_base,
error_handler_base,
callee_stack: None,
redirect_out: None,
redirect_err: None,
matches: vec![],
@ -90,8 +90,6 @@ 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<Box<Stack>>,
/// State set by redirect-out
redirect_out: Option<Redirection>,
/// State set by redirect-err
@ -130,31 +128,6 @@ 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`
#[inline]
fn callee_stack(&mut self) -> &mut Stack {
self.callee_stack.as_deref_mut().unwrap_or(self.stack)
}
/// Create a new callee stack, so that arguments will go onto a new stack that will be used for
/// calls
fn new_callee_stack(&mut self) {
self.drop_callee_stack();
let mut new_stack = Box::new(self.stack.gather_captures(self.engine_state, &[]));
// Increment recursion count on callee stack to prevent recursing too far
new_stack.recursion_count += 1;
self.callee_stack = Some(new_stack);
}
/// Drop the callee stack, so that arguments will be put on the caller stack instead
fn drop_callee_stack(&mut self) {
// this is provided so in case some other cleanup needs to be done, it can be done here.
self.callee_stack = None;
}
}
/// Eval an IR block on the provided slice of registers.
@ -342,41 +315,9 @@ fn eval_instruction<D: DebugContext>(
})
}
}
Instruction::NewCalleeStack => {
ctx.new_callee_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.callee_stack().arguments.push(Argument::Positional {
ctx.stack.arguments.push(Argument::Positional {
span: *span,
val,
ast: ast.clone().map(|ast_ref| ast_ref.0),
@ -385,7 +326,7 @@ fn eval_instruction<D: DebugContext>(
}
Instruction::AppendRest { src } => {
let vals = ctx.collect_reg(*src, *span)?;
ctx.callee_stack().arguments.push(Argument::Spread {
ctx.stack.arguments.push(Argument::Spread {
span: *span,
vals,
ast: ast.clone().map(|ast_ref| ast_ref.0),
@ -394,7 +335,7 @@ fn eval_instruction<D: DebugContext>(
}
Instruction::PushFlag { name } => {
let data = ctx.data.clone();
ctx.callee_stack().arguments.push(Argument::Flag {
ctx.stack.arguments.push(Argument::Flag {
data,
name: *name,
span: *span,
@ -404,7 +345,7 @@ fn eval_instruction<D: DebugContext>(
Instruction::PushNamed { name, src } => {
let val = ctx.collect_reg(*src, *span)?;
let data = ctx.data.clone();
ctx.callee_stack().arguments.push(Argument::Named {
ctx.stack.arguments.push(Argument::Named {
data,
name: *name,
span: *span,
@ -415,7 +356,7 @@ fn eval_instruction<D: DebugContext>(
}
Instruction::PushParserInfo { name, info } => {
let data = ctx.data.clone();
ctx.callee_stack().arguments.push(Argument::ParserInfo {
ctx.stack.arguments.push(Argument::ParserInfo {
data,
name: *name,
info: info.clone(),
@ -831,44 +772,50 @@ fn eval_call<D: DebugContext>(
engine_state,
stack: caller_stack,
args_base,
callee_stack,
redirect_out,
redirect_err,
..
} = ctx;
let stack = callee_stack.as_deref_mut().unwrap_or(caller_stack);
let args_len = stack.arguments.get_len(*args_base);
let args_len = caller_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());
let mut caller_stack = caller_stack.push_redirection(redirect_out.take(), redirect_err.take());
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.
// If the decl is a custom command
let block = engine_state.get_block(block_id);
result = eval_block_with_early_return::<D>(engine_state, &mut stack, block, input);
// Set up a callee stack with the captures and move arguments from the stack into variables
let mut callee_stack = caller_stack.gather_captures(engine_state, &block.captures);
drop(stack);
gather_arguments(
block,
&mut caller_stack,
&mut callee_stack,
*args_base,
args_len,
head,
)?;
// Add one to the recursion count, so we don't recurse too deep. Stack overflows are not
// recoverable in Rust.
callee_stack.recursion_count += 1;
result = eval_block_with_early_return::<D>(engine_state, &mut callee_stack, block, input);
// Move environment variables back into the caller stack scope if requested to do so
if block.redirect_env {
if let Some(callee_stack) = callee_stack {
redirect_env(engine_state, caller_stack, callee_stack);
}
redirect_env(engine_state, &mut caller_stack, &mut callee_stack);
}
} else {
// should this be precalculated? ideally we just use the call builder...
// FIXME: precalculate this and save it somewhere
let span = Span::merge_many(
std::iter::once(head).chain(
stack
caller_stack
.arguments
.get_args(*args_base, args_len)
.into_iter()
@ -885,20 +832,146 @@ fn eval_call<D: DebugContext>(
};
// Run the call
result = decl.run(engine_state, &mut stack, &(&call).into(), input);
drop(stack);
result = decl.run(engine_state, &mut caller_stack, &(&call).into(), input);
};
drop(caller_stack);
// Important that this runs, to reset state post-call:
ctx.stack.arguments.leave_frame(ctx.args_base);
ctx.redirect_out = None;
ctx.redirect_err = None;
ctx.drop_callee_stack();
result
}
fn find_named_var_id(sig: &Signature, name: &[u8], span: Span) -> Result<VarId, ShellError> {
sig.named
.iter()
.find(|n| n.long.as_bytes() == name)
.ok_or_else(|| ShellError::IrEvalError {
msg: format!(
"block does not have an argument named `{}`",
String::from_utf8_lossy(name)
),
span: Some(span),
})
.and_then(|flag| expect_named_var_id(flag, span))
}
fn expect_named_var_id(arg: &Flag, span: Span) -> Result<VarId, ShellError> {
arg.var_id.ok_or_else(|| ShellError::IrEvalError {
msg: format!(
"block signature is missing var id for named arg `{}`",
arg.long
),
span: Some(span),
})
}
fn expect_positional_var_id(arg: &PositionalArg, span: Span) -> Result<VarId, ShellError> {
arg.var_id.ok_or_else(|| ShellError::IrEvalError {
msg: format!(
"block signature is missing var id for positional arg `{}`",
arg.name
),
span: Some(span),
})
}
/// Move arguments from the stack into variables for a custom command
fn gather_arguments(
block: &Block,
caller_stack: &mut Stack,
callee_stack: &mut Stack,
args_base: usize,
args_len: usize,
call_head: Span,
) -> Result<(), ShellError> {
let mut positional_iter = block
.signature
.required_positional
.iter()
.chain(&block.signature.optional_positional);
// Arguments that didn't get consumed by required/optional
let mut rest = vec![];
for arg in caller_stack.arguments.drain_args(args_base, args_len) {
match arg {
Argument::Positional { span, val, .. } => {
if let Some(positional_arg) = positional_iter.next() {
let var_id = expect_positional_var_id(positional_arg, span)?;
callee_stack.add_var(var_id, val);
} else {
rest.push(val);
}
}
Argument::Spread { vals, .. } => {
if let Value::List { vals, .. } = vals {
rest.extend(vals);
} else {
return Err(ShellError::CannotSpreadAsList { span: vals.span() });
}
}
Argument::Flag { data, name, span } => {
let var_id = find_named_var_id(&block.signature, &data[name], span)?;
callee_stack.add_var(var_id, Value::bool(true, span))
}
Argument::Named {
data,
name,
span,
val,
..
} => {
let var_id = find_named_var_id(&block.signature, &data[name], span)?;
callee_stack.add_var(var_id, val)
}
Argument::ParserInfo { .. } => (),
}
}
// Add the collected rest of the arguments if a spread argument exists
if let Some(rest_arg) = &block.signature.rest_positional {
let rest_span = rest.first().map(|v| v.span()).unwrap_or(call_head);
let var_id = expect_positional_var_id(rest_arg, rest_span)?;
callee_stack.add_var(var_id, Value::list(rest, rest_span));
}
// Check for arguments that haven't yet been set and set them to their defaults
for positional_arg in positional_iter {
let var_id = expect_positional_var_id(positional_arg, call_head)?;
callee_stack.add_var(
var_id,
positional_arg
.default_value
.clone()
.unwrap_or(Value::nothing(call_head)),
);
}
for named_arg in &block.signature.named {
if let Some(var_id) = named_arg.var_id {
// For named arguments, we do this check by looking to see if the variable was set yet on
// the stack. This assumes that the stack's variables was previously empty, but that's a
// fair assumption for a brand new callee stack.
if !callee_stack.vars.iter().any(|(id, _)| *id == var_id) {
let val = if named_arg.arg.is_none() {
Value::bool(false, call_head)
} else if let Some(value) = &named_arg.default_value {
value.clone()
} else {
Value::nothing(call_head)
};
callee_stack.add_var(var_id, val);
}
}
}
Ok(())
}
/// Get variable from [`Stack`] or [`EngineState`]
fn get_var(ctx: &EvalContext<'_>, var_id: VarId, span: Span) -> Result<Value, ShellError> {
match var_id {

View File

@ -113,4 +113,10 @@ impl ArgumentStack {
pub fn get_args(&self, base: usize, len: usize) -> &[Argument] {
&self.arguments[base..(base + len)]
}
/// Move arguments for the frame based on the given [`base`](`.get_base()`) and
/// [`len`](`.get_len()`) parameters.
pub fn drain_args(&mut self, base: usize, len: usize) -> impl Iterator<Item = Argument> + '_ {
self.arguments.drain(base..(base + len))
}
}

View File

@ -95,17 +95,6 @@ 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")
}

View File

@ -121,12 +121,6 @@ pub enum Instruction {
LoadEnvOpt { dst: RegId, key: DataSlice },
/// Store the value of an environment variable from the `src` register
StoreEnv { key: DataSlice, src: RegId },
/// 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 (internal) call (spread/rest).