ref_prop: do not collapse a mutable reborrow onto a multi-borrowed place

compiler/rustc_mir_transform/src/ref_prop.rs

@@ -173,6 +173,11 @@ fn compute_replacement<'tcx>(
 
     let fully_replaceable_locals = fully_replaceable_locals(&ssa);
 
+    // Number of borrows that *directly* borrow each local (a borrow whose place is rooted at the
+    // local with no `Deref` in its projection). Used to keep mutable-reference propagation sound;
+    // see `count_direct_borrows` and rust-lang/rust#132898.
+    let direct_borrow_count = count_direct_borrows(body);
+
     // Returns true iff we can use `place` as a pointee.
     //
     // Note that we only need to verify that there is a single `StorageLive` statement, and we do
@@ -276,7 +281,15 @@ fn compute_replacement<'tcx>(
                     place = target.project_deeper(rest, tcx);
                 }
                 assert_ne!(place.local, local);
-                if is_constant_place(place) {
+                // rust-lang/rust#132898: propagating a mutable reference whose referent's root
+                // local has another independent direct borrow can shorten a reborrow's provenance
+                // and introduce UB (the "uniqueness" guard is necessary but not sufficient). Only
+                // block the mutable case where the target is a direct place (no `Deref`); reborrow
+                // targets (`*x`) are unaffected.
+                let blocked_by_aliasing = needs_unique
+                    && !place.projection.iter().any(|e| matches!(e, PlaceElem::Deref))
+                    && direct_borrow_count[place.local] > 1;
+                if is_constant_place(place) && !blocked_by_aliasing {
                     targets[local] = Value::Pointer(place, needs_unique);
                 }
             }
@@ -377,6 +390,26 @@ fn fully_replaceable_locals(ssa: &SsaLocals) -> DenseBitSet<Local> {
     replaceable
 }
 
+/// Count, for each local, how many borrows *directly* borrow it: a borrow whose place is rooted at
+/// the local with no `Deref` in its projection (`_1`, `_1.0` — but not `*_1`). When a local has more
+/// than one such borrow, a mutable reference to one of its (sub-)places is not the unique path to
+/// that memory, so propagating it through a reborrow can change pointer provenance and introduce UB.
+/// See rust-lang/rust#132898.
+fn count_direct_borrows<'tcx>(body: &Body<'tcx>) -> IndexVec<Local, u32> {
+    let mut counts = IndexVec::from_elem(0u32, &body.local_decls);
+    for bb in body.basic_blocks.iter() {
+        for stmt in &bb.statements {
+            if let Some((_, rvalue)) = stmt.kind.as_assign()
+                && let Rvalue::Ref(_, _, place) | Rvalue::RawPtr(_, place) = rvalue
+                && !place.projection.iter().any(|e| matches!(e, PlaceElem::Deref))
+            {
+                counts[place.local] += 1;
+            }
+        }
+    }
+    counts
+}
+
 /// Utility to help performing substitution of `*pattern` by `target`.
 struct Replacer<'tcx> {
     tcx: TyCtxt<'tcx>,

tests/mir-opt/reference_prop_mutable_alias.rs

@@ -0,0 +1,37 @@
+//@ test-mir-pass: ReferencePropagation
+//! Regression test for #132898.
+//!
+//! `ReferencePropagation` must not collapse a reborrow of a mutable reference
+//! (`_2 = &raw mut (*_3)` where `_3 = &mut _1.0`) onto the direct place
+//! (`_2 = &raw mut _1.0`) when the root local `_1` has another independent
+//! borrow. Doing so shortens the pointer's provenance and introduces UB. The
+//! reborrow through `_3` must be preserved.
+
+#![crate_type = "lib"]
+
+struct Foo(u64);
+
+impl Foo {
+    fn add(&mut self, n: u64) -> u64 {
+        self.0 + n
+    }
+}
+
+// EMIT_MIR reference_prop_mutable_alias.two_phase.ReferencePropagation.diff
+pub fn two_phase() {
+    // CHECK-LABEL: fn two_phase(
+    // The `&mut f.0` reference must survive: it carries the provenance the write
+    // travels through. The pass may still propagate the *raw* reborrow into the
+    // direct write (`(*_3) = ...`), which is the pass's intended sound transform;
+    // but it must NOT collapse the chain onto the direct place `_1.0 = ...`,
+    // which is what introduces aliasing UB w.r.t. the 2-phase `&mut _1` below.
+    // CHECK: _{{[0-9]+}} = &mut (_{{[0-9]+}}.0: u64)
+    // CHECK: (*_{{[0-9]+}}) = const 42_u64
+    // CHECK-NOT: (_{{[0-9]+}}.0: u64) = const 42_u64
+    let mut f = Foo(0);
+    let alias = &mut f.0 as *mut u64;
+    let _ = f.add(unsafe {
+        *alias = 42;
+        0
+    });
+}

tests/mir-opt/reference_prop_mutable_alias.two_phase.ReferencePropagation.diff

@@ -0,0 +1,51 @@
+- // MIR for `two_phase` before ReferencePropagation
++ // MIR for `two_phase` after ReferencePropagation
+
+  fn two_phase() -> () {
+      let mut _0: ();
+      let mut _1: Foo;
+      let mut _3: &mut u64;
+      let mut _4: u64;
+      let mut _5: &mut Foo;
+      let mut _6: u64;
+      scope 1 {
+          debug f => _1;
+          let _2: *mut u64;
+          scope 2 {
+-             debug alias => _2;
++             debug alias => _3;
+              scope 3 {
+              }
+          }
+      }
+
+      bb0: {
+          StorageLive(_1);
+          _1 = Foo(const 0_u64);
+-         StorageLive(_2);
+-         StorageLive(_3);
+          _3 = &mut (_1.0: u64);
+-         _2 = &raw mut (*_3);
+-         StorageDead(_3);
+          StorageLive(_4);
+          StorageLive(_5);
+          _5 = &mut _1;
+          StorageLive(_6);
+-         (*_2) = const 42_u64;
++         (*_3) = const 42_u64;
+          _6 = const 0_u64;
+-         _4 = Foo::add(move _5, move _6) -> [return: bb1, unwind continue];
++         _4 = Foo::add(copy _5, move _6) -> [return: bb1, unwind continue];
+      }
+
+      bb1: {
+          StorageDead(_6);
+          StorageDead(_5);
+          StorageDead(_4);
+          _0 = const ();
+-         StorageDead(_2);
+          StorageDead(_1);
+          return;
+      }
+  }
+