add_retag: ensure box-to-raw-ptr casts are preserved for Miri

10 files changed, 91 insertions, 33 deletions

diff --git a/compiler/rustc_const_eval/src/transform/validate.rs b/compiler/rustc_const_eval/src/transform/validate.rs
index f26c8f8592d..68270dab284 100644
--- a/compiler/rustc_const_eval/src/transform/validate.rs
+++ b/compiler/rustc_const_eval/src/transform/validate.rs
@@ -341,7 +341,7 @@ impl<'a, 'tcx> Visitor<'tcx> for CfgChecker<'a, 'tcx> {
                 // FIXME(JakobDegen) The validator should check that `self.mir_phase <
                 // DropsLowered`. However, this causes ICEs with generation of drop shims, which
                 // seem to fail to set their `MirPhase` correctly.
-                if matches!(kind, RetagKind::Raw | RetagKind::TwoPhase) {
+                if matches!(kind, RetagKind::TwoPhase) {
                     self.fail(location, format!("explicit `{kind:?}` is forbidden"));
                 }
             }
@@ -1272,7 +1272,7 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                 // FIXME(JakobDegen) The validator should check that `self.mir_phase <
                 // DropsLowered`. However, this causes ICEs with generation of drop shims, which
                 // seem to fail to set their `MirPhase` correctly.
-                if matches!(kind, RetagKind::Raw | RetagKind::TwoPhase) {
+                if matches!(kind, RetagKind::TwoPhase) {
                     self.fail(location, format!("explicit `{kind:?}` is forbidden"));
                 }
             }
diff --git a/compiler/rustc_mir_transform/src/add_retag.rs b/compiler/rustc_mir_transform/src/add_retag.rs
index 6f668aa4ce8..f880476cec2 100644
--- a/compiler/rustc_mir_transform/src/add_retag.rs
+++ b/compiler/rustc_mir_transform/src/add_retag.rs
@@ -24,7 +24,7 @@ fn may_contain_reference<'tcx>(ty: Ty<'tcx>, depth: u32, tcx: TyCtxt<'tcx>) -> b
         | ty::Str
         | ty::FnDef(..)
         | ty::Never => false,
-        // References
+        // References and Boxes (`noalias` sources)
         ty::Ref(..) => true,
         ty::Adt(..) if ty.is_box() => true,
         ty::Adt(adt, _) if Some(adt.did()) == tcx.lang_items().ptr_unique() => true,
