custom commands, fix let, fix bare $env var

This commit is contained in:
Devyn Cairns 2024-07-03 20:15:00 -07:00
parent d33d21e1dc
commit 5c419a39e8
No known key found for this signature in database
6 changed files with 368 additions and 49 deletions

View File

@ -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: _ } => (),

View File

@ -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,

View File

@ -307,7 +307,7 @@ pub(crate) fn compile_let(
builder,
block,
redirect_modes.with_capture_out(call.head),
None,
Some(io_reg),
io_reg,
)?;

View File

@ -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<D: DebugContext>(
@ -35,6 +35,7 @@ pub fn eval_ir_block<D: DebugContext>(
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<Box<Stack>>,
/// State set by redirect-out
redirect_out: Option<Redirection>,
/// 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<D: DebugContext>(
"{pc:-4}: {}",
instruction.display(ctx.engine_state, ctx.data)
);
match eval_instruction(ctx, instruction, span) {
match eval_instruction::<D>(ctx, instruction, span) {
Ok(InstructionResult::Continue) => {
pc += 1;
}
@ -193,7 +201,7 @@ enum InstructionResult {
}
/// Perform an instruction
fn eval_instruction(
fn eval_instruction<D: DebugContext>(
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::<D>(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<D: DebugContext>(
ctx: &mut EvalContext<'_>,
decl_id: DeclId,
head: Span,
@ -736,18 +784,30 @@ 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());
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::<D>(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(
@ -758,6 +818,7 @@ fn eval_call(
.flat_map(|arg| arg.span()),
),
);
let call = Call {
decl_id,
head,
@ -767,21 +828,47 @@ fn eval_call(
};
// Run the call
let result = decl.run(engine_state, &mut stack, &(&call).into(), input);
// Important that this runs:
decl.run(engine_state, &mut stack, &(&call).into(), input)
};
// 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<Value, ShellError> {
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() {
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::<Vec<(String, Value)>>();
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

View File

@ -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")
}

View File

@ -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<DataSlice> 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<Value> },
/// 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<Expression>,