tail calls: fix copying non-scalar arguments to callee

Author: RalfJung

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -235,7 +235,7 @@ impl<'tcx> CompileTimeInterpCx<'tcx> {
         if self.tcx.is_lang_item(def_id, LangItem::PanicDisplay)
             || self.tcx.is_lang_item(def_id, LangItem::BeginPanic)
         {
-            let args = self.copy_fn_args(args);
+            let args = Self::copy_fn_args(args);
             // &str or &&str
             assert!(args.len() == 1);
 

compiler/rustc_const_eval/src/interpret/call.rs

@@ -16,8 +16,8 @@ use tracing::{info, instrument, trace};
 
 use super::{
     CtfeProvenance, FnVal, ImmTy, InterpCx, InterpResult, MPlaceTy, Machine, OpTy, PlaceTy,
-    Projectable, Provenance, ReturnAction, ReturnContinuation, Scalar, StackPopInfo, interp_ok,
-    throw_ub, throw_ub_custom, throw_unsup_format,
+    Projectable, Provenance, ReturnAction, ReturnContinuation, Scalar, interp_ok, throw_ub,
+    throw_ub_custom, throw_unsup_format,
 };
 use crate::interpret::EnteredTraceSpan;
 use crate::{enter_trace_span, fluent_generated as fluent};
@@ -40,25 +40,22 @@ impl<'tcx, Prov: Provenance> FnArg<'tcx, Prov> {
             FnArg::InPlace(mplace) => &mplace.layout,
         }
     }
-}
 
-impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
     /// Make a copy of the given fn_arg. Any `InPlace` are degenerated to copies, no protection of the
     /// original memory occurs.