@@ -125,15 +125,39 @@ impl<'tcx> MirPass<'tcx> for AddRetag {
             for i in (0..block_data.statements.len()).rev() {
                 let (retag_kind, place) = match block_data.statements[i].kind {
                     // Retag after assignments of reference type.
-                    StatementKind::Assign(box (ref place, ref rvalue)) if needs_retag(place) => {
+                    StatementKind::Assign(box (ref place, ref rvalue)) => {
                         let add_retag = match rvalue {
                             // Ptr-creating operations already do their own internal retagging, no
                             // need to also add a retag statement.
-                            Rvalue::Ref(..) | Rvalue::AddressOf(..) => false,
-                            _ => true,
+                            // *Except* if we are deref'ing a Box, because those get desugared to directly working
+                            // with the inner raw pointer! That's relevant for `AddressOf` as Miri otherwise makes it
+                            // a NOP when the original pointer is already raw.
+                            Rvalue::AddressOf(_mutbl, place) => {
+                                // Using `is_box_global` here is a bit sketchy: if this code is
+                                // generic over the allocator, we'll not add a retag! This is a hack
+                                // to make Stacked Borrows compatible with custom allocator code.
+                                // Long-term, we'll want to move to an aliasing model where "cast to
+                                // raw pointer" is a complete NOP, and then this will no longer be
+                                // an issue.
+                                if place.is_indirect_first_projection()
+                                    && body.local_decls[place.local].ty.is_box_global(tcx)
+                                {
+                                    Some(RetagKind::Raw)
+                                } else {
+                                    None
+                                }
+                            }
+                            Rvalue::Ref(..) => None,
+                            _ => {
+                                if needs_retag(place) {
+                                    Some(RetagKind::Default)
+                                } else {
+                                    None
+                                }
+                            }
                         };
-                        if add_retag {
-                            (RetagKind::Default, *place)
+                        if let Some(kind) = add_retag {
+                            (kind, *place)
                         } else {
                             continue;
                         }
diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs
index 3e5ec323d27..afe228be127 100644
--- a/compiler/rustc_mir_transform/src/lib.rs
+++ b/compiler/rustc_mir_transform/src/lib.rs
@@ -529,11 +529,11 @@ fn run_runtime_lowering_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         // AddMovesForPackedDrops needs to run after drop
         // elaboration.
         &add_moves_for_packed_drops::AddMovesForPackedDrops,
-        // `AddRetag` needs to run after `ElaborateDrops`. Otherwise it should run fairly late,
+        // `AddRetag` needs to run after `ElaborateDrops` but before `ElaborateBoxDerefs`. Otherwise it should run fairly late,
         // but before optimizations begin.
+        &add_retag::AddRetag,
         &elaborate_box_derefs::ElaborateBoxDerefs,
         &coroutine::StateTransform,
-        &add_retag::AddRetag,
         &Lint(known_panics_lint::KnownPanicsLint),
     ];
     pm::run_passes_no_validate(tcx, body, passes, Some(MirPhase::Runtime(RuntimePhase::Initial)));
diff --git a/library/alloc/src/boxed.rs b/library/alloc/src/boxed.rs
index 304f607000b..f47fe560b07 100644
--- a/library/alloc/src/boxed.rs
+++ b/library/alloc/src/boxed.rs
@@ -155,7 +155,6 @@ use core::error::Error;
 use core::fmt;
 use core::future::Future;
 use core::hash::{Hash, Hasher};
-use core::intrinsics::retag_box_to_raw;
 use core::iter::FusedIterator;
 use core::marker::Tuple;
 use core::marker::Unsize;
@@ -165,7 +164,7 @@ use core::ops::{
     CoerceUnsized, Coroutine, CoroutineState, Deref, DerefMut, DispatchFromDyn, Receiver,
 };
 use core::pin::Pin;
-use core::ptr::{self, NonNull, Unique};
+use core::ptr::{self, addr_of_mut, NonNull, Unique};
 use core::task::{Context, Poll};
 
 #[cfg(not(no_global_oom_handling))]
@@ -1111,16 +1110,12 @@ impl<T: ?Sized, A: Allocator> Box<T, A> {
     #[unstable(feature = "allocator_api", issue = "32838")]
     #[inline]
     pub fn into_raw_with_allocator(b: Self) -> (*mut T, A) {
-        // This is the transition point from `Box` to raw pointers. For Stacked Borrows, these casts
-        // are relevant -- if this is a global allocator Box and we just get the pointer from `b.0`,
-        // it will have `Unique` permission, which is not what we want from a raw pointer. We could
-        // fix that by going through `&mut`, but then if this is *not* a global allocator Box, we'd
-        // be adding uniqueness assertions that we do not want. So for Miri's sake we pass this
-        // pointer through an intrinsic for box-to-raw casts, which can do the right thing wrt the
-        // aliasing model.
-        let b = mem::ManuallyDrop::new(b);
+        let mut b = mem::ManuallyDrop::new(b);
+        // We carefully get the raw pointer out in a way that Miri's aliasing model understands what
+        // is happening: using the primitive "deref" of `Box`.
+        let ptr = addr_of_mut!(**b);
         let alloc = unsafe { ptr::read(&b.1) };
-        (unsafe { retag_box_to_raw::<T, A>(b.0.as_ptr()) }, alloc)
+        (ptr, alloc)
     }
 
     #[unstable(
diff --git a/src/tools/miri/src/borrow_tracker/stacked_borrows/mod.rs b/src/tools/miri/src/borrow_tracker/stacked_borrows/mod.rs
index 9130601bbdd..613a245c06f 100644
--- a/src/tools/miri/src/borrow_tracker/stacked_borrows/mod.rs
+++ b/src/tools/miri/src/borrow_tracker/stacked_borrows/mod.rs
@@ -891,9 +891,9 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         let this = self.eval_context_mut();
         let retag_fields = this.machine.borrow_tracker.as_mut().unwrap().get_mut().retag_fields;
         let retag_cause = match kind {
-            RetagKind::Raw | RetagKind::TwoPhase { .. } => unreachable!(), // these can only happen in `retag_ptr_value`
+            RetagKind::TwoPhase { .. } => unreachable!(), // can only happen in `retag_ptr_value`
             RetagKind::FnEntry => RetagCause::FnEntry,
-            RetagKind::Default => RetagCause::Normal,
+            RetagKind::Default | RetagKind::Raw => RetagCause::Normal,
         };
         let mut visitor =
             RetagVisitor { ecx: this, kind, retag_cause, retag_fields, in_field: false };
@@ -959,14 +959,14 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
 
                 // Check the type of this value to see what to do with it (retag, or recurse).
                 match place.layout.ty.kind() {
-                    ty::Ref(..) => {
-                        let new_perm =
-                            NewPermission::from_ref_ty(place.layout.ty, self.kind, self.ecx);
-                        self.retag_ptr_inplace(place, new_perm)?;
-                    }
-                    ty::RawPtr(..) => {
-                        // We do *not* want to recurse into raw pointers -- wide raw pointers have
-                        // fields, and for dyn Trait pointees those can have reference type!
+                    ty::Ref(..) | ty::RawPtr(..) => {
+                        if matches!(place.layout.ty.kind(), ty::Ref(..))
+                            || self.kind == RetagKind::Raw
+                        {
+                            let new_perm =
+                                NewPermission::from_ref_ty(place.layout.ty, self.kind, self.ecx);
+                            self.retag_ptr_inplace(place, new_perm)?;
+                        }
                     }
                     ty::Adt(adt, _) if adt.is_box() => {
                         // Recurse for boxes, they require some tricky handling and will end up in `visit_box` above.
diff --git a/src/tools/miri/tests/fail/both_borrows/newtype_pair_retagging.stack.stderr b/src/tools/miri/tests/fail/both_borrows/newtype_pair_retagging.stack.stderr
index c26c7f397b0..867907e98e6 100644
--- a/src/tools/miri/tests/fail/both_borrows/newtype_pair_retagging.stack.stderr
+++ b/src/tools/miri/tests/fail/both_borrows/newtype_pair_retagging.stack.stderr
@@ -6,7 +6,7 @@ LL |         Box(unsafe { Unique::new_unchecked(raw) }, alloc)
    |
    = help: this indicates a potential bug in the program: it performed an invalid operation, but the Stacked Borrows rules it violated are still experimental
    = help: see https://github.com/rust-lang/unsafe-code-guidelines/blob/master/wip/stacked-borrows.md for further information
-help: <TAG> was created by a SharedReadWrite retag at offsets [0x0..0x4]
+help: <TAG> was created by a Unique retag at offsets [0x0..0x4]
   --> $DIR/newtype_pair_retagging.rs:LL:CC
    |
 LL |     let ptr = Box::into_raw(Box::new(0i32));
diff --git a/src/tools/miri/tests/fail/both_borrows/newtype_retagging.stack.stderr b/src/tools/miri/tests/fail/both_borrows/newtype_retagging.stack.stderr
index ae54da70fe2..56715938e97 100644
--- a/src/tools/miri/tests/fail/both_borrows/newtype_retagging.stack.stderr
+++ b/src/tools/miri/tests/fail/both_borrows/newtype_retagging.stack.stderr
@@ -6,7 +6,7 @@ LL |         Box(unsafe { Unique::new_unchecked(raw) }, alloc)
    |
    = help: this indicates a potential bug in the program: it performed an invalid operation, but the Stacked Borrows rules it violated are still experimental
    = help: see https://github.com/rust-lang/unsafe-code-guidelines/blob/master/wip/stacked-borrows.md for further information
-help: <TAG> was created by a SharedReadWrite retag at offsets [0x0..0x4]
+help: <TAG> was created by a Unique retag at offsets [0x0..0x4]
   --> $DIR/newtype_retagging.rs:LL:CC
    |
 LL |     let ptr = Box::into_raw(Box::new(0i32));
diff --git a/tests/mir-opt/retag.box_to_raw_mut.SimplifyCfg-pre-optimizations.after.panic-abort.mir b/tests/mir-opt/retag.box_to_raw_mut.SimplifyCfg-pre-optimizations.after.panic-abort.mir
new file mode 100644
index 00000000000..0e568f6a568
--- /dev/null
+++ b/tests/mir-opt/retag.box_to_raw_mut.SimplifyCfg-pre-optimizations.after.panic-abort.mir
@@ -0,0 +1,17 @@
+// MIR for `box_to_raw_mut` after SimplifyCfg-pre-optimizations
+
+fn box_to_raw_mut(_1: &mut Box<i32>) -> *mut i32 {
+    debug x => _1;
+    let mut _0: *mut i32;
+    let mut _2: std::boxed::Box<i32>;
+    let mut _3: *const i32;
+
+    bb0: {
+        Retag([fn entry] _1);
+        _2 = deref_copy (*_1);
+        _3 = (((_2.0: std::ptr::Unique<i32>).0: std::ptr::NonNull<i32>).0: *const i32);
+        _0 = &raw mut (*_3);
+        Retag([raw] _0);
+        return;
+    }
+}
diff --git a/tests/mir-opt/retag.box_to_raw_mut.SimplifyCfg-pre-optimizations.after.panic-unwind.mir b/tests/mir-opt/retag.box_to_raw_mut.SimplifyCfg-pre-optimizations.after.panic-unwind.mir
new file mode 100644
index 00000000000..0e568f6a568
--- /dev/null
+++ b/tests/mir-opt/retag.box_to_raw_mut.SimplifyCfg-pre-optimizations.after.panic-unwind.mir
@@ -0,0 +1,17 @@
+// MIR for `box_to_raw_mut` after SimplifyCfg-pre-optimizations
+
+fn box_to_raw_mut(_1: &mut Box<i32>) -> *mut i32 {
+    debug x => _1;
+    let mut _0: *mut i32;
+    let mut _2: std::boxed::Box<i32>;
+    let mut _3: *const i32;
+
+    bb0: {
+        Retag([fn entry] _1);
+        _2 = deref_copy (*_1);
+        _3 = (((_2.0: std::ptr::Unique<i32>).0: std::ptr::NonNull<i32>).0: *const i32);
+        _0 = &raw mut (*_3);
+        Retag([raw] _0);
+        return;
+    }
+}
diff --git a/tests/mir-opt/retag.rs b/tests/mir-opt/retag.rs
index 9cee96e4294..17b3c10abeb 100644
--- a/tests/mir-opt/retag.rs
+++ b/tests/mir-opt/retag.rs
@@ -65,3 +65,8 @@ fn array_casts() {
     let p = &x as *const usize;
     assert_eq!(unsafe { *p.add(1) }, 1);
 }
+
+// EMIT_MIR retag.box_to_raw_mut.SimplifyCfg-pre-optimizations.after.mir
+fn box_to_raw_mut(x: &mut Box<i32>) -> *mut i32 {
+    std::ptr::addr_of_mut!(**x)
+}