Stacked Borrows: raw pointers inherit the tag from their parent pointer

compiler/rustc_middle/src/mir/mod.rs

@@ -1108,7 +1108,7 @@ pub enum LocalInfo<'tcx> {
     /// A temporary created during evaluating `if` predicate, possibly for pattern matching for `let`s,
     /// and subject to Edition 2024 temporary lifetime rules
     IfThenRescopeTemp { if_then: HirId },
-    /// A temporary created during the pass `Derefer` to avoid it's retagging
+    /// A temporary created during the pass `Derefer` to avoid its retagging
     DerefTemp,
     /// A temporary created for borrow checking.
     FakeBorrow,

compiler/rustc_mir_transform/src/add_retag.rs

@@ -4,7 +4,6 @@
 //! of MIR building, and only after this pass we think of the program has having the
 //! normal MIR semantics.
 
-use rustc_hir::LangItem;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, Ty, TyCtxt};
 
@@ -28,7 +27,6 @@ fn may_contain_reference<'tcx>(ty: Ty<'tcx>, depth: u32, tcx: TyCtxt<'tcx>) -> b
         // References and Boxes (`noalias` sources)
         ty::Ref(..) => true,
         ty::Adt(..) if ty.is_box() => true,
-        ty::Adt(adt, _) if tcx.is_lang_item(adt.did(), LangItem::PtrUnique) => true,
         // Compound types: recurse
         ty::Array(ty, _) | ty::Slice(ty) => {
             // This does not branch so we keep the depth the same.
@@ -153,6 +151,13 @@ impl<'tcx> crate::MirPass<'tcx> for AddRetag {
                                 }
                             }
                             Rvalue::Ref(..) => None,
