Use MaybeStorageDead analysis in copy_prop mir-opt to remove fewer storage statements

Author: ohadravid

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -1,8 +1,12 @@
+use std::borrow::Cow;
+
 use rustc_index::IndexSlice;
 use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
+use rustc_mir_dataflow::impls::{MaybeStorageDead, always_storage_live_locals};
+use rustc_mir_dataflow::{Analysis, ResultsCursor};
 use tracing::{debug, instrument};
 
 use crate::ssa::SsaLocals;
@@ -16,7 +20,7 @@ use crate::ssa::SsaLocals;
 ///   _d = move? _c
 /// where each of the locals is only assigned once.
 ///
-/// We want to replace all those locals by `_a`, either copied or moved.
+/// We want to replace all those locals by `_a` (the "head"), either copied or moved.
 pub(super) struct CopyProp;
 
 impl<'tcx> crate::MirPass<'tcx> for CopyProp {
@@ -34,15 +38,40 @@ impl<'tcx> crate::MirPass<'tcx> for CopyProp {
         let fully_moved = fully_moved_locals(&ssa, body);
         debug!(?fully_moved);
 
-        let mut storage_to_remove = DenseBitSet::new_empty(fully_moved.domain_size());
+        let mut head_storage_to_check = DenseBitSet::new_empty(fully_moved.domain_size());
+
         for (local, &head) in ssa.copy_classes().iter_enumerated() {
             if local != head {
-                storage_to_remove.insert(head);
+                // 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 in `maybe_storage_dead`, we have to remove the storage statements for it.
+                head_storage_to_check.insert(head);
             }
         }
 
         let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h);
 
+        let storage_to_remove = if any_replacement {
+            let always_live_locals = &always_storage_live_locals(body);
+
+            let maybe_storage_dead = MaybeStorageDead::new(Cow::Borrowed(always_live_locals))
+                .iterate_to_fixpoint(tcx, body, None)
+                .into_results_cursor(body);
+
+            let mut storage_checker = StorageChecker {
+                copy_classes: ssa.copy_classes(),
+                maybe_storage_dead,
+                head_storage_to_check,
+                storage_to_remove: DenseBitSet::new_empty(fully_moved.domain_size()),
+            };
+
+            storage_checker.visit_body(body);
+
+            storage_checker.storage_to_remove
+        } else {
+            // Will be empty anyway.
+            head_storage_to_check
+        };
+
         Replacer {
             tcx,
             copy_classes: ssa.copy_classes(),
@@ -119,6 +148,7 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
         if self.borrowed_locals.contains(*local) {
             return;
         }
+
         match ctxt {
             // Do not modify the local in storage statements.
             PlaceContext::NonUse(NonUseContext::StorageLive | NonUseContext::StorageDead) => {}
@@ -172,3 +202,29 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
         }
     }
 }
+
+struct StorageChecker<'a, 'tcx> {
+    storage_to_remove: DenseBitSet<Local>,
+    head_storage_to_check: DenseBitSet<Local>,
+    maybe_storage_dead: ResultsCursor<'a, 'tcx, MaybeStorageDead<'a>>,
+    copy_classes: &'a IndexSlice<Local, Local>,
+}
+
+impl<'a, 'tcx> Visitor<'tcx> for StorageChecker<'a, 'tcx> {
+    fn visit_local(&mut self, local: Local, context: PlaceContext, location: Location) {
+        let head = self.copy_classes[local];
+
+        if context.is_use() && self.head_storage_to_check.contains(head) {
+            self.maybe_storage_dead.seek_after_primary_effect(location);
+            if self.maybe_storage_dead.get().contains(head) {
+                debug!(
+                    ?location,
+                    ?local,
+                    ?head,
+                    "found use of local with head at a location in which head is maybe dead, marking head for storage removal"
+                );
+                self.storage_to_remove.insert(head);
+            }
+        }
+    }
+}

tests/mir-opt/copy-prop/cycle.main.CopyProp.panic-abort.diff

@@ -26,7 +26,7 @@
       }
 
       bb1: {
--         StorageLive(_2);
+          StorageLive(_2);
           _2 = copy _1;
 -         StorageLive(_3);
 -         _3 = copy _2;
@@ -46,7 +46,7 @@
           StorageDead(_5);
           _0 = const ();
 -         StorageDead(_3);
--         StorageDead(_2);
+          StorageDead(_2);
           StorageDead(_1);
           return;
       }

tests/mir-opt/copy-prop/cycle.main.CopyProp.panic-unwind.diff

@@ -26,7 +26,7 @@
       }
 
       bb1: {
--         StorageLive(_2);
+          StorageLive(_2);
           _2 = copy _1;
 -         StorageLive(_3);
 -         _3 = copy _2;
@@ -46,7 +46,7 @@
           StorageDead(_5);
           _0 = const ();
 -         StorageDead(_3);
--         StorageDead(_2);
+          StorageDead(_2);
           StorageDead(_1);
           return;
       }

tests/mir-opt/copy-prop/dead_stores_79191.f.CopyProp.after.panic-abort.mir

