Propagate from borrowed locals in CopyProp

7 files changed, 97 insertions, 142 deletions

diff --git a/compiler/rustc_mir_transform/src/ssa.rs b/compiler/rustc_mir_transform/src/ssa.rs
index d3b4b99e932..cd9a7f4a39d 100644
--- a/compiler/rustc_mir_transform/src/ssa.rs
+++ b/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;
 }
diff --git a/tests/mir-opt/copy-prop/borrowed_local.borrowed.CopyProp.panic-abort.diff b/tests/mir-opt/copy-prop/borrowed_local.borrowed.CopyProp.panic-abort.diff
index 285cd0f6527..40e8c06f357 100644
--- a/tests/mir-opt/copy-prop/borrowed_local.borrowed.CopyProp.panic-abort.diff
+++ b/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: {
diff --git a/tests/mir-opt/copy-prop/borrowed_local.borrowed.CopyProp.panic-unwind.diff b/tests/mir-opt/copy-prop/borrowed_local.borrowed.CopyProp.panic-unwind.diff
index f189615ea95..d09c96c0f2b 100644
--- a/tests/mir-opt/copy-prop/borrowed_local.borrowed.CopyProp.panic-unwind.diff
+++ b/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: {
diff --git a/tests/mir-opt/copy-prop/borrowed_local.rs b/tests/mir-opt/copy-prop/borrowed_local.rs
index 68cdc57483a..08c46fea5f0 100644
--- a/tests/mir-opt/copy-prop/borrowed_local.rs
+++ b/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;
diff --git a/tests/mir-opt/copy-prop/write_to_borrowed.main.CopyProp.diff b/tests/mir-opt/copy-prop/write_to_borrowed.main.CopyProp.diff
index baa71501047..eab06b1ba1e 100644
--- a/tests/mir-opt/copy-prop/write_to_borrowed.main.CopyProp.diff
+++ b/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: {
diff --git a/tests/mir-opt/copy-prop/write_to_borrowed.rs b/tests/mir-opt/copy-prop/write_to_borrowed.rs
index 05aa2fba18d..06e617b8bfb 100644
--- a/tests/mir-opt/copy-prop/write_to_borrowed.rs
+++ b/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() }
diff --git a/tests/mir-opt/pre-codegen/slice_filter.variant_a-{closure#0}.PreCodegen.after.mir b/tests/mir-opt/pre-codegen/slice_filter.variant_a-{closure#0}.PreCodegen.after.mir
index 97036745009..cbdd194afd3 100644
--- a/tests/mir-opt/pre-codegen/slice_filter.variant_a-{closure#0}.PreCodegen.after.mir
+++ b/tests/mir-opt/pre-codegen/slice_filter.variant_a-{closure#0}.PreCodegen.after.mir
@@ -10,18 +10,18 @@ fn variant_a::{closure#0}(_1: &mut {closure@$DIR/slice_filter.rs:7:25: 7:39}, _2
     let mut _8: &&usize;
     let _9: &usize;
     let mut _10: &&usize;
-    let mut _15: bool;
+    let mut _13: bool;
+    let mut _14: &&usize;
+    let _15: &usize;
     let mut _16: &&usize;
-    let _17: &usize;
-    let mut _18: &&usize;
+    let mut _19: bool;
+    let mut _20: &&usize;
+    let _21: &usize;
+    let mut _22: &&usize;
     let mut _23: bool;
     let mut _24: &&usize;
     let _25: &usize;
     let mut _26: &&usize;
-    let mut _29: bool;
-    let mut _30: &&usize;
-    let _31: &usize;
-    let mut _32: &&usize;
     scope 1 {
         debug a => _4;
         debug b => _5;
@@ -30,47 +30,39 @@ fn variant_a::{closure#0}(_1: &mut {closure@$DIR/slice_filter.rs:7:25: 7:39}, _2
         scope 2 (inlined std::cmp::impls::<impl PartialOrd for &usize>::le) {
             debug self => _8;
             debug other => _10;
-            let mut _11: &usize;
-            let mut _12: &usize;
             scope 3 (inlined std::cmp::impls::<impl PartialOrd for usize>::le) {
-                debug self => _11;
-                debug other => _12;
-                let mut _13: usize;
-                let mut _14: usize;
+                debug self => _4;
+                debug other => _6;
+                let mut _11: usize;
+                let mut _12: usize;
             }
         }
         scope 4 (inlined std::cmp::impls::<impl PartialOrd for &usize>::le) {
-            debug self => _16;
-            debug other => _18;
-            let mut _19: &usize;
-            let mut _20: &usize;
+            debug self => _14;
+            debug other => _16;
             scope 5 (inlined std::cmp::impls::<impl PartialOrd for usize>::le) {
-                debug self => _19;
-                debug other => _20;
-                let mut _21: usize;
-                let mut _22: usize;
+                debug self => _7;
+                debug other => _5;
+                let mut _17: usize;
+                let mut _18: usize;
             }
         }
         scope 6 (inlined std::cmp::impls::<impl PartialOrd for &usize>::le) {
-            debug self => _24;
-            debug other => _26;
-            let mut _27: &usize;
-            let mut _28: &usize;
+            debug self => _20;
+            debug other => _22;
             scope 7 (inlined std::cmp::impls::<impl PartialOrd for usize>::le) {
-                debug self => _27;
-                debug other => _28;
+                debug self => _6;
+                debug other => _4;
             }
         }
         scope 8 (inlined std::cmp::impls::<impl PartialOrd for &usize>::le) {
-            debug self => _30;
-            debug other => _32;
-            let mut _33: &usize;
-            let mut _34: &usize;
+            debug self => _24;
+            debug other => _26;
             scope 9 (inlined std::cmp::impls::<impl PartialOrd for usize>::le) {
-                debug self => _33;
-                debug other => _34;
-                let mut _35: usize;
-                let mut _36: usize;
+                debug self => _5;
+                debug other => _7;
+                let mut _27: usize;
+                let mut _28: usize;
             }
         }
     }
@@ -81,23 +73,17 @@ fn variant_a::{closure#0}(_1: &mut {closure@$DIR/slice_filter.rs:7:25: 7:39}, _2
         _5 = &((*_3).1: usize);
         _6 = &((*_3).2: usize);
         _7 = &((*_3).3: usize);
-        StorageLive(_15);
+        StorageLive(_13);
         StorageLive(_8);
         _8 = &_4;
         StorageLive(_10);
         StorageLive(_9);
         _9 = copy _6;
         _10 = &_9;
-        StorageLive(_11);
-        StorageLive(_12);
-        _11 = copy _4;
-        _12 = copy _6;
-        _13 = copy ((*_3).0: usize);
-        _14 = copy ((*_3).2: usize);
-        _15 = Le(copy _13, copy _14);
-        StorageDead(_12);
-        StorageDead(_11);
-        switchInt(move _15) -> [0: bb1, otherwise: bb2];
+        _11 = copy ((*_3).0: usize);
+        _12 = copy ((*_3).2: usize);
+        _13 = Le(copy _11, copy _12);
+        switchInt(move _13) -> [0: bb1, otherwise: bb2];
     }
 
     bb1: {
@@ -111,107 +97,89 @@ fn variant_a::{closure#0}(_1: &mut {closure@$DIR/slice_filter.rs:7:25: 7:39}, _2
         StorageDead(_9);
         StorageDead(_10);
         StorageDead(_8);
-        StorageLive(_23);
+        StorageLive(_19);
+        StorageLive(_14);
+        _14 = &_7;
         StorageLive(_16);
-        _16 = &_7;
-        StorageLive(_18);
+        StorageLive(_15);
+        _15 = copy _5;
+        _16 = &_15;
         StorageLive(_17);
-        _17 = copy _5;
-        _18 = &_17;
-        StorageLive(_19);
-        StorageLive(_20);
-        _19 = copy _7;
-        _20 = copy _5;
-        StorageLive(_21);
-        _21 = copy ((*_3).3: usize);
-        StorageLive(_22);
-        _22 = copy ((*_3).1: usize);
-        _23 = Le(move _21, move _22);
-        StorageDead(_22);
-        StorageDead(_21);
-        StorageDead(_20);
-        StorageDead(_19);
-        switchInt(move _23) -> [0: bb3, otherwise: bb8];
+        _17 = copy ((*_3).3: usize);
+        StorageLive(_18);
+        _18 = copy ((*_3).1: usize);
+        _19 = Le(move _17, move _18);
+        StorageDead(_18);
+        StorageDead(_17);
+        switchInt(move _19) -> [0: bb3, otherwise: bb8];
     }
 
     bb3: {
-        StorageDead(_17);
-        StorageDead(_18);
+        StorageDead(_15);
         StorageDead(_16);
+        StorageDead(_14);
         goto -> bb4;
     }
 
     bb4: {
-        StorageLive(_29);
-        StorageLive(_24);
-        _24 = &_6;
-        StorageLive(_26);
-        StorageLive(_25);
-        _25 = copy _4;
-        _26 = &_25;
-        StorageLive(_27);
-        StorageLive(_28);
-        _27 = copy _6;
-        _28 = copy _4;
-        _29 = Le(copy _14, copy _13);
-        StorageDead(_28);
-        StorageDead(_27);
-        switchInt(move _29) -> [0: bb5, otherwise: bb6];
+        StorageLive(_23);
+        StorageLive(_20);
+        _20 = &_6;
+        StorageLive(_22);
+        StorageLive(_21);
+        _21 = copy _4;
+        _22 = &_21;
+        _23 = Le(copy _12, copy _11);
+        switchInt(move _23) -> [0: bb5, otherwise: bb6];
     }
 
     bb5: {
-        StorageDead(_25);
-        StorageDead(_26);
-        StorageDead(_24);
+        StorageDead(_21);
+        StorageDead(_22);
+        StorageDead(_20);
         _0 = const false;
         goto -> bb7;
     }
 
     bb6: {
+        StorageDead(_21);
+        StorageDead(_22);
+        StorageDead(_20);
+        StorageLive(_24);
+        _24 = &_5;
+        StorageLive(_26);
+        StorageLive(_25);
+        _25 = copy _7;
+        _26 = &_25;
+        StorageLive(_27);
+        _27 = copy ((*_3).1: usize);
+        StorageLive(_28);
+        _28 = copy ((*_3).3: usize);
+        _0 = Le(move _27, move _28);
+        StorageDead(_28);
+        StorageDead(_27);
         StorageDead(_25);
         StorageDead(_26);
         StorageDead(_24);
-        StorageLive(_30);
-        _30 = &_5;
-        StorageLive(_32);
-        StorageLive(_31);
-        _31 = copy _7;
-        _32 = &_31;
-        StorageLive(_33);
-        StorageLive(_34);
-        _33 = copy _5;
-        _34 = copy _7;
-        StorageLive(_35);
-        _35 = copy ((*_3).1: usize);
-        StorageLive(_36);
-        _36 = copy ((*_3).3: usize);
-        _0 = Le(move _35, move _36);
-        StorageDead(_36);
-        StorageDead(_35);
-        StorageDead(_34);
-        StorageDead(_33);
-        StorageDead(_31);
-        StorageDead(_32);
-        StorageDead(_30);
         goto -> bb7;
     }
 
     bb7: {
-        StorageDead(_29);
+        StorageDead(_23);
         goto -> bb9;
     }
 
     bb8: {
-        StorageDead(_17);
-        StorageDead(_18);
+        StorageDead(_15);
         StorageDead(_16);
+        StorageDead(_14);
         _0 = const true;
         goto -> bb9;
     }
 
     bb9: {
-        StorageDead(_23);
-        StorageDead(_15);
+        StorageDead(_19);
+        StorageDead(_13);
         return;
     }
 }