retag return places

Author: RalfJung

src/machine.rs

@@ -468,11 +468,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(())
         }
     }
 
@@ -489,6 +488,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::*;
@@ -645,4 +645,34 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
 
         Ok(())
     }
+
+    /// After a stack frame got pushed, retag the return place so that we are sure
+    /// it does not alias with anything.
+    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 reference to use the existing code.
+        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, then make it back into an `ImmTy`.
+        let val = this.retag_reference(val, RefKind::Unique { two_phase: false }, /*protector*/ true)?;
+        let val = ImmTy::from_immediate(val, ptr_layout);
+        // 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(())
+    }
 }