@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
     }
 
     bb0: {
+        StorageLive(_2);
         _2 = copy _1;
         _1 = const 5_usize;
         _1 = copy _2;
@@ -21,6 +22,7 @@ fn f(_1: usize) -> usize {
 
     bb1: {
         StorageDead(_4);
+        StorageDead(_2);
         return;
     }
 }

tests/mir-opt/copy-prop/dead_stores_79191.f.CopyProp.after.panic-unwind.mir

@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
     }
 
     bb0: {
+        StorageLive(_2);
         _2 = copy _1;
         _1 = const 5_usize;
         _1 = copy _2;
@@ -21,6 +22,7 @@ fn f(_1: usize) -> usize {
 
     bb1: {
         StorageDead(_4);
+        StorageDead(_2);
         return;
     }
 }

tests/mir-opt/copy-prop/dead_stores_better.f.CopyProp.after.panic-abort.mir

@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
     }
 
     bb0: {
+        StorageLive(_2);
         _2 = copy _1;
         _1 = const 5_usize;
         _1 = copy _2;
@@ -21,6 +22,7 @@ fn f(_1: usize) -> usize {
 
     bb1: {
         StorageDead(_4);
+        StorageDead(_2);
         return;
     }
 }

tests/mir-opt/copy-prop/dead_stores_better.f.CopyProp.after.panic-unwind.mir

@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
     }
 
     bb0: {
+        StorageLive(_2);
         _2 = copy _1;
         _1 = const 5_usize;
         _1 = copy _2;
@@ -21,6 +22,7 @@ fn f(_1: usize) -> usize {
 
     bb1: {
         StorageDead(_4);
+        StorageDead(_2);
         return;
     }
 }

tests/mir-opt/copy-prop/issue_107511.main.CopyProp.panic-abort.diff

@@ -86,7 +86,7 @@
       }
 
       bb6: {
--         StorageLive(_16);
+          StorageLive(_16);
           _16 = copy ((_11 as Some).0: usize);
           StorageLive(_17);
 -         StorageLive(_18);
@@ -116,7 +116,7 @@
           StorageDead(_17);
 -         StorageDead(_18);
 -         _10 = const ();
--         StorageDead(_16);
+          StorageDead(_16);
           StorageDead(_13);
           StorageDead(_11);
 -         StorageDead(_10);

tests/mir-opt/copy-prop/issue_107511.main.CopyProp.panic-unwind.diff

@@ -86,7 +86,7 @@
       }
 
       bb6: {
--         StorageLive(_16);
+          StorageLive(_16);
           _16 = copy ((_11 as Some).0: usize);
           StorageLive(_17);
 -         StorageLive(_18);
@@ -116,7 +116,7 @@
           StorageDead(_17);
 -         StorageDead(_18);
 -         _10 = const ();
--         StorageDead(_16);
+          StorageDead(_16);
           StorageDead(_13);
           StorageDead(_11);
 -         StorageDead(_10);

tests/mir-opt/copy-prop/issue_141649.main.CopyProp.panic-abort.diff

@@ -0,0 +1,100 @@
+- // MIR for `main` before CopyProp
++ // MIR for `main` after CopyProp
+  
+  fn main() -> () {
+      let mut _0: ();
+      let _1: ();
+      let _2: main::S;
+      let _3: ();
+      let mut _4: main::S;
+      let _5: ();
+      let _6: main::S;
+      let _7: ();
+      let mut _8: main::S;
+      let _9: ();
+      let _10: main::C;
+      let _11: ();
+      let mut _12: main::C;
+      let _13: main::C;
+      let _14: ();
+      let mut _15: main::C;
+      scope 1 {
+          debug s1 => _2;
+      }
+      scope 2 {
+          debug s2 => _6;
+      }
+      scope 3 {
+          debug c1 => _10;
+      }
+      scope 4 {
+          debug c2 => _13;
+      }
+  
+      bb0: {
+-         StorageLive(_1);
+          StorageLive(_2);
+          _2 = S(const 1_usize, const 2_usize);
+          StorageLive(_3);
+-         StorageLive(_4);
+-         _4 = move _2;
+-         _3 = std::mem::drop::<S>(move _4) -> [return: bb1, unwind unreachable];
++         _3 = std::mem::drop::<S>(move _2) -> [return: bb1, unwind unreachable];
+      }
+  
+      bb1: {
+-         StorageDead(_4);
+          StorageDead(_3);
+-         _1 = const ();
+          StorageDead(_2);
+-         StorageDead(_1);
+-         StorageLive(_5);
+          StorageLive(_6);
+          _6 = S(const 3_usize, const 4_usize);
+          StorageLive(_7);
+-         StorageLive(_8);
+-         _8 = move _6;
+-         _7 = std::mem::drop::<S>(move _8) -> [return: bb2, unwind unreachable];
++         _7 = std::mem::drop::<S>(move _6) -> [return: bb2, unwind unreachable];
+      }
+  
+      bb2: {
+-         StorageDead(_8);
+          StorageDead(_7);
+-         _5 = const ();
+          StorageDead(_6);
+-         StorageDead(_5);
+-         StorageLive(_9);
+          StorageLive(_10);
+          _10 = C(const 1_usize, const 2_usize);
+          StorageLive(_11);
+-         StorageLive(_12);
+-         _12 = copy _10;
+-         _11 = std::mem::drop::<C>(move _12) -> [return: bb3, unwind unreachable];
++         _11 = std::mem::drop::<C>(copy _10) -> [return: bb3, unwind unreachable];
+      }
+  
+      bb3: {
+-         StorageDead(_12);
+          StorageDead(_11);
+-         _9 = const ();
+          StorageDead(_10);
+-         StorageDead(_9);
+          StorageLive(_13);
+          _13 = C(const 3_usize, const 4_usize);
+          StorageLive(_14);
+-         StorageLive(_15);
+-         _15 = copy _13;
+-         _14 = std::mem::drop::<C>(move _15) -> [return: bb4, unwind unreachable];
++         _14 = std::mem::drop::<C>(copy _13) -> [return: bb4, unwind unreachable];
+      }
+  
+      bb4: {
+-         StorageDead(_15);
+          StorageDead(_14);
+          _0 = const ();
+          StorageDead(_13);
+          return;
+      }
+  }
+