retag return places

src/machine.rs

@@ -481,11 +481,10 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for Evaluator<'tcx> {
         kind: mir::RetagKind,
         place: PlaceTy<'tcx, Tag>,
     ) -> InterpResult<'tcx> {
-        if ecx.memory.extra.stacked_borrows.is_none() {
-            // No tracking.
-            Ok(())
-        } else {
+        if ecx.memory.extra.stacked_borrows.is_some() {
             ecx.retag(kind, place)
+        } else {
+            Ok(())
         }
     }
 
@@ -502,6 +501,15 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for Evaluator<'tcx> {
         Ok(frame.with_extra(extra))
     }
 
+    #[inline(always)]
+    fn after_stack_push(ecx: &mut InterpCx<'mir, 'tcx, Self>) -> InterpResult<'tcx> {
+        if ecx.memory.extra.stacked_borrows.is_some() {
+            ecx.retag_return_place()
+        } else {
+            Ok(())
+        }
+    }
+
     #[inline(always)]
     fn after_stack_pop(
         ecx: &mut InterpCx<'mir, 'tcx, Self>,

src/stacked_borrows.rs

@@ -11,7 +11,7 @@ use log::trace;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_middle::mir::RetagKind;
 use rustc_middle::ty;
-use rustc_target::abi::Size;
+use rustc_target::abi::{LayoutOf, Size};
 use rustc_hir::Mutability;
 
 use crate::*;
@@ -569,7 +569,7 @@ trait EvalContextPrivExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
         val: ImmTy<'tcx, Tag>,
         kind: RefKind,
         protect: bool,
-    ) -> InterpResult<'tcx, Immediate<Tag>> {
+    ) -> InterpResult<'tcx, ImmTy<'tcx, Tag>> {
         let this = self.eval_context_mut();
         // We want a place for where the ptr *points to*, so we get one.
         let place = this.ref_to_mplace(val)?;
@@ -582,7 +582,7 @@ trait EvalContextPrivExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
         let place = this.mplace_access_checked(place)?;
         if size == Size::ZERO {
             // Nothing to do for ZSTs.
-            return Ok(*val);
+            return Ok(val);
         }
 
         // Compute new borrow.
@@ -603,7 +603,7 @@ trait EvalContextPrivExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
         let new_place = place.replace_tag(new_tag);
 
         // Return new pointer.
-        Ok(new_place.to_ref())
+        Ok(ImmTy::from_immediate(new_place.to_ref(), val.layout))
     }
 }
 
@@ -640,9 +640,42 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
             // Fast path.
             let val = this.read_immediate(this.place_to_op(place)?)?;
             let val = this.retag_reference(val, mutbl, protector)?;
-            this.write_immediate(val, place)?;
+            this.write_immediate(*val, place)?;
         }
 
         Ok(())
     }
+
+    /// After a stack frame got pushed, retag the return place so that we are sure
+    /// it does not alias with anything.
+    /// 
+    /// This is a HACK because there is nothing in MIR that would make the retag
+    /// explicit. Also see https://github.com/rust-lang/rust/issues/71117.
+    fn retag_return_place(&mut self) -> InterpResult<'tcx> {
+        let this = self.eval_context_mut();
+        let return_place = if let Some(return_place) = this.frame_mut().return_place {
+            return_place
+        } else {
+            // No return place, nothing to do.
+            return Ok(());
+        };
+        if return_place.layout.is_zst() {
+            // There may not be any memory here, nothing to do.
+            return Ok(());
+        }
+        // We need this to be in-memory to use tagged pointers.
+        let return_place = this.force_allocation(return_place)?;
+
+        // We have to turn the place into a pointer to use the existing code.
+        // (The pointer type does not matter, so we use a raw pointer.)
+        let ptr_layout = this.layout_of(this.tcx.mk_mut_ptr(return_place.layout.ty))?;
+        let val = ImmTy::from_immediate(return_place.to_ref(), ptr_layout);
+        // Reborrow it.
+        let val = this.retag_reference(val, RefKind::Unique { two_phase: false }, /*protector*/ true)?;
+        // And use reborrowed pointer for return place.
+        let return_place = this.ref_to_mplace(val)?;
+        this.frame_mut().return_place = Some(return_place.into());
+
+        Ok(())
+    }
 }