Auto merge of #1330 - RalfJung:retag-return-place, r=RalfJung

retag return place

@eddyb suggested that return places should be treated like unique references for Stacked Borrows. That is implemented by this patch, but it is unfortunately quite the hack because otherwise we are retagging *references*, not places.

@eddyb does this roughly correspond to what you had in mind? (Except for whatever it is you think should happen with argument passing, which is a much bigger issue.) Also, do you think there is any way we can *test* this?

Needs https://github.com/rust-lang/rust/pull/71100 to land.

rust-version

@@ -1 +1 @@
-47f49695dfb4fe9e584239fdc59c771887148a57
+df768c5c8fcb361c4dc94b4c776d6a78c12862e1

src/machine.rs

@@ -481,30 +481,42 @@ fn retag(
         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(())
         }
     }
 
     #[inline(always)]
-    fn stack_push(ecx: &mut InterpCx<'mir, 'tcx, Self>) -> InterpResult<'tcx, FrameData<'tcx>> {
+    fn init_frame_extra(
+        ecx: &mut InterpCx<'mir, 'tcx, Self>,
+        frame: Frame<'mir, 'tcx, Tag>,
+    ) -> InterpResult<'tcx, Frame<'mir, 'tcx, Tag, FrameData<'tcx>>> {
         let stacked_borrows = ecx.memory.extra.stacked_borrows.as_ref();
         let call_id = stacked_borrows.map_or(NonZeroU64::new(1).unwrap(), |stacked_borrows| {
             stacked_borrows.borrow_mut().new_call()
         });
-        Ok(FrameData { call_id, catch_unwind: None })
+        let extra = FrameData { call_id, catch_unwind: None };
+        Ok(frame.with_extra(extra))
     }
 
     #[inline(always)]
-    fn stack_pop(
+    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>,
-        extra: FrameData<'tcx>,
+        frame: Frame<'mir, 'tcx, Tag, FrameData<'tcx>>,
         unwinding: bool,
     ) -> InterpResult<'tcx, StackPopJump> {
-        ecx.handle_stack_pop(extra, unwinding)
+        ecx.handle_stack_pop(frame.extra, unwinding)
     }
 
     #[inline(always)]

src/stacked_borrows.rs

@@ -11,7 +11,7 @@
 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 @@ fn retag_reference(
         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 @@ fn retag_reference(
         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 @@ fn retag_reference(
         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 @@ fn qualify(ty: ty::Ty<'_>, kind: RetagKind) -> Option<(RefKind, bool)> {
             // 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(())
+    }
 }