-    pub fn copy_fn_arg(&self, arg: &FnArg<'tcx, M::Provenance>) -> OpTy<'tcx, M::Provenance> {
-        match arg {
+    pub fn copy_fn_arg(&self) -> OpTy<'tcx, Prov> {
+        match self {
             FnArg::Copy(op) => op.clone(),
             FnArg::InPlace(mplace) => mplace.clone().into(),
         }
     }
+}
 
+impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
     /// Make a copy of the given fn_args. Any `InPlace` are degenerated to copies, no protection of the
     /// original memory occurs.
-    pub fn copy_fn_args(
-        &self,
-        args: &[FnArg<'tcx, M::Provenance>],
-    ) -> Vec<OpTy<'tcx, M::Provenance>> {
-        args.iter().map(|fn_arg| self.copy_fn_arg(fn_arg)).collect()
+    pub fn copy_fn_args(args: &[FnArg<'tcx, M::Provenance>]) -> Vec<OpTy<'tcx, M::Provenance>> {
+        args.iter().map(|fn_arg| fn_arg.copy_fn_arg()).collect()
     }
 
     /// Helper function for argument untupling.
@@ -310,7 +307,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         // We work with a copy of the argument for now; if this is in-place argument passing, we
         // will later protect the source it comes from. This means the callee cannot observe if we
         // did in-place of by-copy argument passing, except for pointer equality tests.
-        let caller_arg_copy = self.copy_fn_arg(caller_arg);
+        let caller_arg_copy = caller_arg.copy_fn_arg();
         if !already_live {
             let local = callee_arg.as_local().unwrap();
             let meta = caller_arg_copy.meta();
@@ -559,7 +556,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 if let Some(fallback) = M::call_intrinsic(
                     self,
                     instance,
-                    &self.copy_fn_args(args),
+                    &Self::copy_fn_args(args),
                     destination,
                     target,
                     unwind,
@@ -646,7 +643,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 // An `InPlace` does nothing here, we keep the original receiver intact. We can't
                 // really pass the argument in-place anyway, and we are constructing a new
                 // `Immediate` receiver.
-                let mut receiver = self.copy_fn_arg(&args[0]);
+                let mut receiver = args[0].copy_fn_arg();
                 let receiver_place = loop {
                     match receiver.layout.ty.kind() {
                         ty::Ref(..) | ty::RawPtr(..) => {
@@ -765,41 +762,50 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         with_caller_location: bool,
     ) -> InterpResult<'tcx> {
         trace!("init_fn_tail_call: {:#?}", fn_val);
-
         // This is the "canonical" implementation of tails calls,
         // a pop of the current stack frame, followed by a normal call
         // which pushes a new stack frame, with the return address from
         // the popped stack frame.
         //
-        // Note that we are using `pop_stack_frame_raw` and not `return_from_current_stack_frame`,
-        // as the latter "executes" the goto to the return block, but we don't want to,
+        // Note that we cannot use `return_from_current_stack_frame`,
+        // as that "executes" the goto to the return block, but we don't want to,
         // only the tail called function should return to the current return block.
-        let StackPopInfo { return_action, return_cont, return_place } =
-            self.pop_stack_frame_raw(false, |_this, _return_place| {
-                // This function's return value is just discarded, the tail-callee will fill in the return place instead.
-                interp_ok(())
-            })?;
 
-        assert_eq!(return_action, ReturnAction::Normal);
-
-        // Take the "stack pop cleanup" info, and use that to initiate the next call.
-        let ReturnContinuation::Goto { ret, unwind } = return_cont else {
-            bug!("can't tailcall as root");
+        // The arguments need to all be copied since the current stack frame will be removed
+        // before the callee even starts executing.
+        // FIXME(explicit_tail_calls): does this match what codegen does?
+        let args = args.iter().map(|fn_arg| FnArg::Copy(fn_arg.copy_fn_arg())).collect::<Vec<_>>();
+        // Remove the frame from the stack.
+        let frame = self.pop_stack_frame_raw()?;
+        // Remember where this frame would have returned to.
+        let ReturnContinuation::Goto { ret, unwind } = frame.return_cont() else {
+            bug!("can't tailcall as root of the stack");
         };
-
+        // There's no return value to deal with! Instead, we forward the old return place
+        // to the new function.
         // FIXME(explicit_tail_calls):
         //   we should check if both caller&callee can/n't unwind,
         //   see <https://github.com/rust-lang/rust/pull/113128#issuecomment-1614979803>
 
+        // Now push the new stack frame.
         self.init_fn_call(
             fn_val,
             (caller_abi, caller_fn_abi),
-            args,
+            &*args,
             with_caller_location,
-            &return_place,
+            frame.return_place(),
             ret,
             unwind,
-        )
+        )?;
+
+        // Finally, clear the local variables. Has to be done after pushing to support
+        // non-scalar arguments.
+        // FIXME(explicit_tail_calls): revisit this once codegen supports indirect arguments,
+        // to ensure the semantics are compatible.
+        let return_action = self.cleanup_stack_frame(/* unwinding */ false, frame)?;
+        assert_eq!(return_action, ReturnAction::Normal);
+
+        interp_ok(())
     }
 
     pub(super) fn init_drop_in_place_call(
@@ -894,14 +900,18 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         // local's value out.
         let return_op =
             self.local_to_op(mir::RETURN_PLACE, None).expect("return place should always be live");
-        // Do the actual pop + copy.
-        let stack_pop_info = self.pop_stack_frame_raw(unwinding, |this, return_place| {
-            this.copy_op_allow_transmute(&return_op, return_place)?;
-            trace!("return value: {:?}", this.dump_place(return_place));
-            interp_ok(())
-        })?;
-
-        match stack_pop_info.return_action {
+        // Remove the frame from the stack.
+        let frame = self.pop_stack_frame_raw()?;
+        // Copy the return value and remember the return continuation.
+        if !unwinding {
+            self.copy_op_allow_transmute(&return_op, frame.return_place())?;
+            trace!("return value: {:?}", self.dump_place(frame.return_place()));
+        }
+        let return_cont = frame.return_cont();
+        // Finish popping the stack frame.
+        let return_action = self.cleanup_stack_frame(unwinding, frame)?;
+        // Jump to the next block.
+        match return_action {
             ReturnAction::Normal => {}
             ReturnAction::NoJump => {
                 // The hook already did everything.
@@ -919,7 +929,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         // Normal return, figure out where to jump.
         if unwinding {
             // Follow the unwind edge.
-            match stack_pop_info.return_cont {
+            match return_cont {
                 ReturnContinuation::Goto { unwind, .. } => {
                     // This must be the very last thing that happens, since it can in fact push a new stack frame.
                     self.unwind_to_block(unwind)
@@ -930,7 +940,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             }
         } else {
             // Follow the normal return edge.
-            match stack_pop_info.return_cont {
+            match return_cont {
                 ReturnContinuation::Goto { ret, .. } => self.return_to_block(ret),
                 ReturnContinuation::Stop { .. } => {
                     assert!(

compiler/rustc_const_eval/src/interpret/mod.rs

@@ -36,7 +36,7 @@ pub use self::operand::{ImmTy, Immediate, OpTy};
 pub use self::place::{MPlaceTy, MemPlaceMeta, PlaceTy, Writeable};
 use self::place::{MemPlace, Place};
 pub use self::projection::{OffsetMode, Projectable};
-pub use self::stack::{Frame, FrameInfo, LocalState, ReturnContinuation, StackPopInfo};
+pub use self::stack::{Frame, FrameInfo, LocalState, ReturnContinuation};
 pub use self::util::EnteredTraceSpan;
 pub(crate) use self::util::create_static_alloc;
 pub use self::validity::{CtfeValidationMode, RangeSet, RefTracking};

compiler/rustc_const_eval/src/interpret/stack.rs

@@ -80,7 +80,7 @@ pub struct Frame<'tcx, Prov: Provenance = CtfeProvenance, Extra = ()> {
     /// and its layout in the caller. This place is to be interpreted relative to the
     /// *caller's* stack frame. We use a `PlaceTy` instead of an `MPlaceTy` since this
     /// avoids having to move *all* return places into Miri's memory.
-    pub return_place: PlaceTy<'tcx, Prov>,
+    return_place: PlaceTy<'tcx, Prov>,
 
     /// The list of locals for this stack frame, stored in order as
     /// `[return_ptr, arguments..., variables..., temporaries...]`.
@@ -122,19 +122,6 @@ pub enum ReturnContinuation {
     Stop { cleanup: bool },
 }
 
-/// Return type of [`InterpCx::pop_stack_frame_raw`].
-pub struct StackPopInfo<'tcx, Prov: Provenance> {
-    /// Additional information about the action to be performed when returning from the popped
-    /// stack frame.
-    pub return_action: ReturnAction,
-
-    /// [`return_cont`](Frame::return_cont) of the popped stack frame.
-    pub return_cont: ReturnContinuation,
-
-    /// [`return_place`](Frame::return_place) of the popped stack frame.
-    pub return_place: PlaceTy<'tcx, Prov>,
-}
-
 /// State of a local variable including a memoized layout
 #[derive(Clone)]
 pub struct LocalState<'tcx, Prov: Provenance = CtfeProvenance> {
@@ -286,6 +273,14 @@ impl<'tcx, Prov: Provenance, Extra> Frame<'tcx, Prov, Extra> {
         self.instance
     }
 
+    pub fn return_place(&self) -> &PlaceTy<'tcx, Prov> {
+        &self.return_place
+    }
+
+    pub fn return_cont(&self) -> ReturnContinuation {
+        self.return_cont
+    }
+
     /// Return the `SourceInfo` of the current instruction.
     pub fn current_source_info(&self) -> Option<&mir::SourceInfo> {
         self.loc.left().map(|loc| self.body.source_info(loc))
@@ -410,35 +405,26 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         interp_ok(())
     }
 
-    /// Low-level helper that pops a stack frame from the stack and returns some information about
-    /// it.
-    ///
-    /// This also deallocates locals, if necessary.
-    /// `copy_ret_val` gets called after the frame has been taken from the stack but before the locals have been deallocated.
-    ///
-    /// [`M::before_stack_pop`] and [`M::after_stack_pop`] are called by this function
-    /// automatically.
-    ///
-    /// The high-level version of this is `return_from_current_stack_frame`.
-    ///
-    /// [`M::before_stack_pop`]: Machine::before_stack_pop
-    /// [`M::after_stack_pop`]: Machine::after_stack_pop
+    /// Low-level helper that pops a stack frame from the stack without any cleanup.
+    /// This invokes `before_stack_pop`.
+    /// After calling this function, you need to deal with the return value, and then
+    /// invoke `cleanup_stack_frame`.
     pub(super) fn pop_stack_frame_raw(
         &mut self,
-        unwinding: bool,
-        copy_ret_val: impl FnOnce(&mut Self, &PlaceTy<'tcx, M::Provenance>) -> InterpResult<'tcx>,
-    ) -> InterpResult<'tcx, StackPopInfo<'tcx, M::Provenance>> {
+    ) -> InterpResult<'tcx, Frame<'tcx, M::Provenance, M::FrameExtra>> {
         M::before_stack_pop(self)?;
         let frame =
             self.stack_mut().pop().expect("tried to pop a stack frame, but there were none");
+        interp_ok(frame)
+    }
 
-        // Copy return value (unless we are unwinding).
-        if !unwinding {
-            copy_ret_val(self, &frame.return_place)?;
-        }
-
+    /// Deallocate local variables in the stack frame, and invoke `after_stack_pop`.
+    pub(super) fn cleanup_stack_frame(
+        &mut self,
+        unwinding: bool,
+        frame: Frame<'tcx, M::Provenance, M::FrameExtra>,
+    ) -> InterpResult<'tcx, ReturnAction> {
         let return_cont = frame.return_cont;
-        let return_place = frame.return_place.clone();
 
         // Cleanup: deallocate locals.
         // Usually we want to clean up (deallocate locals), but in a few rare cases we don't.
@@ -448,7 +434,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             ReturnContinuation::Stop { cleanup, .. } => cleanup,
         };
 
-        let return_action = if cleanup {
+        if cleanup {
             // We need to take the locals out, since we need to mutate while iterating.
             for local in &frame.locals {
                 self.deallocate_local(local.value)?;
@@ -457,13 +443,11 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             // Call the machine hook, which determines the next steps.
             let return_action = M::after_stack_pop(self, frame, unwinding)?;
             assert_ne!(return_action, ReturnAction::NoCleanup);
-            return_action
+            interp_ok(return_action)
         } else {
             // We also skip the machine hook when there's no cleanup. This not a real "pop" anyway.
-            ReturnAction::NoCleanup
-        };
-
-        interp_ok(StackPopInfo { return_action, return_cont, return_place })
+            interp_ok(ReturnAction::NoCleanup)
+        }
     }
 
     /// In the current stack frame, mark all locals as live that are not arguments and don't have

src/tools/miri/src/concurrency/thread.rs

@@ -343,8 +343,8 @@ impl VisitProvenance for Thread<'_> {
 
 impl VisitProvenance for Frame<'_, Provenance, FrameExtra<'_>> {
     fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        let return_place = self.return_place();
         let Frame {
-            return_place,
             locals,
             extra,
             // There are some private fields we cannot access; they contain no tags.

src/tools/miri/src/diagnostics.rs

@@ -497,7 +497,7 @@ pub fn report_result<'tcx>(
     for (i, frame) in ecx.active_thread_stack().iter().enumerate() {
         trace!("-------------------");
         trace!("Frame {}", i);
-        trace!("    return: {:?}", frame.return_place);
+        trace!("    return: {:?}", frame.return_place());
         for (i, local) in frame.locals.iter().enumerate() {
             trace!("    local {}: {:?}", i, local);
         }

src/tools/miri/src/machine.rs

@@ -1228,7 +1228,7 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> {
             // to run extra MIR), and Ok(Some(body)) if we found MIR to run for the
             // foreign function
             // Any needed call to `goto_block` will be performed by `emulate_foreign_item`.
-            let args = ecx.copy_fn_args(args); // FIXME: Should `InPlace` arguments be reset to uninit?
+            let args = MiriInterpCx::copy_fn_args(args); // FIXME: Should `InPlace` arguments be reset to uninit?
             let link_name = Symbol::intern(ecx.tcx.symbol_name(instance).name);
             return ecx.emulate_foreign_item(link_name, abi, &args, dest, ret, unwind);
         }
@@ -1255,7 +1255,7 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> {
         ret: Option<mir::BasicBlock>,
         unwind: mir::UnwindAction,
     ) -> InterpResult<'tcx> {
-        let args = ecx.copy_fn_args(args); // FIXME: Should `InPlace` arguments be reset to uninit?
+        let args = MiriInterpCx::copy_fn_args(args); // FIXME: Should `InPlace` arguments be reset to uninit?
         ecx.emulate_dyn_sym(fn_val, abi, &args, dest, ret, unwind)
     }
 

src/tools/miri/tests/fail/tail_calls/dangling-local-var.stderr

@@ -14,8 +14,8 @@ LL |     let local = 0;
 help: ALLOC was deallocated here:
   --> tests/fail/tail_calls/dangling-local-var.rs:LL:CC
    |
-LL |     f(std::ptr::null());
-   |     ^^^^^^^^^^^^^^^^^^^
+LL |     let _val = unsafe { *x };
+   |         ^^^^
    = note: BACKTRACE (of the first span):
    = note: inside `g` at tests/fail/tail_calls/dangling-local-var.rs:LL:CC
 note: inside `main`

src/tools/miri/tests/pass/tail_call.rs …i/tests/pass/function_calls/tail_call.rssrc/tools/miri/tests/pass/tail_call.rs renamed to src/tools/miri/tests/pass/function_calls/tail_call.rs src/tools/miri/tests/pass/tail_call.rs renamed to src/tools/miri/tests/pass/function_calls/tail_call.rs

@@ -4,6 +4,7 @@
 fn main() {
     assert_eq!(factorial(10), 3_628_800);
     assert_eq!(mutually_recursive_identity(1000), 1000);
+    non_scalar();
 }
 
 fn factorial(n: u32) -> u32 {
@@ -37,3 +38,14 @@ fn mutually_recursive_identity(x: u32) -> u32 {
 
     switch(x, 0)
 }
+
+fn non_scalar() {
+    fn f(x: [usize; 2], i: u32) {
+        if i == 0 {
+            return;
+        }
+        become f(x, i - 1);
+    }
+
+    f([5, 5], 2);
+}