Rollup merge of rust-lang#150473 - RalfJung:interpret-tail-call, r=WaffleLapkin

tail calls: fix copying non-scalar arguments to callee

Alternative to rust-lang#144933: when invoking a tail call with a non-scalar argument, we need to delay freeing the caller's local variables until after the callee is initialized, so that we can copy things from the caller to the callee.

Fixes rust-lang#144820... but as the FIXMEs in the code show, it's not clear to me whether these are the right semantics.
r? @WaffleLapkin

Author: jhpratt

compiler/rustc_const_eval/src/const_eval/machine.rs

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

@@ -19,8 +19,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,
+    Projectable, Provenance, ReturnAction, ReturnContinuation, Scalar, interp_ok, throw_ub,
+    throw_ub_custom,
 };
 use crate::enter_trace_span;
 use crate::interpret::EnteredTraceSpan;
@@ -43,25 +43,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.
@@ -319,7 +316,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();
@@ -616,7 +613,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,
@@ -703,7 +700,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(..) => {
@@ -824,41 +821,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,#144855): 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,#144855): 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(
@@ -953,14 +959,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.
@@ -978,7 +988,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)
@@ -989,7 +999,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

@@ -81,7 +81,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...]`.
@@ -127,19 +127,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> {
@@ -292,6 +279,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))
@@ -417,35 +412,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.
@@ -455,7 +441,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             ReturnContinuation::Stop { cleanup, .. } => cleanup,
         };
 
-        let return_action = if cleanup {
+        if cleanup {
             for local in &frame.locals {
                 self.deallocate_local(local.value)?;
             }
@@ -466,13 +452,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
@@ -655,7 +639,7 @@ impl<'a, 'tcx: 'a, M: Machine<'tcx>> InterpCx<'tcx, M> {
             let (_idx, callee_abi) = callee_abis.next().unwrap();
             assert!(self.check_argument_compat(caller_abi, callee_abi)?);
             // FIXME: do we have to worry about in-place argument passing?
-            let op = self.copy_fn_arg(fn_arg);
+            let op = fn_arg.copy_fn_arg();
             let mplace = self.allocate(op.layout, MemoryKind::Stack)?;
             self.copy_op(&op, &mplace)?;
 

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

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

@@ -1235,7 +1235,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);
         }
@@ -1262,7 +1262,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: stack backtrace:
            0: g
                at tests/fail/tail_calls/dangling-local-var.rs:LL:CC

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);
+}