Use a single bitset to check the storage in copy_prop, only check reachable blocks

Author: ohadravid

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -32,30 +32,20 @@ impl<'tcx> crate::MirPass<'tcx> for CopyProp {
 
         let typing_env = body.typing_env(tcx);
         let ssa = SsaLocals::new(tcx, body, typing_env);
-        let borrowed_locals = ssa.borrowed_locals().clone();
 
-        debug!(?borrowed_locals);
+        debug!(borrowed_locals = ?ssa.borrowed_locals());
         debug!(copy_classes = ?ssa.copy_classes());
 
         let mut any_replacement = false;
         let fully_moved = fully_moved_locals(&ssa, body);
         debug!(?fully_moved);
 
-        let mut head_storage_to_check = DenseBitSet::new_empty(fully_moved.domain_size());
         let mut storage_to_remove = DenseBitSet::new_empty(fully_moved.domain_size());
 
         for (local, &head) in ssa.copy_classes().iter_enumerated() {
             if local != head {
                 any_replacement = true;
-                // We need to determine if we can keep the head's storage statements (which enables better optimizations).
-                // For every local's usage location, if the head is maybe-uninitialized, we'll need to remove it's storage statements.
-                head_storage_to_check.insert(head);
-
-                if borrowed_locals.contains(local) {
-                    // To keep the storage of a head, we require that none of the locals in it's copy class are borrowed,
-                    // since otherwise we cannot easily identify when it is used.
-                    storage_to_remove.insert(head);
-                }
+                storage_to_remove.insert(head);
             }
         }
 
@@ -66,25 +56,31 @@ impl<'tcx> crate::MirPass<'tcx> for CopyProp {
         let fully_moved = fully_moved_locals(&ssa, body);
         debug!(?fully_moved);
 
-        // Debug builds have no use for the storage statements, so avoid extra work.
-        let storage_to_remove = if any_replacement && tcx.sess.emit_lifetime_markers() {
+        // We can determine if we can keep the head's storage statements (which enables better optimizations).
+        // For every local's usage location, we'll to remove it's storage statements only if the head is maybe-uninitialized,
+        // or if the local is borrowed (since we cannot easily identify when it is used).
+        let storage_to_remove = if tcx.sess.emit_lifetime_markers() {
+            storage_to_remove.clear();
+
             let maybe_uninit = MaybeUninitializedLocals::new()
                 .iterate_to_fixpoint(tcx, body, Some("mir_opt::copy_prop"))
                 .into_results_cursor(body);
 
             let mut storage_checker = StorageChecker {
                 maybe_uninit,
                 copy_classes: ssa.copy_classes(),
-                head_storage_to_check,
+                borrowed_locals: ssa.borrowed_locals(),
                 storage_to_remove,
             };
 
-            storage_checker.visit_body(body);
+            for (bb, data) in traversal::reachable(body) {
+                storage_checker.visit_basic_block_data(bb, data);
+            }
 
             storage_checker.storage_to_remove
         } else {
             // Conservatively remove all storage statements for the head locals.
-            head_storage_to_check
+            storage_to_remove
         };
 
         debug!(?storage_to_remove);
@@ -200,7 +196,7 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
 struct StorageChecker<'a, 'tcx> {
     maybe_uninit: ResultsCursor<'a, 'tcx, MaybeUninitializedLocals>,
     copy_classes: &'a IndexSlice<Local, Local>,
-    head_storage_to_check: DenseBitSet<Local>,
+    borrowed_locals: &'a DenseBitSet<Local>,
     storage_to_remove: DenseBitSet<Local>,
 }
 
@@ -222,23 +218,28 @@ impl<'a, 'tcx> Visitor<'tcx> for StorageChecker<'a, 'tcx> {
 
         let head = self.copy_classes[local];
 
-        // The head must be initialized at the location of the local, otherwise we must remove it's storage statements.
-        if self.head_storage_to_check.contains(head) {
-            self.maybe_uninit.seek_before_primary_effect(loc);
-
-            if self.maybe_uninit.get().contains(head) {
-                debug!(
-                    ?loc,
-                    ?context,
-                    ?local,
-                    ?head,
-                    "found a head at a location in which it is maybe uninit, marking head for storage statement removal"
-                );
-                self.storage_to_remove.insert(head);
-
-                // Once we found a use of the head that is maybe uninit, we do not need to check it again.
-                self.head_storage_to_check.remove(head);
-            }
+        // If the local is the head, or if we already marked it for deletion, we do not need to check it.
+        if head == local || self.storage_to_remove.contains(head) {
+            return;
+        }
+
+        // If the local is borrowed, we cannot easily determine if it is used, so we have to remove the storage statements.
+        if self.borrowed_locals.contains(local) {
+            self.storage_to_remove.insert(head);
+            return;
+        }
+
+        self.maybe_uninit.seek_before_primary_effect(loc);
+
+        if self.maybe_uninit.get().contains(head) {
+            debug!(
+                ?loc,
+                ?context,
+                ?local,
+                ?head,
+                "found a head at a location in which it is maybe uninit, marking head for storage statement removal"
+            );
+            self.storage_to_remove.insert(head);
         }
     }
 }

tests/mir-opt/copy-prop/copy_prop_storage_unreachable.f.CopyProp.diff

@@ -0,0 +1,27 @@
+- // MIR for `f` before CopyProp
++ // MIR for `f` after CopyProp
+  
+  fn f(_1: &mut usize) -> () {
+      let mut _0: ();
+      let mut _2: usize;
+      let mut _3: usize;
+  
+      bb0: {
+          StorageLive(_2);
+          _2 = const 42_usize;
+-         _3 = copy _2;
+-         (*_1) = copy _3;
++         (*_1) = copy _2;
+          StorageDead(_2);
+          return;
+      }
+  
+      bb1: {
+          StorageLive(_2);
+-         (*_1) = copy _3;
++         (*_1) = copy _2;
+          StorageDead(_2);
+          return;
+      }
+  }
+  

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

@@ -0,0 +1,35 @@
+// skip-filecheck
+//@ test-mir-pass: CopyProp
+
+#![feature(custom_mir, core_intrinsics)]
+
+// Check that we do not remove the storage statements if the head
+// is uninitialized in an unreachable block.
+
+use std::intrinsics::mir::*;
+
+// EMIT_MIR copy_prop_storage_unreachable.f.CopyProp.diff
+
+#[custom_mir(dialect = "runtime", phase = "post-cleanup")]
+pub fn f(_1: &mut usize) {
+    mir! {
+        let _2: usize;
+        let _3: usize;
+        {
+            StorageLive(_2);
+            _2 = 42;
+            _3 = _2;
+            (*_1) = _3;
+            StorageDead(_2);
+            Return()
+        }
+        bb1 = {
+            // Ensure that _2 is considered uninitialized by `MaybeUninitializedLocals`.
+            StorageLive(_2);
+            // Use of _3 (in an unreachable block) when definition of _2 is unavailable.
+            (*_1) = _3;
+            StorageDead(_2);
+            Return()
+        }
+    }
+}

tests/mir-opt/pre-codegen/range_iter.forward_loop.PreCodegen.after.panic-abort.mir

@@ -49,6 +49,7 @@ fn forward_loop(_1: u32, _2: u32, _3: impl Fn(u32)) -> () {
 
     bb1: {
         StorageLive(_9);
+        StorageLive(_7);
         StorageLive(_6);
         StorageLive(_5);
         _5 = copy _4;
@@ -59,6 +60,7 @@ fn forward_loop(_1: u32, _2: u32, _3: impl Fn(u32)) -> () {
 
     bb2: {
         StorageDead(_6);
+        StorageDead(_7);
         StorageDead(_9);
         StorageDead(_4);
         drop(_3) -> [return: bb3, unwind unreachable];
@@ -69,15 +71,14 @@ fn forward_loop(_1: u32, _2: u32, _3: impl Fn(u32)) -> () {
     }
 
     bb4: {
-        StorageLive(_7);
         _7 = copy _4;
         StorageLive(_8);
         _8 = AddUnchecked(copy _7, const 1_u32);
         _4 = move _8;
         StorageDead(_8);
         _9 = Option::<u32>::Some(copy _7);
-        StorageDead(_7);
         StorageDead(_6);
+        StorageDead(_7);
         StorageLive(_10);
         _10 = copy ((_9 as Some).0: u32);
         StorageLive(_11);

tests/mir-opt/pre-codegen/range_iter.forward_loop.PreCodegen.after.panic-unwind.mir

@@ -49,6 +49,7 @@ fn forward_loop(_1: u32, _2: u32, _3: impl Fn(u32)) -> () {
 
     bb1: {
         StorageLive(_9);
+        StorageLive(_7);
         StorageLive(_6);
         StorageLive(_5);
         _5 = copy _4;
@@ -59,6 +60,7 @@ fn forward_loop(_1: u32, _2: u32, _3: impl Fn(u32)) -> () {
 
     bb2: {
         StorageDead(_6);
+        StorageDead(_7);
         StorageDead(_9);
         StorageDead(_4);
         drop(_3) -> [return: bb3, unwind continue];
@@ -69,15 +71,14 @@ fn forward_loop(_1: u32, _2: u32, _3: impl Fn(u32)) -> () {
     }
 
     bb4: {
-        StorageLive(_7);
         _7 = copy _4;
         StorageLive(_8);
         _8 = AddUnchecked(copy _7, const 1_u32);
         _4 = move _8;
         StorageDead(_8);
         _9 = Option::<u32>::Some(copy _7);
-        StorageDead(_7);
         StorageDead(_6);
+        StorageDead(_7);
         StorageLive(_10);
         _10 = copy ((_9 as Some).0: u32);
         StorageLive(_11);

tests/mir-opt/pre-codegen/range_iter.range_iter_next.PreCodegen.after.panic-abort.mir

@@ -26,6 +26,7 @@ fn range_iter_next(_1: &mut std::ops::Range<u32>) -> Option<u32> {
     }
 
     bb0: {
+        StorageLive(_5);
         StorageLive(_4);
         StorageLive(_2);
         _2 = copy ((*_1).0: u32);
@@ -43,19 +44,18 @@ fn range_iter_next(_1: &mut std::ops::Range<u32>) -> Option<u32> {
     }
 
     bb2: {
-        StorageLive(_5);
         _5 = copy ((*_1).0: u32);
         StorageLive(_6);
         _6 = AddUnchecked(copy _5, const 1_u32);
         ((*_1).0: u32) = move _6;
         StorageDead(_6);
         _0 = Option::<u32>::Some(copy _5);
-        StorageDead(_5);
         goto -> bb3;
     }
 
     bb3: {
         StorageDead(_4);
+        StorageDead(_5);
         return;
     }
 }

tests/mir-opt/pre-codegen/range_iter.range_iter_next.PreCodegen.after.panic-unwind.mir

@@ -26,6 +26,7 @@ fn range_iter_next(_1: &mut std::ops::Range<u32>) -> Option<u32> {
     }
 
     bb0: {
+        StorageLive(_5);
         StorageLive(_4);
         StorageLive(_2);
         _2 = copy ((*_1).0: u32);
@@ -43,19 +44,18 @@ fn range_iter_next(_1: &mut std::ops::Range<u32>) -> Option<u32> {
     }
 
     bb2: {
-        StorageLive(_5);
         _5 = copy ((*_1).0: u32);
         StorageLive(_6);
         _6 = AddUnchecked(copy _5, const 1_u32);
         ((*_1).0: u32) = move _6;
         StorageDead(_6);
         _0 = Option::<u32>::Some(copy _5);
-        StorageDead(_5);
         goto -> bb3;
     }
 
     bb3: {
         StorageDead(_4);
+        StorageDead(_5);
         return;
     }
 }