Rollup merge of #122647 - RalfJung:box-to-raw-retag, r=oli-obk

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

In https://github.com/rust-lang/rust/pull/122233 I added `retag_box_to_raw` not realizing that we can already do `addr_of_mut!(*bx)` to turn a box into a raw pointer without an intermediate reference. We just need to ensure this information is preserved past the ElaborateBoxDerefs pass.

r? ``@oli-obk``

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"));
                 }
             }

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -658,10 +658,6 @@ pub fn check_intrinsic_type(
             sym::simd_shuffle => (3, 0, vec![param(0), param(0), param(1)], param(2)),
             sym::simd_shuffle_generic => (2, 1, vec![param(0), param(0)], param(1)),
 
-            sym::retag_box_to_raw => {
-                (2, 0, vec![Ty::new_mut_ptr(tcx, param(0))], Ty::new_mut_ptr(tcx, param(0)))
-            }
-
             other => {
                 tcx.dcx().emit_err(UnrecognizedIntrinsicFunction { span, name: other });
                 return;

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;
                         }

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)));

compiler/rustc_span/src/symbol.rs

@@ -1463,7 +1463,6 @@ symbols! {
         residual,
         result,
         resume,
-        retag_box_to_raw,
         return_position_impl_trait_in_trait,
         return_type_notation,
         rhs,

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(

library/core/src/intrinsics.rs

@@ -2709,19 +2709,6 @@ pub unsafe fn vtable_size(_ptr: *const ()) -> usize {
     unreachable!()
 }
 
-/// Retag a box pointer as part of casting it to a raw pointer. This is the `Box` equivalent of
-/// `(x: &mut T) as *mut T`. The input pointer must be the pointer of a `Box` (passed as raw pointer
-/// to avoid all questions around move semantics and custom allocators), and `A` must be the `Box`'s
-/// allocator.
-#[unstable(feature = "core_intrinsics", issue = "none")]
-#[rustc_nounwind]
-#[cfg_attr(not(bootstrap), rustc_intrinsic)]
-#[cfg_attr(bootstrap, inline)]
-pub unsafe fn retag_box_to_raw<T: ?Sized, A>(ptr: *mut T) -> *mut T {
-    // Miri needs to adjust the provenance, but for regular codegen this is not needed
-    ptr
-}
-
 // Some functions are defined here because they accidentally got made
 // available in this module on stable. See <https://github.com/rust-lang/rust/issues/15702>.
 // (`transmute` also falls into this category, but it cannot be wrapped due to the

src/tools/miri/src/borrow_tracker/mod.rs

@@ -5,7 +5,7 @@ use std::num::NonZero;
 use smallvec::SmallVec;
 
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
-use rustc_middle::{mir::RetagKind, ty::Ty};
+use rustc_middle::mir::RetagKind;
 use rustc_target::abi::Size;
 
 use crate::*;
@@ -291,19 +291,6 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
     }
 
-    fn retag_box_to_raw(
-        &mut self,
-        val: &ImmTy<'tcx, Provenance>,
-        alloc_ty: Ty<'tcx>,
-    ) -> InterpResult<'tcx, ImmTy<'tcx, Provenance>> {
-        let this = self.eval_context_mut();
-        let method = this.machine.borrow_tracker.as_ref().unwrap().borrow().borrow_tracker_method;
-        match method {
-            BorrowTrackerMethod::StackedBorrows => this.sb_retag_box_to_raw(val, alloc_ty),
-            BorrowTrackerMethod::TreeBorrows => this.tb_retag_box_to_raw(val, alloc_ty),
-        }
-    }
-
     fn retag_place_contents(
         &mut self,
         kind: RetagKind,

src/tools/miri/src/borrow_tracker/stacked_borrows/mod.rs

@@ -865,24 +865,6 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         this.sb_retag_reference(val, new_perm, RetagInfo { cause, in_field: false })
     }
 
-    fn sb_retag_box_to_raw(
-        &mut self,
-        val: &ImmTy<'tcx, Provenance>,
-        alloc_ty: Ty<'tcx>,
-    ) -> InterpResult<'tcx, ImmTy<'tcx, Provenance>> {
-        let this = self.eval_context_mut();
-        let is_global_alloc = alloc_ty.ty_adt_def().is_some_and(|adt| {
-            let global_alloc = this.tcx.require_lang_item(rustc_hir::LangItem::GlobalAlloc, None);
-            adt.did() == global_alloc
-        });
-        if is_global_alloc {
-            // Retag this as-if it was a mutable reference.
-            this.sb_retag_ptr_value(RetagKind::Raw, val)
-        } else {
-            Ok(val.clone())
-        }
-    }
-
     fn sb_retag_place_contents(
         &mut self,
         kind: RetagKind,
@@ -891,9 +873,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 +941,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.

src/tools/miri/src/borrow_tracker/tree_borrows/mod.rs

@@ -392,15 +392,6 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
     }
 
-    fn tb_retag_box_to_raw(
-        &mut self,
-        val: &ImmTy<'tcx, Provenance>,
-        _alloc_ty: Ty<'tcx>,
-    ) -> InterpResult<'tcx, ImmTy<'tcx, Provenance>> {
-        // Casts to raw pointers are NOPs in Tree Borrows.
-        Ok(val.clone())
-    }
-
     /// Retag all pointers that are stored in this place.
     fn tb_retag_place_contents(
         &mut self,

src/tools/miri/src/shims/intrinsics/mod.rs

@@ -140,18 +140,6 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
 
                 this.write_pointer(Pointer::new(ptr.provenance, masked_addr), dest)?;
             }
-            "retag_box_to_raw" => {
-                let [ptr] = check_arg_count(args)?;
-                let alloc_ty = generic_args[1].expect_ty();
-
-                let val = this.read_immediate(ptr)?;
-                let new_val = if this.machine.borrow_tracker.is_some() {
-                    this.retag_box_to_raw(&val, alloc_ty)?
-                } else {
-                    val
-                };
-                this.write_immediate(*new_val, dest)?;
-            }
 
             // We want to return either `true` or `false` at random, or else something like
             // ```

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));

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));

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;
+    }
+}

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;
+    }
+}

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)
+}