Rollup merge of #105317 - RalfJung:retag-rework, r=oli-obk

make retagging work even with 'unstable' places

This is based on top of #105301. Only the last two commits are new.

While investigating rust-lang/unsafe-code-guidelines#381 I realized that we would have caught this issue much earlier if the add_retag pass wouldn't bail out on assignments of the form `*ptr = ...`.

So this PR changes our retag strategy:
- When a new reference is created via `Rvalue::Ref` (or a raw ptr via `Rvalue::AddressOf`), we do the retagging as part of just executing that address-taking operation.
- For everything else, we still insert retags -- these retags basically serve to ensure that references stored in local variables (and their fields) are always freshly tagged, so skipping this for assignments like `*ptr = ...` is less egregious.
r? ```@oli-obk```

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -373,9 +373,21 @@ pub trait Machine<'mir, 'tcx>: Sized {
         Ok(())
     }
 
-    /// Executes a retagging operation.
+    /// Executes a retagging operation for a single pointer.
+    /// Returns the possibly adjusted pointer.
     #[inline]
-    fn retag(
+    fn retag_ptr_value(
+        _ecx: &mut InterpCx<'mir, 'tcx, Self>,
+        _kind: mir::RetagKind,
+        val: &ImmTy<'tcx, Self::Provenance>,
+    ) -> InterpResult<'tcx, ImmTy<'tcx, Self::Provenance>> {
+        Ok(val.clone())
+    }
+
+    /// Executes a retagging operation on a compound value.
+    /// Replaces all pointers stored in the given place.
+    #[inline]
+    fn retag_place_contents(
         _ecx: &mut InterpCx<'mir, 'tcx, Self>,
         _kind: mir::RetagKind,
         _place: &PlaceTy<'tcx, Self::Provenance>,

compiler/rustc_const_eval/src/interpret/step.rs

@@ -8,7 +8,7 @@ use rustc_middle::mir;
 use rustc_middle::mir::interpret::{InterpResult, Scalar};
 use rustc_middle::ty::layout::LayoutOf;
 
-use super::{InterpCx, Machine};
+use super::{ImmTy, InterpCx, Machine};
 
 /// Classify whether an operator is "left-homogeneous", i.e., the LHS has the
 /// same type as the result.
@@ -108,7 +108,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             // Stacked Borrows.
             Retag(kind, place) => {
                 let dest = self.eval_place(**place)?;
-                M::retag(self, *kind, &dest)?;
+                M::retag_place_contents(self, *kind, &dest)?;
             }
 
             Intrinsic(box ref intrinsic) => self.emulate_nondiverging_intrinsic(intrinsic)?,
@@ -247,10 +247,41 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 self.write_scalar(Scalar::from_machine_usize(len, self), &dest)?;
             }
 
-            AddressOf(_, place) | Ref(_, _, place) => {
+            Ref(_, borrow_kind, place) => {
                 let src = self.eval_place(place)?;
                 let place = self.force_allocation(&src)?;
-                self.write_immediate(place.to_ref(self), &dest)?;
+                let val = ImmTy::from_immediate(place.to_ref(self), dest.layout);
+                // A fresh reference was created, make sure it gets retagged.
+                let val = M::retag_ptr_value(
+                    self,
+                    if borrow_kind.allows_two_phase_borrow() {
+                        mir::RetagKind::TwoPhase
+                    } else {
+                        mir::RetagKind::Default
+                    },
+                    &val,
+                )?;
+                self.write_immediate(*val, &dest)?;
+            }
+
+            AddressOf(_, place) => {
+                // Figure out whether this is an addr_of of an already raw place.
+                let place_base_raw = if place.has_deref() {
+                    let ty = self.frame().body.local_decls[place.local].ty;
+                    ty.is_unsafe_ptr()
+                } else {
+                    // Not a deref, and thus not raw.
+                    false
+                };
+
+                let src = self.eval_place(place)?;
+                let place = self.force_allocation(&src)?;
+                let mut val = ImmTy::from_immediate(place.to_ref(self), dest.layout);
+                if !place_base_raw {
+                    // If this was not already raw, it needs retagging.
+                    val = M::retag_ptr_value(self, mir::RetagKind::Raw, &val)?;
+                }
+                self.write_immediate(*val, &dest)?;
             }
 
             NullaryOp(null_op, ty) => {

compiler/rustc_middle/src/mir/syntax.rs

@@ -400,7 +400,7 @@ impl std::fmt::Display for NonDivergingIntrinsic<'_> {
 #[derive(Copy, Clone, TyEncodable, TyDecodable, Debug, PartialEq, Eq, Hash, HashStable)]
 #[rustc_pass_by_value]
 pub enum RetagKind {
-    /// The initial retag when entering a function.
+    /// The initial retag of arguments when entering a function.
     FnEntry,
     /// Retag preparing for a two-phase borrow.
     TwoPhase,

compiler/rustc_mir_transform/src/add_retag.rs

@@ -10,16 +10,6 @@ use rustc_middle::ty::{self, Ty, TyCtxt};
 
 pub struct AddRetag;
 
-/// Determines whether this place is "stable": Whether, if we evaluate it again
-/// after the assignment, we can be sure to obtain the same place value.
-/// (Concurrent accesses by other threads are no problem as these are anyway non-atomic
-/// copies.  Data races are UB.)
-fn is_stable(place: PlaceRef<'_>) -> bool {
-    // Which place this evaluates to can change with any memory write,
-    // so cannot assume deref to be stable.
-    !place.has_deref()
-}
-
 /// Determine whether this type may contain a reference (or box), and thus needs retagging.
 /// We will only recurse `depth` times into Tuples/ADTs to bound the cost of this.
 fn may_contain_reference<'tcx>(ty: Ty<'tcx>, depth: u32, tcx: TyCtxt<'tcx>) -> bool {
@@ -69,22 +59,10 @@ impl<'tcx> MirPass<'tcx> for AddRetag {
         let basic_blocks = body.basic_blocks.as_mut();
         let local_decls = &body.local_decls;
         let needs_retag = |place: &Place<'tcx>| {
-            // FIXME: Instead of giving up for unstable places, we should introduce
-            // a temporary and retag on that.
-            is_stable(place.as_ref())
+            !place.has_deref() // we're not eally interested in stores to "outside" locations, they are hard to keep track of anyway
                 && may_contain_reference(place.ty(&*local_decls, tcx).ty, /*depth*/ 3, tcx)
                 && !local_decls[place.local].is_deref_temp()
         };
-        let place_base_raw = |place: &Place<'tcx>| {
-            // If this is a `Deref`, get the type of what we are deref'ing.
-            if place.has_deref() {
-                let ty = &local_decls[place.local].ty;
-                ty.is_unsafe_ptr()
-            } else {
-                // Not a deref, and thus not raw.
-                false
-            }
-        };
 
         // PART 1
         // Retag arguments at the beginning of the start block.
@@ -108,7 +86,7 @@ impl<'tcx> MirPass<'tcx> for AddRetag {
         }
 
         // PART 2
-        // Retag return values of functions.  Also escape-to-raw the argument of `drop`.
+        // Retag return values of functions.
         // We collect the return destinations because we cannot mutate while iterating.
         let returns = basic_blocks
             .iter_mut()
@@ -140,30 +118,25 @@ impl<'tcx> MirPass<'tcx> for AddRetag {
         }
 
         // PART 3
-        // Add retag after assignment.
+        // Add retag after assignments where data "enters" this function: the RHS is behind a deref and the LHS is not.
         for block_data in basic_blocks {
             // We want to insert statements as we iterate.  To this end, we
             // iterate backwards using indices.
             for i in (0..block_data.statements.len()).rev() {
                 let (retag_kind, place) = match block_data.statements[i].kind {
-                    // Retag-as-raw after escaping to a raw pointer, if the referent
-                    // is not already a raw pointer.
-                    StatementKind::Assign(box (lplace, Rvalue::AddressOf(_, ref rplace)))
-                        if !place_base_raw(rplace) =>
-                    {
-                        (RetagKind::Raw, lplace)
-                    }
                     // Retag after assignments of reference type.
                     StatementKind::Assign(box (ref place, ref rvalue)) if needs_retag(place) => {
-                        let kind = match rvalue {
-                            Rvalue::Ref(_, borrow_kind, _)
-                                if borrow_kind.allows_two_phase_borrow() =>
-                            {
-                                RetagKind::TwoPhase
-                            }
-                            _ => RetagKind::Default,
+                        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,
                         };
-                        (kind, *place)
+                        if add_retag {
+                            (RetagKind::Default, *place)
+                        } else {
+                            continue;
+                        }
                     }
                     // Do nothing for the rest
                     _ => continue,

compiler/rustc_mir_transform/src/generator.rs

@@ -985,16 +985,6 @@ fn create_generator_drop_shim<'tcx>(
         tcx.mk_ptr(ty::TypeAndMut { ty: gen_ty, mutbl: hir::Mutability::Mut }),
         source_info,
     );
-    if tcx.sess.opts.unstable_opts.mir_emit_retag {
-        // Alias tracking must know we changed the type
-        body.basic_blocks_mut()[START_BLOCK].statements.insert(
-            0,
-            Statement {
-                source_info,
-                kind: StatementKind::Retag(RetagKind::Raw, Box::new(Place::from(SELF_ARG))),
-            },
-        )
-    }
 
     // Make sure we remove dead blocks to remove
     // unrelated code from the resume part of the function

compiler/rustc_mir_transform/src/shim.rs

@@ -177,16 +177,6 @@ fn build_drop_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, ty: Option<Ty<'tcx>>)
     if ty.is_some() {
         // The first argument (index 0), but add 1 for the return value.
         let dropee_ptr = Place::from(Local::new(1 + 0));
-        if tcx.sess.opts.unstable_opts.mir_emit_retag {
-            // Function arguments should be retagged, and we make this one raw.
-            body.basic_blocks_mut()[START_BLOCK].statements.insert(
-                0,
-                Statement {
-                    source_info,
-                    kind: StatementKind::Retag(RetagKind::Raw, Box::new(dropee_ptr)),
-                },
-            );
-        }
         let patch = {
             let param_env = tcx.param_env_reveal_all_normalized(def_id);
             let mut elaborator =

src/test/mir-opt/inline/inline_retag.bar.Inline.after.mir

@@ -38,9 +38,7 @@ fn bar() -> bool {
                                          // + literal: Const { ty: &i32, val: Unevaluated(bar, [], Some(promoted[1])) }
         Retag(_10);                      // scope 1 at $DIR/inline_retag.rs:+2:7: +2:9
         _4 = &(*_10);                    // scope 1 at $DIR/inline_retag.rs:+2:7: +2:9
-        Retag(_4);                       // scope 1 at $DIR/inline_retag.rs:+2:7: +2:9
         _3 = &(*_4);                     // scope 1 at $DIR/inline_retag.rs:+2:7: +2:9
-        Retag(_3);                       // scope 1 at $DIR/inline_retag.rs:+2:7: +2:9
         StorageLive(_6);                 // scope 1 at $DIR/inline_retag.rs:+2:11: +2:14
         StorageLive(_7);                 // scope 1 at $DIR/inline_retag.rs:+2:11: +2:14
         _9 = const _;                    // scope 1 at $DIR/inline_retag.rs:+2:11: +2:14
@@ -49,9 +47,7 @@ fn bar() -> bool {
                                          // + literal: Const { ty: &i32, val: Unevaluated(bar, [], Some(promoted[0])) }
         Retag(_9);                       // scope 1 at $DIR/inline_retag.rs:+2:11: +2:14
         _7 = &(*_9);                     // scope 1 at $DIR/inline_retag.rs:+2:11: +2:14
-        Retag(_7);                       // scope 1 at $DIR/inline_retag.rs:+2:11: +2:14
         _6 = &(*_7);                     // scope 1 at $DIR/inline_retag.rs:+2:11: +2:14
-        Retag(_6);                       // scope 1 at $DIR/inline_retag.rs:+2:11: +2:14
         Retag(_3);                       // scope 2 at $DIR/inline_retag.rs:16:8: 16:9
         Retag(_6);                       // scope 2 at $DIR/inline_retag.rs:16:17: 16:18
         StorageLive(_11);                // scope 2 at $DIR/inline_retag.rs:17:5: 17:7

src/test/mir-opt/retag.array_casts.SimplifyCfg-elaborate-drops.after.mir

@@ -68,9 +68,7 @@ fn array_casts() -> () {
         StorageLive(_3);                 // scope 1 at $DIR/retag.rs:+2:13: +2:19
         StorageLive(_4);                 // scope 1 at $DIR/retag.rs:+2:13: +2:19
         _4 = &mut _1;                    // scope 1 at $DIR/retag.rs:+2:13: +2:19
-        Retag(_4);                       // scope 1 at $DIR/retag.rs:+2:13: +2:19
         _3 = &raw mut (*_4);             // scope 1 at $DIR/retag.rs:+2:13: +2:19
-        Retag([raw] _3);                 // scope 1 at $DIR/retag.rs:+2:13: +2:19
         _2 = move _3 as *mut usize (Pointer(ArrayToPointer)); // scope 1 at $DIR/retag.rs:+2:13: +2:33
         StorageDead(_3);                 // scope 1 at $DIR/retag.rs:+2:32: +2:33
         StorageDead(_4);                 // scope 1 at $DIR/retag.rs:+2:33: +2:34
@@ -96,9 +94,7 @@ fn array_casts() -> () {
         StorageLive(_10);                // scope 4 at $DIR/retag.rs:+6:13: +6:15
         StorageLive(_11);                // scope 4 at $DIR/retag.rs:+6:13: +6:15
         _11 = &_8;                       // scope 4 at $DIR/retag.rs:+6:13: +6:15
-        Retag(_11);                      // scope 4 at $DIR/retag.rs:+6:13: +6:15
         _10 = &raw const (*_11);         // scope 4 at $DIR/retag.rs:+6:13: +6:15
-        Retag([raw] _10);                // scope 4 at $DIR/retag.rs:+6:13: +6:15
         _9 = move _10 as *const usize (Pointer(ArrayToPointer)); // scope 4 at $DIR/retag.rs:+6:13: +6:31
         StorageDead(_10);                // scope 4 at $DIR/retag.rs:+6:30: +6:31
         StorageDead(_11);                // scope 4 at $DIR/retag.rs:+6:31: +6:32
@@ -119,15 +115,13 @@ fn array_casts() -> () {
         StorageDead(_17);                // scope 6 at $DIR/retag.rs:+7:33: +7:34
         _15 = (*_16);                    // scope 6 at $DIR/retag.rs:+7:25: +7:34
         _14 = &_15;                      // scope 5 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
-        Retag(_14);                      // scope 5 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         StorageLive(_18);                // scope 5 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         _35 = const _;                   // scope 5 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
                                          // mir::Constant
                                          // + span: $SRC_DIR/core/src/macros/mod.rs:LL:COL
                                          // + literal: Const { ty: &usize, val: Unevaluated(array_casts, [], Some(promoted[0])) }
         Retag(_35);                      // scope 5 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         _18 = &(*_35);                   // scope 5 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
-        Retag(_18);                      // scope 5 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         Deinit(_13);                     // scope 5 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         (_13.0: &usize) = move _14;      // scope 5 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         (_13.1: &usize) = move _18;      // scope 5 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
@@ -164,15 +158,11 @@ fn array_casts() -> () {
         StorageLive(_30);                // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         StorageLive(_31);                // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         _31 = &(*_20);                   // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
-        Retag(_31);                      // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         _30 = &(*_31);                   // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
-        Retag(_30);                      // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         StorageLive(_32);                // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         StorageLive(_33);                // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         _33 = &(*_21);                   // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
-        Retag(_33);                      // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         _32 = &(*_33);                   // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
-        Retag(_32);                      // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         StorageLive(_34);                // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         Deinit(_34);                     // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL
         discriminant(_34) = 0;           // scope 8 at $SRC_DIR/core/src/macros/mod.rs:LL:COL

src/test/mir-opt/retag.core.ptr-drop_in_place.Test.SimplifyCfg-make_shim.after.mir

@@ -6,7 +6,6 @@ fn std::ptr::drop_in_place(_1: *mut Test) -> () {
     let mut _3: ();                      // in scope 0 at $SRC_DIR/core/src/ptr/mod.rs:+0:1: +0:56
 
     bb0: {
-        Retag([raw] _1);                 // scope 0 at $SRC_DIR/core/src/ptr/mod.rs:+0:1: +0:56
         _2 = &mut (*_1);                 // scope 0 at $SRC_DIR/core/src/ptr/mod.rs:+0:1: +0:56
         _3 = <Test as Drop>::drop(move _2) -> bb1; // scope 0 at $SRC_DIR/core/src/ptr/mod.rs:+0:1: +0:56
                                          // mir::Constant

src/test/mir-opt/retag.main-{closure#0}.SimplifyCfg-elaborate-drops.after.mir

@@ -15,7 +15,6 @@ fn main::{closure#0}(_1: &[closure@main::{closure#0}], _2: &i32) -> &i32 {
         _3 = _2;                         // scope 0 at $DIR/retag.rs:+1:18: +1:19
         Retag(_3);                       // scope 0 at $DIR/retag.rs:+1:18: +1:19
         _0 = &(*_2);                     // scope 1 at $DIR/retag.rs:+2:9: +2:10
-        Retag(_0);                       // scope 1 at $DIR/retag.rs:+2:9: +2:10
         StorageDead(_3);                 // scope 0 at $DIR/retag.rs:+3:5: +3:6
         return;                          // scope 0 at $DIR/retag.rs:+3:6: +3:6
     }