try/catch with catch as an expression rather than literal block

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

@@ -197,11 +197,8 @@ impl BlockBuilder {
                 stream: _,
                 end_index: _,
             } => self.mark_register(*dst)?,
-            Instruction::PushErrorHandler { index: _ } => (),
+            Instruction::OnError { index: _ } => (),
-            Instruction::PushErrorHandlerVar {
+            Instruction::OnErrorInto { index: _, dst } => self.mark_register(*dst)?,
-                index: _,
-                error_var: _,
-            } => (),
             Instruction::PopErrorHandler => (),
             Instruction::Return { src } => self.free_register(*src)?,
         }
@@ -318,8 +315,8 @@ impl BlockBuilder {
                 | Instruction::Iterate {
                     end_index: index, ..
                 }
-                | Instruction::PushErrorHandler { index }
+                | Instruction::OnError { index }
-                | Instruction::PushErrorHandlerVar { index, .. },
+                | Instruction::OnErrorInto { index, .. },
             ) => {
                 *index = target_index;
                 Ok(())

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

@@ -1,8 +1,8 @@
 use nu_protocol::{
-    ast::{Call, Expr, Expression},
+    ast::{Block, Call, Expr, Expression},
     engine::StateWorkingSet,
     ir::Instruction,
-    IntoSpanned, RegId,
+    IntoSpanned, RegId, VarId,
 };
 
 use super::{compile_block, compile_expression, BlockBuilder, CompileError, RedirectModes};
@@ -186,14 +186,26 @@ pub(crate) fn compile_try(
     redirect_modes: RedirectModes,
     io_reg: RegId,
 ) -> Result<(), CompileError> {
-    // Pseudocode:
+    // Pseudocode (literal block):
     //
-    //      push-error-handler-var ERR, $err     // or without var
+    //      on-error-with ERR, %io_reg           // or without
     //      %io_reg <- <...block...> <- %io_reg
     //      pop-error-handler
     //      jump END
-    // ERR: %io_reg <- <...catch block...>       // set to empty if none
+    // ERR: store-variable $err_var, %io_reg     // or without
+    //      %io_reg <- <...catch block...>       // set to empty if no catch block
     // END:
+    //
+    // with expression that can't be inlined:
+    //
+    //      %closure_reg <- <catch_expr>
+    //      on-error-with ERR, %io_reg
+    //      %io_reg <- <...block...> <- %io_reg
+    //      pop-error-handler
+    //      jump END
+    // ERR: push-positional %closure_reg
+    //      push-positional %io_reg
+    //      call "do", %io_reg
     let invalid = || CompileError::InvalidKeywordCall {
         keyword: "try".into(),
         span: call.head,
@@ -203,29 +215,68 @@ pub(crate) fn compile_try(
     let block_id = block_arg.as_block().ok_or_else(invalid)?;
     let block = working_set.get_block(block_id);
 
-    let catch_block = match call.positional_nth(1) {
+    let catch_expr = match call.positional_nth(1) {
-        Some(kw_expr) => {
+        Some(kw_expr) => Some(kw_expr.as_keyword().ok_or_else(invalid)?),
-            let catch_expr = kw_expr.as_keyword().ok_or_else(invalid)?;
-            let catch_block_id = catch_expr.as_block().ok_or_else(invalid)?;
-            Some(working_set.get_block(catch_block_id))
-        }
         None => None,
     };
-    let catch_var_id = catch_block
-        .and_then(|b| b.signature.get_positional(0))
-        .and_then(|v| v.var_id);
 
-    // Put the error handler placeholder
+    // We have two ways of executing `catch`: if it was provided as a literal, we can inline it.
-    let error_handler_index = if let Some(catch_var_id) = catch_var_id {
+    // Otherwise, we have to evaluate the expression and keep it as a register, and then call `do`.
+    enum CatchType<'a> {
+        Block {
+            block: &'a Block,
+            var_id: Option<VarId>,
+        },
+        Closure {
+            closure_reg: RegId,
+        },
+    }
+
+    let catch_type = catch_expr
+        .map(|catch_expr| match catch_expr.as_block() {
+            Some(block_id) => {
+                let block = working_set.get_block(block_id);
+                let var_id = block.signature.get_positional(0).and_then(|v| v.var_id);
+                Ok(CatchType::Block { block, var_id })
+            }
+            None => {
+                // We have to compile the catch_expr and use it as a closure
+                let closure_reg = builder.next_register()?;
+                compile_expression(
+                    working_set,
+                    builder,
+                    catch_expr,
+                    redirect_modes.with_capture_out(catch_expr.span),
+                    None,
+                    closure_reg,
+                )?;
+                Ok(CatchType::Closure { closure_reg })
+            }
+        })
+        .transpose()?;
+
+    // Put the error handler placeholder. If the catch argument is a non-block expression or a block
+    // that takes an argument, we should capture the error into `io_reg` since we safely don't need
+    // that.
+    let error_handler_index = if matches!(
+        catch_type,
+        Some(
+            CatchType::Block {
+                var_id: Some(_),
+                ..
+            } | CatchType::Closure { .. }
+        )
+    ) {
         builder.push(
-            Instruction::PushErrorHandlerVar {
+            Instruction::OnErrorInto {
                 index: usize::MAX,
-                error_var: catch_var_id,
+                dst: io_reg,
             }
             .into_spanned(call.head),
         )?
     } else {
-        builder.push(Instruction::PushErrorHandler { index: usize::MAX }.into_spanned(call.head))?
+        // Otherwise, we don't need the error value.
+        builder.push(Instruction::OnError { index: usize::MAX }.into_spanned(call.head))?
     };
 
     // Compile the block
@@ -242,8 +293,8 @@ pub(crate) fn compile_try(
     builder.push(Instruction::PopErrorHandler.into_spanned(call.head))?;
 
     // Jump over the failure case
-    let jump_index =
+    let catch_span = catch_expr.map(|e| e.span).unwrap_or(call.head);
-        builder.jump_placeholder(catch_block.and_then(|b| b.span).unwrap_or(call.head))?;
+    let jump_index = builder.jump_placeholder(catch_span)?;
 
     // This is the error handler - go back and set the right branch destination
     builder.set_branch_target(error_handler_index, builder.next_instruction_index())?;
@@ -251,19 +302,54 @@ pub(crate) fn compile_try(
     // Mark out register as likely not clean - state in error handler is not well defined
     builder.mark_register(io_reg)?;
 
-    // If we have a catch block, compile that
+    // Now compile whatever is necessary for the error handler
-    if let Some(catch_block) = catch_block {
+    match catch_type {
-        compile_block(
+        Some(CatchType::Block { block, var_id }) => {
-            working_set,
+            if let Some(var_id) = var_id {
-            builder,
+                // Error will be in io_reg
-            catch_block,
+                builder.mark_register(io_reg)?;
-            redirect_modes,
+                builder.push(
-            None,
+                    Instruction::StoreVariable {
-            io_reg,
+                        var_id,
-        )?;
+                        src: io_reg,
-    } else {
+                    }
-        // Otherwise just set out to empty.
+                    .into_spanned(catch_span),
-        builder.load_empty(io_reg)?;
+                )?;
+            }
+            // Compile the block, now that the variable is set
+            compile_block(working_set, builder, block, redirect_modes, None, io_reg)?;
+        }
+        Some(CatchType::Closure { closure_reg }) => {
+            // We should call `do`. Error will be in io_reg
+            let do_decl_id = working_set.find_decl(b"do").ok_or_else(|| {
+                CompileError::MissingRequiredDeclaration {
+                    decl_name: "do".into(),
+                    span: call.head,
+                }
+            })?;
+            builder.mark_register(io_reg)?;
+
+            // Push the closure and the error
+            builder
+                .push(Instruction::PushPositional { src: closure_reg }.into_spanned(catch_span))?;
+            builder.push(Instruction::PushPositional { src: io_reg }.into_spanned(catch_span))?;
+
+            // Empty input to the block
+            builder.load_empty(io_reg)?;
+
+            // Call `do $closure $err`
+            builder.push(
+                Instruction::Call {
+                    decl_id: do_decl_id,
+                    src_dst: io_reg,
+                }
+                .into_spanned(catch_span),
+            )?;
+        }
+        None => {
+            // Just set out to empty.
+            builder.load_empty(io_reg)?;
+        }
     }
 
     // This is the end - if we succeeded, should jump here

crates/nu-engine/src/eval_ir.rs

@@ -6,8 +6,8 @@ use nu_protocol::{
     debugger::DebugContext,
     engine::{Argument, Closure, EngineState, ErrorHandler, Redirection, Stack},
     ir::{Call, DataSlice, Instruction, IrBlock, Literal, RedirectMode},
-    DeclId, IntoPipelineData, IntoSpanned, ListStream, OutDest, PipelineData, Range, Record, RegId,
+    record, DeclId, IntoPipelineData, IntoSpanned, ListStream, OutDest, PipelineData, Range,
-    ShellError, Span, Value, VarId,
+    Record, RegId, ShellError, Span, Spanned, Value, VarId,
 };
 
 use crate::eval::is_automatic_env_var;
@@ -141,7 +141,7 @@ fn eval_ir_block_impl<D: DebugContext>(
             Err(err) => {
                 if let Some(error_handler) = ctx.stack.error_handlers.pop(ctx.error_handler_base) {
                     // If an error handler is set, branch there
-                    error_handler.prepare_stack(ctx.engine_state, &mut ctx.stack, err);
+                    prepare_error_handler(ctx, error_handler, err.into_spanned(*span));
                     pc = error_handler.handler_index;
                 } else {
                     // If not, exit the block with the error
@@ -161,6 +161,26 @@ fn eval_ir_block_impl<D: DebugContext>(
     })
 }
 
+/// Prepare the context for an error handler
+fn prepare_error_handler(
+    ctx: &mut EvalContext<'_>,
+    error_handler: ErrorHandler,
+    error: Spanned<ShellError>,
+) {
+    if let Some(reg_id) = error_handler.error_register {
+        // Create the error value and put it in the register
+        let value = Value::record(
+            record! {
+                "msg" => Value::string(format!("{}", error.item), error.span),
+                "debug" => Value::string(format!("{:?}", error.item), error.span),
+                "raw" => Value::error(error.item, error.span),
+            },
+            error.span,
+        );
+        ctx.put_reg(reg_id, PipelineData::Value(value, None));
+    }
+}
+
 /// The result of performing an instruction. Describes what should happen next
 #[derive(Debug)]
 enum InstructionResult {
@@ -485,17 +505,17 @@ fn eval_instruction(
             stream,
             end_index,
         } => eval_iterate(ctx, *dst, *stream, *end_index),
-        Instruction::PushErrorHandler { index } => {
+        Instruction::OnError { index } => {
             ctx.stack.error_handlers.push(ErrorHandler {
                 handler_index: *index,
-                error_variable: None,
+                error_register: None,
             });
             Ok(Continue)
         }
-        Instruction::PushErrorHandlerVar { index, error_var } => {
+        Instruction::OnErrorInto { index, dst } => {
             ctx.stack.error_handlers.push(ErrorHandler {
                 handler_index: *index,
-                error_variable: Some(*error_var),
+                error_register: Some(*dst),
             });
             Ok(Continue)
         }

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

@@ -1,32 +1,12 @@
-use crate::{record, ShellError, Value, VarId};
+use crate::RegId;
-
-use super::{EngineState, Stack};
 
 /// Describes an error handler stored during IR evaluation.
 #[derive(Debug, Clone, Copy)]
 pub struct ErrorHandler {
     /// Instruction index within the block that will handle the error
     pub handler_index: usize,
-    /// Variable to put the error information into, when an error occurs
+    /// Register to put the error information into, when an error occurs
-    pub error_variable: Option<VarId>,
+    pub error_register: Option<RegId>,
-}
-
-impl ErrorHandler {
-    /// Add `error_variable` to the stack with the error value.
-    pub fn prepare_stack(&self, engine_state: &EngineState, stack: &mut Stack, error: ShellError) {
-        if let Some(var_id) = self.error_variable {
-            let span = engine_state.get_var(var_id).declaration_span;
-            let value = Value::record(
-                record! {
-                    "msg" => Value::string(format!("{}", error), span),
-                    "debug" => Value::string(format!("{:?}", error), span),
-                    "raw" => Value::error(error, span),
-                },
-                span,
-            );
-            stack.add_var(var_id, value);
-        }
-    }
 }
 
 /// Keeps track of error handlers pushed during evaluation of an IR block.

crates/nu-protocol/src/errors/compile_error.rs

@@ -107,6 +107,13 @@ pub enum CompileError {
         span: Span,
     },
 
+    #[error("Missing required declaration: `{decl_name}`")]
+    MissingRequiredDeclaration {
+        decl_name: String,
+        #[label("`{decl_name}` must be in scope to compile this expression")]
+        span: Span,
+    },
+
     #[error("TODO: {msg}")]
     #[diagnostic(code(nu::compile::todo), help("IR compilation is a work in progress"))]
     Todo {

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

@@ -46,7 +46,7 @@ pub struct FmtInstruction<'a> {
 
 impl<'a> fmt::Display for FmtInstruction<'a> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        const WIDTH: usize = 22;
+        const WIDTH: usize = 20;
 
         match self.instruction {
             Instruction::LoadLiteral { dst, lit } => {
@@ -170,16 +170,11 @@ impl<'a> fmt::Display for FmtInstruction<'a> {
             } => {
                 write!(f, "{:WIDTH$} {dst}, {stream}, end {end_index}", "iterate")
             }
-            Instruction::PushErrorHandler { index } => {
+            Instruction::OnError { index } => {
-                write!(f, "{:WIDTH$} {index}", "push-error-handler")
+                write!(f, "{:WIDTH$} {index}", "on-error")
             }
-            Instruction::PushErrorHandlerVar { index, error_var } => {
+            Instruction::OnErrorInto { index, dst } => {
-                let error_var = FmtVar::new(self.engine_state, *error_var);
+                write!(f, "{:WIDTH$} {index}, {dst}", "on-error-into")
-                write!(
-                    f,
-                    "{:WIDTH$} {index}, {error_var}",
-                    "push-error-handler-var"
-                )
             }
             Instruction::PopErrorHandler => {
                 write!(f, "{:WIDTH$}", "pop-error-handler")

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

@@ -150,9 +150,10 @@ pub enum Instruction {
         end_index: usize,
     },
     /// Push an error handler, without capturing the error value
-    PushErrorHandler { index: usize },
+    OnError { index: usize },
-    /// Push an error handler, capturing the error value into the `error_var`
+    /// Push an error handler, capturing the error value into `dst`. If the error handler is not
-    PushErrorHandlerVar { index: usize, error_var: VarId },
+    /// called, the register should be freed manually.
+    OnErrorInto { index: usize, dst: RegId },
     /// Pop an error handler. This is not necessary when control flow is directed to the error
     /// handler due to an error.
     PopErrorHandler,

tests/eval/mod.rs

@@ -359,3 +359,25 @@ fn try_catch_var() {
         Eq("foo"),
     )
 }
+
+#[test]
+fn try_catch_with_non_literal_closure_no_var() {
+    test_eval(
+        r#"
+            let error_handler = { || "pass" }
+            try { error make { msg: foobar } } catch $error_handler
+        "#,
+        Eq("pass"),
+    )
+}
+
+#[test]
+fn try_catch_with_non_literal_closure() {
+    test_eval(
+        r#"
+            let error_handler = { |err| $err.msg }
+            try { error make { msg: foobar } } catch $error_handler
+        "#,
+        Eq("foobar"),
+    )
+}