Rollup merge of #122076 - WaffleLapkin:mplace-args, r=RalfJung

Tweak the way we protect in-place function arguments in interpreters

Use `MPlaceTy` instead of `PlaceTy` in `FnArg` and ignore (copy) locals in an earlier step ("Locals that don't have their address taken are as protected as they can ever be").

This seems to be crucial for tail call support (as they can't refer to caller's locals which are killed when replacing the stack frame).

r? `@RalfJung`
cc `@oli-obk`
see #121273 (comment)

Author: matthiaskrgr

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -227,7 +227,7 @@ impl<'mir, 'tcx: 'mir> CompileTimeEvalContext<'mir, 'tcx> {
         if self.tcx.has_attr(def_id, sym::rustc_const_panic_str)
             || Some(def_id) == self.tcx.lang_items().begin_panic_fn()
         {
-            let args = self.copy_fn_args(args)?;
+            let args = self.copy_fn_args(args);
             // &str or &&str
             assert!(args.len() == 1);
 
@@ -254,7 +254,7 @@ impl<'mir, 'tcx: 'mir> CompileTimeEvalContext<'mir, 'tcx> {
 
             return Ok(Some(new_instance));
         } else if Some(def_id) == self.tcx.lang_items().align_offset_fn() {
-            let args = self.copy_fn_args(args)?;
+            let args = self.copy_fn_args(args);
             // For align_offset, we replace the function call if the pointer has no address.
             match self.align_offset(instance, &args, dest, ret)? {
                 ControlFlow::Continue(()) => return Ok(Some(instance)),

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -472,11 +472,11 @@ pub trait Machine<'mir, 'tcx: 'mir>: Sized {
     /// argument/return value was actually copied or passed in-place..
     fn protect_in_place_function_argument(
         ecx: &mut InterpCx<'mir, 'tcx, Self>,
-        place: &PlaceTy<'tcx, Self::Provenance>,
+        mplace: &MPlaceTy<'tcx, Self::Provenance>,
     ) -> InterpResult<'tcx> {
         // Without an aliasing model, all we can do is put `Uninit` into the place.
         // Conveniently this also ensures that the place actually points to suitable memory.
-        ecx.write_uninit(place)
+        ecx.write_uninit(mplace)
     }
 
     /// Called immediately before a new stack frame gets pushed.

compiler/rustc_const_eval/src/interpret/terminator.rs

@@ -1,5 +1,7 @@
 use std::borrow::Cow;
 
+use either::Either;
+
 use rustc_middle::{
     mir,
     ty::{
@@ -29,28 +31,25 @@ pub enum FnArg<'tcx, Prov: Provenance = CtfeProvenance> {
     Copy(OpTy<'tcx, Prov>),
     /// Allow for the argument to be passed in-place: destroy the value originally stored at that place and
     /// make the place inaccessible for the duration of the function call.
-    InPlace(PlaceTy<'tcx, Prov>),
+    InPlace(MPlaceTy<'tcx, Prov>),
 }
 
 impl<'tcx, Prov: Provenance> FnArg<'tcx, Prov> {
     pub fn layout(&self) -> &TyAndLayout<'tcx> {
         match self {
             FnArg::Copy(op) => &op.layout,
-            FnArg::InPlace(place) => &place.layout,
+            FnArg::InPlace(mplace) => &mplace.layout,
         }
     }
 }
 
 impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, '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>,
-    ) -> InterpResult<'tcx, OpTy<'tcx, M::Provenance>> {
+    pub fn copy_fn_arg(&self, arg: &FnArg<'tcx, M::Provenance>) -> OpTy<'tcx, M::Provenance> {
         match arg {
-            FnArg::Copy(op) => Ok(op.clone()),
-            FnArg::InPlace(place) => self.place_to_op(place),
+            FnArg::Copy(op) => op.clone(),
+            FnArg::InPlace(mplace) => mplace.clone().into(),
         }
     }
 
@@ -59,7 +58,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
     pub fn copy_fn_args(
         &self,
         args: &[FnArg<'tcx, M::Provenance>],
-    ) -> InterpResult<'tcx, Vec<OpTy<'tcx, M::Provenance>>> {
+    ) -> Vec<OpTy<'tcx, M::Provenance>> {
         args.iter().map(|fn_arg| self.copy_fn_arg(fn_arg)).collect()
     }
 
@@ -70,7 +69,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
     ) -> InterpResult<'tcx, FnArg<'tcx, M::Provenance>> {
         Ok(match arg {
             FnArg::Copy(op) => FnArg::Copy(self.project_field(op, field)?),
-            FnArg::InPlace(place) => FnArg::InPlace(self.project_field(place, field)?),
+            FnArg::InPlace(mplace) => FnArg::InPlace(self.project_field(mplace, field)?),
         })
     }
 
@@ -238,10 +237,36 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
     ) -> InterpResult<'tcx, Vec<FnArg<'tcx, M::Provenance>>> {
         ops.iter()
             .map(|op| {
-                Ok(match &op.node {
-                    mir::Operand::Move(place) => FnArg::InPlace(self.eval_place(*place)?),
-                    _ => FnArg::Copy(self.eval_operand(&op.node, None)?),
-                })
+                let arg = match &op.node {
+                    mir::Operand::Copy(_) | mir::Operand::Constant(_) => {
+                        // Make a regular copy.
+                        let op = self.eval_operand(&op.node, None)?;
+                        FnArg::Copy(op)
+                    }
+                    mir::Operand::Move(place) => {
+                        // If this place lives in memory, preserve its location.
+                        // We call `place_to_op` which will be an `MPlaceTy` whenever there exists
+                        // an mplace for this place. (This is in contrast to `PlaceTy::as_mplace_or_local`
+                        // which can return a local even if that has an mplace.)
+                        let place = self.eval_place(*place)?;
+                        let op = self.place_to_op(&place)?;
+
+                        match op.as_mplace_or_imm() {
+                            Either::Left(mplace) => FnArg::InPlace(mplace),
+                            Either::Right(_imm) => {
+                                // This argument doesn't live in memory, so there's no place
+                                // to make inaccessible during the call.
+                                // We rely on there not being any stray `PlaceTy` that would let the
+                                // caller directly access this local!
+                                // This is also crucial for tail calls, where we want the `FnArg` to
+                                // stay valid when the old stack frame gets popped.
+                                FnArg::Copy(op)
+                            }
+                        }
+                    }
+                };
+
+                Ok(arg)
             })
             .collect()
     }
@@ -451,7 +476,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, '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 = self.copy_fn_arg(caller_arg);
         if !already_live {
             let local = callee_arg.as_local().unwrap();
             let meta = caller_arg_copy.meta();
@@ -469,8 +494,8 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         // specifically.)
         self.copy_op_allow_transmute(&caller_arg_copy, &callee_arg)?;
         // If this was an in-place pass, protect the place it comes from for the duration of the call.
-        if let FnArg::InPlace(place) = caller_arg {
-            M::protect_in_place_function_argument(self, place)?;
+        if let FnArg::InPlace(mplace) = caller_arg {
+            M::protect_in_place_function_argument(self, mplace)?;
         }
         Ok(())
     }
@@ -517,7 +542,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 M::call_intrinsic(
                     self,
                     instance,
-                    &self.copy_fn_args(args)?,
+                    &self.copy_fn_args(args),
                     destination,
                     target,
                     unwind,
@@ -594,8 +619,8 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                             .map(|arg| (
                                 arg.layout().ty,
                                 match arg {
-                                    FnArg::Copy(op) => format!("copy({:?})", *op),
-                                    FnArg::InPlace(place) => format!("in-place({:?})", *place),
+                                    FnArg::Copy(op) => format!("copy({op:?})"),
+                                    FnArg::InPlace(mplace) => format!("in-place({mplace:?})"),
                                 }
                             ))
                             .collect::<Vec<_>>()
@@ -717,8 +742,9 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                             callee_ty: callee_fn_abi.ret.layout.ty
                         });
                     }
+
                     // Protect return place for in-place return value passing.
-                    M::protect_in_place_function_argument(self, &destination.clone().into())?;
+                    M::protect_in_place_function_argument(self, &destination)?;
 
                     // Don't forget to mark "initially live" locals as live.
                     self.storage_live_for_always_live_locals()?;
@@ -741,7 +767,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, '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 = self.copy_fn_arg(&args[0]);
                 let receiver_place = loop {
                     match receiver.layout.ty.kind() {
                         ty::Ref(..) | ty::RawPtr(..) => {

src/tools/miri/src/machine.rs

@@ -8,7 +8,6 @@ use std::fmt;
 use std::path::Path;
 use std::process;
 
-use either::Either;
 use rand::rngs::StdRng;
 use rand::Rng;
 use rand::SeedableRng;
@@ -962,7 +961,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, '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 = ecx.copy_fn_args(args); // FIXME: Should `InPlace` arguments be reset to uninit?
             let link_name = ecx.item_link_name(instance.def_id());
             return ecx.emulate_foreign_item(link_name, abi, &args, dest, ret, unwind);
         }
@@ -981,7 +980,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, '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 = ecx.copy_fn_args(args); // FIXME: Should `InPlace` arguments be reset to uninit?
         ecx.emulate_dyn_sym(fn_val, abi, &args, dest, ret, unwind)
     }
 
@@ -1334,18 +1333,12 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
 
     fn protect_in_place_function_argument(
         ecx: &mut InterpCx<'mir, 'tcx, Self>,
-        place: &PlaceTy<'tcx, Provenance>,
+        place: &MPlaceTy<'tcx, Provenance>,
     ) -> InterpResult<'tcx> {
         // If we have a borrow tracker, we also have it set up protection so that all reads *and
         // writes* during this call are insta-UB.
         let protected_place = if ecx.machine.borrow_tracker.is_some() {
-            // Have to do `to_op` first because a `Place::Local` doesn't imply the local doesn't have an address.
-            if let Either::Left(place) = ecx.place_to_op(place)?.as_mplace_or_imm() {
-                ecx.protect_place(&place)?.into()
-            } else {
-                // Locals that don't have their address taken are as protected as they can ever be.
-                place.clone()
-            }
+            ecx.protect_place(&place)?.into()
         } else {
             // No borrow tracker.
             place.clone()