Propagate from borrowed locals in CopyProp

compiler/rustc_mir_transform/src/ssa.rs

@@ -293,10 +293,6 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor<'_, 'tcx> {
 fn compute_copy_classes(ssa: &mut SsaLocals, body: &Body<'_>) {
     let mut direct_uses = std::mem::take(&mut ssa.direct_uses);
     let mut copies = IndexVec::from_fn_n(|l| l, body.local_decls.len());
-    // We must not unify two locals that are borrowed. But this is fine if one is borrowed and
-    // the other is not. This bitset is keyed by *class head* and contains whether any member of
-    // the class is borrowed.
-    let mut borrowed_classes = ssa.borrowed_locals().clone();
 
     for (local, rvalue, _) in ssa.assignments(body) {
         let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place))
@@ -322,8 +318,12 @@ fn compute_copy_classes(ssa: &mut SsaLocals, body: &Body<'_>) {
         // visited before `local`, and we just have to copy the representing local.
         let head = copies[rhs];
 
-        // Do not unify borrowed locals.
-        if borrowed_classes.contains(local) || borrowed_classes.contains(head) {
+        // When propagating from `head` to `local` we need to ensure that changes to the address
+        // are not observable, so at most one the locals involved can be borrowed. Additionally, we
+        // need to ensure that the definition of `head` dominates all uses of `local`. When `local`
+        // is borrowed, there might exist an indirect use of `local` that isn't dominated by the
+        // definition, so we have to reject copy propagation.
+        if ssa.borrowed_locals().contains(local) {
             continue;
         }
 
@@ -339,21 +339,14 @@ fn compute_copy_classes(ssa: &mut SsaLocals, body: &Body<'_>) {
                     *h = RETURN_PLACE;
                 }
             }
-            if borrowed_classes.contains(head) {
-                borrowed_classes.insert(RETURN_PLACE);
-            }
         } else {
             copies[local] = head;
-            if borrowed_classes.contains(local) {
-                borrowed_classes.insert(head);
-            }
         }
         direct_uses[rhs] -= 1;
     }
 
     debug!(?copies);
     debug!(?direct_uses);
-    debug!(?borrowed_classes);
 
     // Invariant: `copies` must point to the head of an equivalence class.
     #[cfg(debug_assertions)]
@@ -362,13 +355,6 @@ fn compute_copy_classes(ssa: &mut SsaLocals, body: &Body<'_>) {
     }
     debug_assert_eq!(copies[RETURN_PLACE], RETURN_PLACE);
 
-    // Invariant: `borrowed_classes` must be true if any member of the class is borrowed.
-    #[cfg(debug_assertions)]
-    for &head in copies.iter() {
-        let any_borrowed = ssa.borrowed_locals.iter().any(|l| copies[l] == head);
-        assert_eq!(borrowed_classes.contains(head), any_borrowed);
-    }
-
     ssa.direct_uses = direct_uses;
     ssa.copy_classes = copies;
 }

tests/mir-opt/copy-prop/borrowed_local.borrowed.CopyProp.panic-abort.diff

@@ -7,13 +7,14 @@
       let mut _3: &T;
 
       bb0: {
-          _2 = copy _1;
+-         _2 = copy _1;
           _3 = &_1;
           _0 = opaque::<&T>(copy _3) -> [return: bb1, unwind unreachable];
       }
 
       bb1: {
-          _0 = opaque::<T>(copy _2) -> [return: bb2, unwind unreachable];
+-         _0 = opaque::<T>(copy _2) -> [return: bb2, unwind unreachable];
++         _0 = opaque::<T>(copy _1) -> [return: bb2, unwind unreachable];
       }
 
       bb2: {

tests/mir-opt/copy-prop/borrowed_local.borrowed.CopyProp.panic-unwind.diff

@@ -7,13 +7,14 @@
       let mut _3: &T;
 
       bb0: {
-          _2 = copy _1;
+-         _2 = copy _1;
           _3 = &_1;
           _0 = opaque::<&T>(copy _3) -> [return: bb1, unwind continue];
       }
 
       bb1: {
-          _0 = opaque::<T>(copy _2) -> [return: bb2, unwind continue];
+-         _0 = opaque::<T>(copy _2) -> [return: bb2, unwind continue];
++         _0 = opaque::<T>(copy _1) -> [return: bb2, unwind continue];
       }
 
       bb2: {

tests/mir-opt/copy-prop/borrowed_local.rs

@@ -50,11 +50,10 @@ fn compare_address() -> bool {
 fn borrowed<T: Copy + Freeze>(x: T) -> bool {
     // CHECK-LABEL: fn borrowed(
     // CHECK: bb0: {
-    // CHECK-NEXT: _2 = copy _1;
     // CHECK-NEXT: _3 = &_1;
     // CHECK-NEXT: _0 = opaque::<&T>(copy _3)
     // CHECK: bb1: {
-    // CHECK-NEXT: _0 = opaque::<T>(copy _2)
+    // CHECK-NEXT: _0 = opaque::<T>(copy _1)
     mir! {
         {
             let a = x;

tests/mir-opt/copy-prop/write_to_borrowed.main.CopyProp.diff

@@ -16,10 +16,11 @@
           _3 = const 'b';
           _5 = copy _3;
           _6 = &_3;
-          _4 = copy _5;
+-         _4 = copy _5;
           (*_1) = copy (*_6);
           _6 = &_5;
-          _7 = dump_var::<char>(copy _4) -> [return: bb1, unwind unreachable];
+-         _7 = dump_var::<char>(copy _4) -> [return: bb1, unwind unreachable];
++         _7 = dump_var::<char>(copy _5) -> [return: bb1, unwind unreachable];
       }
 
       bb1: {

tests/mir-opt/copy-prop/write_to_borrowed.rs

@@ -27,13 +27,12 @@ fn main() {
             _5 = _3;
             // CHECK-NEXT: _6 = &_3;
             _6 = &_3;
-            // CHECK-NEXT: _4 = copy _5;
             _4 = _5;
             // CHECK-NEXT: (*_1) = copy (*_6);
             *_1 = *_6;
             // CHECK-NEXT: _6 = &_5;
             _6 = &_5;
-            // CHECK-NEXT: _7 = dump_var::<char>(copy _4)
+            // CHECK-NEXT: _7 = dump_var::<char>(copy _5)
             Call(_7 = dump_var(_4), ReturnTo(bb1), UnwindUnreachable())
         }
         bb1 = { Return() }