+                            // Also skip coercions whose source is already a local. This is crucial
+                            // to avoid unnecessary retags in `my_array.len()` calls.
+                            Rvalue::Cast(
+                                CastKind::PointerCoercion(_, CoercionSource::Implicit),
+                                Operand::Copy(p) | Operand::Move(p),
+                                _,
+                            ) if p.as_local().is_some() => None,
                             _ => {
                                 if needs_retag(place) {
                                     Some(RetagKind::Default)

compiler/rustc_mir_transform/src/lib.rs

@@ -628,6 +628,7 @@ fn run_runtime_lowering_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
 fn run_runtime_cleanup_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     let passes: &[&dyn MirPass<'tcx>] = &[
         &lower_intrinsics::LowerIntrinsics,
+        &lower_slice_len::LowerSliceLenCalls,
         &remove_place_mention::RemovePlaceMention,
         &simplify::SimplifyCfg::PreOptimizations,
     ];
@@ -671,9 +672,6 @@ pub(crate) fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'
             &check_null::CheckNull,
             // Before inlining: trim down MIR with passes to reduce inlining work.
 
-            // Has to be done before inlining, otherwise actual call will be almost always inlined.
-            // Also simple, so can just do first.
-            &lower_slice_len::LowerSliceLenCalls,
             // Perform instsimplify before inline to eliminate some trivial calls (like clone
             // shims).
             &instsimplify::InstSimplify::BeforeInline,

compiler/rustc_mir_transform/src/lower_slice_len.rs

@@ -2,14 +2,18 @@
 //! It should run before inlining!
 
 use rustc_hir::def_id::DefId;
+use rustc_index::IndexVec;
 use rustc_middle::mir::*;
-use rustc_middle::ty::TyCtxt;
+use rustc_middle::ty::adjustment::PointerCoercion;
+use rustc_middle::ty::{Ty, TyCtxt};
 
 pub(super) struct LowerSliceLenCalls;
 
 impl<'tcx> crate::MirPass<'tcx> for LowerSliceLenCalls {
     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
-        sess.mir_opt_level() > 0
+        // This pass also removes some otherwise-problematic retags, so we
+        // always enable it when retags matter.
+        sess.mir_opt_level() > 0 || sess.opts.unstable_opts.mir_emit_retag
     }
 
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
@@ -23,7 +27,7 @@ impl<'tcx> crate::MirPass<'tcx> for LowerSliceLenCalls {
         let basic_blocks = body.basic_blocks.as_mut_preserves_cfg();
         for block in basic_blocks {
             // lower `<[_]>::len` calls
-            lower_slice_len_call(block, slice_len_fn_item_def_id);
+            lower_slice_len_call(tcx, &mut body.local_decls, block, slice_len_fn_item_def_id);
         }
     }
 
@@ -32,7 +36,12 @@ impl<'tcx> crate::MirPass<'tcx> for LowerSliceLenCalls {
     }
 }
 
-fn lower_slice_len_call<'tcx>(block: &mut BasicBlockData<'tcx>, slice_len_fn_item_def_id: DefId) {
+fn lower_slice_len_call<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
+    block: &mut BasicBlockData<'tcx>,
+    slice_len_fn_item_def_id: DefId,
+) {
     let terminator = block.terminator();
     if let TerminatorKind::Call {
         func,
@@ -50,11 +59,12 @@ fn lower_slice_len_call<'tcx>(block: &mut BasicBlockData<'tcx>, slice_len_fn_ite
         // perform modifications from something like:
         //     _5 = core::slice::<impl [u8]>::len(move _6) -> bb1
         // into:
-        //     _5 = PtrMetadata(move _6)
+        //     _5 = PtrMetadata(move _6);
         //     goto bb1
 
         // make new RValue for Len
-        let r_value = Rvalue::UnaryOp(UnOp::PtrMetadata, arg.node.clone());
+        let arg = arg.node.clone();
+        let r_value = Rvalue::UnaryOp(UnOp::PtrMetadata, arg.clone());
         let len_statement_kind = StatementKind::Assign(Box::new((*destination, r_value)));
         let add_statement =
             Statement { kind: len_statement_kind, source_info: terminator.source_info };
@@ -64,5 +74,75 @@ fn lower_slice_len_call<'tcx>(block: &mut BasicBlockData<'tcx>, slice_len_fn_ite
 
         block.statements.push(add_statement);
         block.terminator_mut().kind = new_terminator_kind;
+
+        if tcx.sess.opts.unstable_opts.mir_emit_retag {
+            // If we care about retags, we want there to be *no retag* for this `len` call.
+            // That means we have to clean up the MIR a little: the MIR will now often be:
+            //     _6 = &*_4;
+            //     _5 = PtrMetadata(move _6);
+            // Change that to:
+            //     _6 = &raw const *_4;
+            //     _5 = PtrMetadata(move _6);
+            let len = block.statements.len();
+            if len >= 2
+                && let &mut StatementKind::Assign(box (rebor_dest, ref mut rebor_r)) =
+                    &mut block.statements[len - 2].kind
+                && let Some(retag_local) = rebor_dest.as_local()
+                && !local_decls[retag_local].is_user_variable()
+                // The LHS of this previous assignment is the `arg` from above.
+                && matches!(arg, Operand::Copy(p) | Operand::Move(p) if p.as_local() == Some(retag_local))
+                // The RHS is a reference-taking operation.
+                && let Rvalue::Ref(_, BorrowKind::Shared, rebor_orig) = *rebor_r
+            {
+                let orig_inner_ty = rebor_orig.ty(local_decls, tcx).ty;
+                // Change the previous statement to use `&raw` instead of `&`.
+                *rebor_r = Rvalue::RawPtr(RawPtrKind::FakeForPtrMetadata, rebor_orig);
+                // Change the type of the local to match.
+                local_decls[retag_local].ty = Ty::new_ptr(tcx, orig_inner_ty, Mutability::Not);
+            }
+            // There's a second pattern we need to recognize for `array.len()` calls. There,
+            // the MIR is:
+            //     _4 = &*_3;
+            //     _6 = move _4 as &[T] (PointerCoercion(Unsize, Implicit));
+            //     StorageDead(_4);
+            //     _5 = PtrMetadata(move _6);
+            // Change that to:
+            //     _4 = &raw const *_3;
+            //     _6 = move _4 as *const [T] (PointerCoercion(Unsize, Implicit));
+            //     StorageDead(_4);
+            //     _5 = PtrMetadata(move _6);
+            if len >= 4
+                && let [reborrow_stmt, unsize_stmt, storage_dead_stmt] =
+                    block.statements.get_disjoint_mut([len - 4, len - 3, len - 2]).unwrap()
+                // Check that the 3 statements have the right shape.
+                && let &mut StatementKind::Assign(box (rebor_dest, ref mut rebor_r)) = &mut reborrow_stmt.kind
+                && let Some(retag_local) = rebor_dest.as_local()
+                && !local_decls[retag_local].is_user_variable()
+                && let &mut StatementKind::Assign(box (unsized_dest, ref mut unsize_r)) = &mut unsize_stmt.kind
+                && let Some(unsized_local) = unsized_dest.as_local()
+                && !local_decls[unsized_local].is_user_variable()
+                && storage_dead_stmt.kind == StatementKind::StorageDead(retag_local)
+                // And check that they have the right operands.
+                && matches!(arg, Operand::Copy(p) | Operand::Move(p) if p.as_local() == Some(unsized_local))
+                && let &mut Rvalue::Cast(
+                    CastKind::PointerCoercion(PointerCoercion::Unsize, CoercionSource::Implicit),
+                    ref cast_op,
+                    ref mut cast_ty,
+                ) = unsize_r
+                && let Operand::Copy(unsize_src) | Operand::Move(unsize_src) = cast_op
+                && unsize_src.as_local() == Some(retag_local)
+                && let Rvalue::Ref(_, BorrowKind::Shared, rebor_orig) = *rebor_r
+            {
+                let orig_inner_ty = rebor_orig.ty(local_decls, tcx).ty;
+                let unsize_inner_ty = cast_ty.builtin_deref(true).unwrap();
+                // Change the reborrow to use `&raw` instead of `&`.
+                *rebor_r = Rvalue::RawPtr(RawPtrKind::FakeForPtrMetadata, rebor_orig);
+                // Change the unsize coercion in the same vein.
+                *cast_ty = Ty::new_ptr(tcx, unsize_inner_ty, Mutability::Not);
+                // Change the type of the locals to match.
+                local_decls[retag_local].ty = Ty::new_ptr(tcx, orig_inner_ty, Mutability::Not);
+                local_decls[unsized_local].ty = *cast_ty;
+            }
+        }
     }
 }

library/alloctests/tests/sync.rs

@@ -332,15 +332,15 @@ fn weak_self_cyclic() {
 #[test]
 fn drop_arc() {
     let mut canary = AtomicUsize::new(0);
-    let x = Arc::new(Canary(&mut canary as *mut AtomicUsize));
+    let x = Arc::new(Canary(&raw mut canary));
     drop(x);
     assert!(canary.load(Acquire) == 1);
 }
 
 #[test]
 fn drop_arc_weak() {
     let mut canary = AtomicUsize::new(0);
-    let arc = Arc::new(Canary(&mut canary as *mut AtomicUsize));
+    let arc = Arc::new(Canary(&raw mut canary));
     let arc_weak = Arc::downgrade(&arc);
     assert!(canary.load(Acquire) == 0);
     drop(arc);