From 689cc496635ce2bdf478b421f5ff5ce0f2b09ca2 Mon Sep 17 00:00:00 2001
From: Devyn Cairns <devyn.cairns@gmail.com>
Date: Sun, 7 Jul 2024 20:17:55 -0700
Subject: [PATCH] support calling with the redirection mode of the caller

---
 crates/nu-engine/src/compile/mod.rs      |  9 ++++----
 crates/nu-engine/src/compile/redirect.rs |  7 ++++++
 crates/nu-engine/src/eval_ir.rs          | 28 +++++++++++++++---------
 crates/nu-protocol/src/ir/display.rs     |  1 +
 crates/nu-protocol/src/ir/mod.rs         |  2 ++
 5 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/crates/nu-engine/src/compile/mod.rs b/crates/nu-engine/src/compile/mod.rs
index 7117d71984..faa4d8a3d8 100644
--- a/crates/nu-engine/src/compile/mod.rs
+++ b/crates/nu-engine/src/compile/mod.rs
@@ -25,20 +25,19 @@ const BLOCK_INPUT: RegId = RegId(0);
 pub fn compile(working_set: &StateWorkingSet, block: &Block) -> Result<IrBlock, CompileError> {
     let mut builder = BlockBuilder::new();
 
+    let span = block.span.unwrap_or(Span::unknown());
+
     compile_block(
         working_set,
         &mut builder,
         block,
-        RedirectModes::default(),
+        RedirectModes::caller(span),
         Some(BLOCK_INPUT),
         BLOCK_INPUT,
     )?;
 
     // A complete block has to end with a `return`
-    builder.push(
-        Instruction::Return { src: BLOCK_INPUT }
-            .into_spanned(block.span.unwrap_or(Span::unknown())),
-    )?;
+    builder.push(Instruction::Return { src: BLOCK_INPUT }.into_spanned(span))?;
 
     Ok(builder.finish())
 }
diff --git a/crates/nu-engine/src/compile/redirect.rs b/crates/nu-engine/src/compile/redirect.rs
index 5ed943e5ea..8b470c8ef2 100644
--- a/crates/nu-engine/src/compile/redirect.rs
+++ b/crates/nu-engine/src/compile/redirect.rs
@@ -21,6 +21,13 @@ impl RedirectModes {
         }
     }
 
+    pub(crate) fn caller(span: Span) -> RedirectModes {
+        RedirectModes {
+            out: Some(RedirectMode::Caller.into_spanned(span)),
+            err: Some(RedirectMode::Caller.into_spanned(span)),
+        }
+    }
+
     pub(crate) fn with_pipe_out(&self, span: Span) -> Self {
         RedirectModes {
             out: Some(RedirectMode::Pipe.into_spanned(span)),
diff --git a/crates/nu-engine/src/eval_ir.rs b/crates/nu-engine/src/eval_ir.rs
index a7c1b405aa..4962031d37 100644
--- a/crates/nu-engine/src/eval_ir.rs
+++ b/crates/nu-engine/src/eval_ir.rs
@@ -371,13 +371,11 @@ fn eval_instruction<D: DebugContext>(
             Ok(Continue)
         }
         Instruction::RedirectOut { mode } => {
-            let out_dest = eval_redirection(ctx, mode, *span)?;
-            ctx.redirect_out = Some(out_dest);
+            ctx.redirect_out = eval_redirection(ctx, mode, *span, RedirectionStream::Out)?;
             Ok(Continue)
         }
         Instruction::RedirectErr { mode } => {
-            let out_dest = eval_redirection(ctx, mode, *span)?;
-            ctx.redirect_err = Some(out_dest);
+            ctx.redirect_err = eval_redirection(ctx, mode, *span, RedirectionStream::Err)?;
             Ok(Continue)
         }
         Instruction::Call { decl_id, src_dst } => {
@@ -1093,17 +1091,23 @@ fn collect(data: PipelineData, fallback_span: Span) -> Result<PipelineData, Shel
     Ok(PipelineData::Value(value, metadata))
 }
 
+enum RedirectionStream {
+    Out,
+    Err,
+}
+
 /// Set up a [`Redirection`] from a [`RedirectMode`]
 fn eval_redirection(
     ctx: &mut EvalContext<'_>,
     mode: &RedirectMode,
     span: Span,
-) -> Result<Redirection, ShellError> {
+    which: RedirectionStream,
+) -> Result<Option<Redirection>, ShellError> {
     match mode {
-        RedirectMode::Pipe => Ok(Redirection::Pipe(OutDest::Pipe)),
-        RedirectMode::Capture => Ok(Redirection::Pipe(OutDest::Capture)),
-        RedirectMode::Null => Ok(Redirection::Pipe(OutDest::Null)),
-        RedirectMode::Inherit => Ok(Redirection::Pipe(OutDest::Inherit)),
+        RedirectMode::Pipe => Ok(Some(Redirection::Pipe(OutDest::Pipe))),
+        RedirectMode::Capture => Ok(Some(Redirection::Pipe(OutDest::Capture))),
+        RedirectMode::Null => Ok(Some(Redirection::Pipe(OutDest::Null))),
+        RedirectMode::Inherit => Ok(Some(Redirection::Pipe(OutDest::Inherit))),
         RedirectMode::File { path, append } => {
             let path = ctx.collect_reg(*path, span)?;
             let path_expanded =
@@ -1118,8 +1122,12 @@ fn eval_redirection(
                 .create(true)
                 .open(path_expanded)
                 .map_err(|err| err.into_spanned(span))?;
-            Ok(Redirection::File(file.into()))
+            Ok(Some(Redirection::File(file.into())))
         }
+        RedirectMode::Caller => Ok(match which {
+            RedirectionStream::Out => ctx.stack.pipe_stdout().cloned().map(Redirection::Pipe),
+            RedirectionStream::Err => ctx.stack.pipe_stderr().cloned().map(Redirection::Pipe),
+        }),
     }
 }
 
diff --git a/crates/nu-protocol/src/ir/display.rs b/crates/nu-protocol/src/ir/display.rs
index 4b24543929..3b88844bca 100644
--- a/crates/nu-protocol/src/ir/display.rs
+++ b/crates/nu-protocol/src/ir/display.rs
@@ -254,6 +254,7 @@ impl fmt::Display for RedirectMode {
             RedirectMode::Null => write!(f, "null"),
             RedirectMode::Inherit => write!(f, "inherit"),
             RedirectMode::File { path, append } => write!(f, "file({path}, append={append})"),
+            RedirectMode::Caller => write!(f, "caller"),
         }
     }
 }
diff --git a/crates/nu-protocol/src/ir/mod.rs b/crates/nu-protocol/src/ir/mod.rs
index a202c7b70c..931d21bd88 100644
--- a/crates/nu-protocol/src/ir/mod.rs
+++ b/crates/nu-protocol/src/ir/mod.rs
@@ -300,6 +300,8 @@ pub enum RedirectMode {
         path: RegId,
         append: bool,
     },
+    /// Use the redirection mode requested by the caller, for a pre-return call.
+    Caller,
 }
 
 /// Just a hack to allow `Arc<[u8]>` to be serialized and deserialized