make timeit and metadata work through hack that's ok for now

crates/nu-command/src/debug/metadata.rs

@@ -28,6 +28,10 @@ impl Command for Metadata {
             .category(Category::Debug)
     }
 
+    fn requires_ast_for_arguments(&self) -> bool {
+        true
+    }
+
     fn run(
         &self,
         engine_state: &EngineState,
@@ -35,8 +39,7 @@ impl Command for Metadata {
         call: &Call,
         input: PipelineData,
     ) -> Result<PipelineData, ShellError> {
-        let call = call.assert_ast_call()?; // FIXME
+        let arg = call.positional_nth(stack, 0);
-        let arg = call.positional_nth(0);
         let head = call.head;
 
         match arg {

crates/nu-command/src/debug/timeit.rs

@@ -32,6 +32,10 @@ impl Command for TimeIt {
         vec!["timing", "timer", "benchmark", "measure"]
     }
 
+    fn requires_ast_for_arguments(&self) -> bool {
+        true
+    }
+
     fn run(
         &self,
         engine_state: &EngineState,
@@ -39,15 +43,14 @@ impl Command for TimeIt {
         call: &Call,
         input: PipelineData,
     ) -> Result<PipelineData, ShellError> {
-        // FIXME: get working with IR. I don't think this should actually be so AST dependent
+        // reset outdest, so the command can write to stdout and stderr.
-        let call = call.assert_ast_call()?;
+        let stack = &mut stack.push_redirection(None, None);
-        let command_to_run = call.positional_nth(0);
+
+        let command_to_run = call.positional_nth(stack, 0);
 
         // Get the start time after all other computation has been done.
         let start_time = Instant::now();
 
-        // reset outdest, so the command can write to stdout and stderr.
-        let stack = &mut stack.push_redirection(None, None);
         if let Some(command_to_run) = command_to_run {
             if let Some(block_id) = command_to_run.as_block() {
                 let eval_block = get_eval_block(engine_state);
@@ -55,7 +58,8 @@ impl Command for TimeIt {
                 eval_block(engine_state, stack, block, input)?
             } else {
                 let eval_expression_with_input = get_eval_expression_with_input(engine_state);
-                eval_expression_with_input(engine_state, stack, command_to_run, input)?.0
+                let expression = &command_to_run.clone();
+                eval_expression_with_input(engine_state, stack, expression, input)?.0
             }
         } else {
             PipelineData::empty()

crates/nu-engine/src/compile/builder.rs

@@ -1,5 +1,5 @@
 use nu_protocol::{
-    ir::{DataSlice, Instruction, IrBlock, Literal, RedirectMode},
+    ir::{DataSlice, Instruction, IrAstRef, IrBlock, Literal, RedirectMode},
     CompileError, IntoSpanned, RegId, Span, Spanned,
 };
 
@@ -9,6 +9,7 @@ pub(crate) struct BlockBuilder {
     pub(crate) instructions: Vec<Instruction>,
     pub(crate) spans: Vec<Span>,
     pub(crate) data: Vec<u8>,
+    pub(crate) ast: Vec<Option<IrAstRef>>,
     pub(crate) register_allocation_state: Vec<bool>,
 }
 
@@ -19,6 +20,7 @@ impl BlockBuilder {
             instructions: vec![],
             spans: vec![],
             data: vec![],
+            ast: vec![],
             register_allocation_state: vec![true],
         }
     }
@@ -223,9 +225,15 @@ impl BlockBuilder {
         let index = self.next_instruction_index();
         self.instructions.push(instruction.item);
         self.spans.push(instruction.span);
+        self.ast.push(None);
         Ok(index)
     }
 
+    /// Set the AST of the last instruction. Separate method because it's rarely used.
+    pub(crate) fn set_last_ast(&mut self, ast_ref: Option<IrAstRef>) {
+        *self.ast.last_mut().expect("no last instruction") = ast_ref;
+    }
+
     /// Load a register with a literal.
     pub(crate) fn load_literal(
         &mut self,
@@ -369,6 +377,7 @@ impl BlockBuilder {
             instructions: self.instructions,
             spans: self.spans,
             data: self.data.into(),
+            ast: self.ast,
             register_count: self.register_allocation_state.len(),
         }
     }

crates/nu-engine/src/compile/call.rs

@@ -1,9 +1,9 @@
-use std::iter::repeat;
+use std::{iter::repeat, sync::Arc};
 
 use nu_protocol::{
     ast::{Argument, Block, Call, Expression, ExternalArgument},
     engine::StateWorkingSet,
-    ir::{Instruction, Literal},
+    ir::{Instruction, IrAstRef, Literal},
     IntoSpanned, RegId, Span, Spanned,
 };
 
@@ -83,6 +83,9 @@ pub(crate) fn compile_call(
         );
     }
 
+    // Keep AST if the decl needs it.
+    let requires_ast = decl.requires_ast_for_arguments();
+
     // It's important that we evaluate the args first before trying to set up the argument
     // state for the call.
     //
@@ -90,9 +93,9 @@ pub(crate) fn compile_call(
     // the argument state, but we'd have to check all of that first and it just isn't really worth
     // it.
     enum CompiledArg<'a> {
-        Positional(RegId, Span),
+        Positional(RegId, Span, Option<IrAstRef>),
-        Named(&'a str, Option<RegId>, Span),
+        Named(&'a str, Option<RegId>, Span, Option<IrAstRef>),
-        Spread(RegId, Span),
+        Spread(RegId, Span, Option<IrAstRef>),
     }
 
     let mut compiled_args = vec![];
@@ -116,16 +119,29 @@ pub(crate) fn compile_call(
             })
             .transpose()?;
 
+        let ast_ref = arg
+            .expr()
+            .filter(|_| requires_ast)
+            .map(|expr| IrAstRef(Arc::new(expr.clone())));
+
         match arg {
-            Argument::Positional(_) | Argument::Unknown(_) => compiled_args.push(
+            Argument::Positional(_) | Argument::Unknown(_) => {
-                CompiledArg::Positional(arg_reg.expect("expr() None in non-Named"), arg.span()),
+                compiled_args.push(CompiledArg::Positional(
-            ),
+                    arg_reg.expect("expr() None in non-Named"),
-            Argument::Named((name, _, _)) => {
+                    arg.span(),
-                compiled_args.push(CompiledArg::Named(name.item.as_str(), arg_reg, arg.span()))
+                    ast_ref,
+                ))
             }
+            Argument::Named((name, _, _)) => compiled_args.push(CompiledArg::Named(
+                name.item.as_str(),
+                arg_reg,
+                arg.span(),
+                ast_ref,
+            )),
             Argument::Spread(_) => compiled_args.push(CompiledArg::Spread(
                 arg_reg.expect("expr() None in non-Named"),
                 arg.span(),
+                ast_ref,
             )),
         }
     }
@@ -133,19 +149,23 @@ pub(crate) fn compile_call(
     // Now that the args are all compiled, set up the call state (argument stack and redirections)
     for arg in compiled_args {
         match arg {
-            CompiledArg::Positional(reg, span) => {
+            CompiledArg::Positional(reg, span, ast_ref) => {
                 builder.push(Instruction::PushPositional { src: reg }.into_spanned(span))?;
+                builder.set_last_ast(ast_ref);
             }
-            CompiledArg::Named(name, Some(reg), span) => {
+            CompiledArg::Named(name, Some(reg), span, ast_ref) => {
                 let name = builder.data(name)?;
                 builder.push(Instruction::PushNamed { name, src: reg }.into_spanned(span))?;
+                builder.set_last_ast(ast_ref);
             }
-            CompiledArg::Named(name, None, span) => {
+            CompiledArg::Named(name, None, span, ast_ref) => {
                 let name = builder.data(name)?;
                 builder.push(Instruction::PushFlag { name }.into_spanned(span))?;
+                builder.set_last_ast(ast_ref);
             }
-            CompiledArg::Spread(reg, span) => {
+            CompiledArg::Spread(reg, span, ast_ref) => {
                 builder.push(Instruction::AppendRest { src: reg }.into_spanned(span))?;
+                builder.set_last_ast(ast_ref);
             }
         }
     }

crates/nu-engine/src/eval_ir.rs

@@ -5,7 +5,7 @@ use nu_protocol::{
     ast::{Bits, Block, Boolean, CellPath, Comparison, Math, Operator},
     debugger::DebugContext,
     engine::{Argument, Closure, EngineState, ErrorHandler, Matcher, Redirection, Stack},
-    ir::{Call, DataSlice, Instruction, IrBlock, Literal, RedirectMode},
+    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,
 };
@@ -166,11 +166,12 @@ fn eval_ir_block_impl<D: DebugContext>(
     while pc < ir_block.instructions.len() {
         let instruction = &ir_block.instructions[pc];
         let span = &ir_block.spans[pc];
+        let ast = &ir_block.ast[pc];
         log::trace!(
             "{pc:-4}: {}",
             instruction.display(ctx.engine_state, ctx.data)
         );
-        match eval_instruction::<D>(ctx, instruction, span) {
+        match eval_instruction::<D>(ctx, instruction, span, ast) {
             Ok(InstructionResult::Continue) => {
                 pc += 1;
             }
@@ -236,6 +237,7 @@ fn eval_instruction<D: DebugContext>(
     ctx: &mut EvalContext<'_>,
     instruction: &Instruction,
     span: &Span,
+    ast: &Option<IrAstRef>,
 ) -> Result<InstructionResult, ShellError> {
     use self::InstructionResult::*;
 
@@ -367,16 +369,20 @@ fn eval_instruction<D: DebugContext>(
         }
         Instruction::PushPositional { src } => {
             let val = ctx.collect_reg(*src, *span)?;
-            ctx.callee_stack()
+            ctx.callee_stack().arguments.push(Argument::Positional {
-                .arguments
+                span: *span,
-                .push(Argument::Positional { span: *span, val });
+                val,
+                ast: ast.clone().map(|ast_ref| ast_ref.0),
+            });
             Ok(Continue)
         }
         Instruction::AppendRest { src } => {
             let vals = ctx.collect_reg(*src, *span)?;
-            ctx.callee_stack()
+            ctx.callee_stack().arguments.push(Argument::Spread {
-                .arguments
+                span: *span,
-                .push(Argument::Spread { span: *span, vals });
+                vals,
+                ast: ast.clone().map(|ast_ref| ast_ref.0),
+            });
             Ok(Continue)
         }
         Instruction::PushFlag { name } => {
@@ -396,6 +402,7 @@ fn eval_instruction<D: DebugContext>(
                 name: *name,
                 span: *span,
                 val,
+                ast: ast.clone().map(|ast_ref| ast_ref.0),
             });
             Ok(Continue)
         }

crates/nu-parser/src/known_external.rs

@@ -166,6 +166,7 @@ fn ir_call_to_extern_call(
                         std::str::from_utf8(&data[*name]).expect("invalid flag name"),
                         *span,
                     ),
+                    ast: None,
                 };
                 extern_call.add_argument(stack, name_arg);
             }
@@ -174,6 +175,7 @@ fn ir_call_to_extern_call(
                 name,
                 span,
                 val,
+                ..
             } => {
                 let name_arg = engine::Argument::Positional {
                     span: *span,
@@ -181,10 +183,12 @@ fn ir_call_to_extern_call(
                         std::str::from_utf8(&data[*name]).expect("invalid arg name"),
                         *span,
                     ),
+                    ast: None,
                 };
                 let val_arg = engine::Argument::Positional {
                     span: *span,
                     val: val.clone(),
+                    ast: None,
                 };
                 extern_call.add_argument(stack, name_arg);
                 extern_call.add_argument(stack, val_arg);

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

@@ -6,9 +6,17 @@ use crate::{ast::Expression, ir::DataSlice, Span, Value};
 #[derive(Debug, Clone)]
 pub enum Argument {
     /// A positional argument
-    Positional { span: Span, val: Value },
+    Positional {
+        span: Span,
+        val: Value,
+        ast: Option<Arc<Expression>>,
+    },
     /// A spread argument, e.g. `...$args`
-    Spread { span: Span, vals: Value },
+    Spread {
+        span: Span,
+        vals: Value,
+        ast: Option<Arc<Expression>>,
+    },
     /// A named argument with no value, e.g. `--flag`
     Flag {
         data: Arc<[u8]>,
@@ -21,6 +29,7 @@ pub enum Argument {
         name: DataSlice,
         span: Span,
         val: Value,
+        ast: Option<Arc<Expression>>,
     },
     /// Information generated by the parser for use by certain keyword commands
     ParserInfo {
@@ -45,6 +54,18 @@ impl Argument {
             Argument::ParserInfo { .. } => None,
         }
     }
+
+    /// The original AST [`Expression`] for the argument's value. This is not usually available;
+    /// declarations have to opt-in if they require this.
+    pub fn ast_expression(&self) -> Option<&Arc<Expression>> {
+        match self {
+            Argument::Positional { ast, .. } => ast.as_ref(),
+            Argument::Spread { ast, .. } => ast.as_ref(),
+            Argument::Flag { .. } => None,
+            Argument::Named { ast, .. } => ast.as_ref(),
+            Argument::ParserInfo { .. } => None,
+        }
+    }
 }
 
 /// Stores the argument context for calls in IR evaluation.

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

@@ -178,6 +178,17 @@ impl Call<'_> {
             CallImpl::IrBox(call) => by_ir(call, stack, starting_pos),
         }
     }
+
+    /// Get the original AST expression for a positional argument. Does not usually work for IR
+    /// unless the decl specified `requires_ast_for_arguments()`
+    pub fn positional_nth<'a>(&'a self, stack: &'a Stack, index: usize) -> Option<&'a Expression> {
+        match &self.inner {
+            CallImpl::AstRef(call) => call.positional_nth(index),
+            CallImpl::AstBox(call) => call.positional_nth(index),
+            CallImpl::IrRef(call) => call.positional_ast(stack, index).map(|arc| arc.as_ref()),
+            CallImpl::IrBox(call) => call.positional_ast(stack, index).map(|arc| arc.as_ref()),
+        }
+    }
 }
 
 impl CallImpl<'_> {

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

@@ -124,6 +124,12 @@ pub trait Command: Send + Sync + CommandClone {
     fn pipe_redirection(&self) -> (Option<OutDest>, Option<OutDest>) {
         (None, None)
     }
+
+    /// Return true if the AST nodes for the arguments are required for IR evaluation. This is
+    /// currently inefficient so is not generally done.
+    fn requires_ast_for_arguments(&self) -> bool {
+        false
+    }
 }
 
 pub trait CommandClone {

crates/nu-protocol/src/ir/call.rs

@@ -141,6 +141,24 @@ impl Call {
         self.positional_iter(stack).nth(index)
     }
 
+    /// Get the AST node for a positional argument by index. Not usually available unless the decl
+    /// required it.
+    pub fn positional_ast<'a>(
+        &self,
+        stack: &'a Stack,
+        index: usize,
+    ) -> Option<&'a Arc<Expression>> {
+        self.arguments(stack)
+            .iter()
+            .filter_map(|arg| match arg {
+                Argument::Positional { ast, .. } => Some(ast),
+                _ => None,
+            })
+            .nth(index)
+            .map(|option| option.as_ref())
+            .flatten()
+    }
+
     /// Returns every argument to the rest parameter, as well as whether each argument
     /// is spread or a normal positional argument (true for spread, false for normal)
     pub fn rest_iter<'a>(
@@ -230,12 +248,26 @@ impl CallBuilder {
 
     /// Add a positional argument to the [`Stack`] and reference it from the [`Call`].
     pub fn add_positional(&mut self, stack: &mut Stack, span: Span, val: Value) -> &mut Self {
-        self.add_argument(stack, Argument::Positional { span, val })
+        self.add_argument(
+            stack,
+            Argument::Positional {
+                span,
+                val,
+                ast: None,
+            },
+        )
     }
 
     /// Add a spread argument to the [`Stack`] and reference it from the [`Call`].
     pub fn add_spread(&mut self, stack: &mut Stack, span: Span, vals: Value) -> &mut Self {
-        self.add_argument(stack, Argument::Spread { span, vals })
+        self.add_argument(
+            stack,
+            Argument::Spread {
+                span,
+                vals,
+                ast: None,
+            },
+        )
     }
 
     /// Add a flag (no-value named) argument to the [`Stack`] and reference it from the [`Call`].
@@ -268,6 +300,7 @@ impl CallBuilder {
                 name,
                 span,
                 val,
+                ast: None,
             },
         )
     }

crates/nu-protocol/src/ir/mod.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{fmt, sync::Arc};
 
 use crate::{
     ast::{CellPath, Expression, Operator, Pattern, RangeInclusion},
@@ -14,15 +14,28 @@ mod display;
 pub use call::*;
 pub use display::{FmtInstruction, FmtIrBlock};
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Clone, Serialize, Deserialize)]
 pub struct IrBlock {
     pub instructions: Vec<Instruction>,
     pub spans: Vec<Span>,
     #[serde(with = "serde_arc_u8_array")]
     pub data: Arc<[u8]>,
+    pub ast: Vec<Option<IrAstRef>>,
     pub register_count: usize,
 }
 
+impl fmt::Debug for IrBlock {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        // the ast field is too verbose and doesn't add much
+        f.debug_struct("IrBlock")
+            .field("instructions", &self.instructions)
+            .field("spans", &self.spans)
+            .field("data", &self.data)
+            .field("register_count", &self.register_count)
+            .finish_non_exhaustive()
+    }
+}
+
 impl IrBlock {
     /// Returns a value that can be formatted with [`Display`](std::fmt::Display) to show a detailed
     /// listing of the instructions contained within this [`IrBlock`].
@@ -57,6 +70,29 @@ impl std::ops::Index<DataSlice> for [u8] {
     }
 }
 
+/// A possible reference into the abstract syntax tree for an instruction. This is not present for
+/// most instructions and is just added when needed.
+#[derive(Debug, Clone)]
+pub struct IrAstRef(pub Arc<Expression>);
+
+impl Serialize for IrAstRef {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        self.0.as_ref().serialize(serializer)
+    }
+}
+
+impl<'de> Deserialize<'de> for IrAstRef {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        Expression::deserialize(deserializer).map(|expr| IrAstRef(Arc::new(expr)))
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum Instruction {
     /// Load a literal value into the `dst` register