Remove fewer Storage calls in copy_prop

compiler/rustc_middle/src/mir/mod.rs

@@ -1189,8 +1189,8 @@ impl<'tcx> LocalDecl<'tcx> {
 
     /// Returns `true` if this is a DerefTemp
     pub fn is_deref_temp(&self) -> bool {
-        match self.local_info() {
-            LocalInfo::DerefTemp => true,
+        match self.local_info.as_ref() {
+            ClearCrossCrate::Set(box LocalInfo::DerefTemp) => true,
             _ => false,
         }
     }

compiler/rustc_mir_dataflow/src/move_paths/builder.rs

@@ -1,9 +1,9 @@
 use std::mem;
 
 use rustc_index::IndexVec;
+use rustc_middle::bug;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, Ty, TyCtxt, TypeVisitableExt};
-use rustc_middle::{bug, span_bug};
 use smallvec::{SmallVec, smallvec};
 use tracing::debug;
 
@@ -344,13 +344,14 @@ impl<'a, 'tcx, F: Fn(Ty<'tcx>) -> bool> MoveDataBuilder<'a, 'tcx, F> {
         match &stmt.kind {
             StatementKind::Assign(box (place, Rvalue::CopyForDeref(reffed))) => {
                 let local = place.as_local().unwrap();
-                assert!(self.body.local_decls[local].is_deref_temp());
+                if self.body.local_decls[local].is_deref_temp() {
+                    let rev_lookup = &mut self.data.rev_lookup;
 
-                let rev_lookup = &mut self.data.rev_lookup;
-
-                rev_lookup.un_derefer.insert(local, reffed.as_ref());
-                let base_local = rev_lookup.un_derefer.deref_chain(local).first().unwrap().local;
-                rev_lookup.locals[local] = rev_lookup.locals[base_local];
+                    rev_lookup.un_derefer.insert(local, reffed.as_ref());
+                    let base_local =
+                        rev_lookup.un_derefer.deref_chain(local).first().unwrap().local;
+                    rev_lookup.locals[local] = rev_lookup.locals[base_local];
+                }
             }
             StatementKind::Assign(box (place, rval)) => {
                 self.create_move_path(*place);
@@ -375,13 +376,9 @@ impl<'a, 'tcx, F: Fn(Ty<'tcx>) -> bool> MoveDataBuilder<'a, 'tcx, F> {
                     self.gather_move(Place::from(*local));
                 }
             }
-            StatementKind::SetDiscriminant { .. } | StatementKind::Deinit(..) => {
-                span_bug!(
-                    stmt.source_info.span,
-                    "SetDiscriminant/Deinit should not exist during borrowck"
-                );
-            }
-            StatementKind::Retag { .. }
+            StatementKind::SetDiscriminant { .. }
+            | StatementKind::Deinit(..)
+            | StatementKind::Retag { .. }
             | StatementKind::AscribeUserType(..)
             | StatementKind::PlaceMention(..)
             | StatementKind::Coverage(..)

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -3,6 +3,9 @@ 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::MaybeUninitializedPlaces;
+use rustc_mir_dataflow::move_paths::{HasMoveData, MoveData};
+use rustc_mir_dataflow::{Analysis, ResultsCursor};
 use tracing::{debug, instrument};
 
 use crate::ssa::SsaLocals;
@@ -16,7 +19,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,25 +37,51 @@ 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 maybe-uninitialized, we'll need to remove it's storage statements.
+                head_storage_to_check.insert(head);
             }
         }
 
         let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h);
 
+        let storage_to_remove = DenseBitSet::new_empty(fully_moved.domain_size());
+
         Replacer {
             tcx,
             copy_classes: ssa.copy_classes(),
             fully_moved,
             borrowed_locals: ssa.borrowed_locals(),
-            storage_to_remove,
         }
         .visit_body_preserves_cfg(body);
 
+        debug!(?head_storage_to_check);
+
         if any_replacement {
+            // Debug builds have no use for the storage statements, so avoid extra work.
+            let storage_to_remove = if tcx.sess.emit_lifetime_markers() {
+                let move_data = MoveData::gather_moves(body, tcx, |_| true);
+
+                let maybe_uninit = MaybeUninitializedPlaces::new(tcx, body, &move_data)
+                    .iterate_to_fixpoint(tcx, body, Some("mir_opt::copy_prop"))
+                    .into_results_cursor(body);
+
+                let mut storage_checker =
+                    StorageChecker { maybe_uninit, head_storage_to_check, storage_to_remove };
+
+                storage_checker.visit_body(body);
+
+                storage_checker.storage_to_remove
+            } else {
+                // Conservatively remove all storage statements for the head locals.
+                head_storage_to_check
+            };
+            StorageRemover { tcx, storage_to_remove }.visit_body_preserves_cfg(body);
+
             crate::simplify::remove_unused_definitions(body);
         }
     }
@@ -101,7 +130,6 @@ fn fully_moved_locals(ssa: &SsaLocals, body: &Body<'_>) -> DenseBitSet<Local> {
 struct Replacer<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     fully_moved: DenseBitSet<Local>,
-    storage_to_remove: DenseBitSet<Local>,
     borrowed_locals: &'a DenseBitSet<Local>,
     copy_classes: &'a IndexSlice<Local, Local>,
 }
@@ -119,6 +147,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) => {}
@@ -152,14 +181,6 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
     }
 
     fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) {
-        // When removing storage statements, we need to remove both (#107511).
-        if let StatementKind::StorageLive(l) | StatementKind::StorageDead(l) = stmt.kind
-            && self.storage_to_remove.contains(l)
-        {
-            stmt.make_nop();
-            return;
-        }
-
         self.super_statement(stmt, loc);
 
         // Do not leave tautological assignments around.
@@ -172,3 +193,72 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
         }
     }
 }
+
+struct StorageChecker<'a, 'tcx> {
+    storage_to_remove: DenseBitSet<Local>,
+    head_storage_to_check: DenseBitSet<Local>,
+    maybe_uninit: ResultsCursor<'a, 'tcx, MaybeUninitializedPlaces<'a, 'tcx>>,
+}
+
+impl<'a, 'tcx> Visitor<'tcx> for StorageChecker<'a, 'tcx> {
+    fn visit_local(&mut self, local: Local, context: PlaceContext, loc: Location) {
+        // We don't need to check storage statements and statements for which the local doesn't need to be initialized.
+        match context {
+            PlaceContext::MutatingUse(
+                MutatingUseContext::Store
+                | MutatingUseContext::Call
+                | MutatingUseContext::AsmOutput,
+            )
+            | PlaceContext::NonUse(_) => {
+                return;
+            }
+            _ => {}
+        };
+
+        // 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(local) {
+            if let Some(move_idx) =
+                self.maybe_uninit.analysis().move_data().rev_lookup.find_local(local)
+            {
+                self.maybe_uninit.seek_before_primary_effect(loc);
+
+                if self.maybe_uninit.get().contains(move_idx) {
+                    debug!(
+                        ?loc,
+                        ?context,
+                        ?local,
+                        ?move_idx,
+                        "found a head at a location in which it is maybe uninit, marking head for storage statement removal"
+                    );
+                    self.storage_to_remove.insert(local);
+
+                    // 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(local);
+                }
+            }
+        }
+    }
+}
+
+struct StorageRemover<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    storage_to_remove: DenseBitSet<Local>,
+}
+
+impl<'tcx> MutVisitor<'tcx> for StorageRemover<'tcx> {
+    fn tcx(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+
+    fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) {
+        // When removing storage statements, we need to remove both (#107511).
+        if let StatementKind::StorageLive(l) | StatementKind::StorageDead(l) = stmt.kind
+            && self.storage_to_remove.contains(l)
+        {
+            stmt.make_nop();
+            return;
+        }
+
+        self.super_statement(stmt, loc);
+    }
+}

tests/mir-opt/copy-prop/copy_prop_storage_dead_twice.dead_twice.CopyProp.panic-abort.diff

@@ -0,0 +1,28 @@
+- // MIR for `dead_twice` before CopyProp
++ // MIR for `dead_twice` after CopyProp
+
+  fn dead_twice(_1: T) -> T {
+      let mut _0: T;
+      let mut _2: T;
+      let mut _3: T;
+      let mut _4: T;
+
+      bb0: {
+-         StorageLive(_2);
+          _2 = opaque::<T>(move _1) -> [return: bb1, unwind unreachable];
+      }
+
+      bb1: {
+-         _4 = move _2;
+-         StorageDead(_2);
+-         StorageLive(_2);
+-         _0 = opaque::<T>(move _4) -> [return: bb2, unwind unreachable];
++         _0 = opaque::<T>(move _2) -> [return: bb2, unwind unreachable];
+      }
+
+      bb2: {
+-         StorageDead(_2);
+          return;
+      }
+  }
+

tests/mir-opt/copy-prop/copy_prop_storage_dead_twice.dead_twice.CopyProp.panic-unwind.diff

@@ -0,0 +1,28 @@
+- // MIR for `dead_twice` before CopyProp
++ // MIR for `dead_twice` after CopyProp
+
+  fn dead_twice(_1: T) -> T {
+      let mut _0: T;
+      let mut _2: T;
+      let mut _3: T;
+      let mut _4: T;
+
+      bb0: {
+-         StorageLive(_2);
+          _2 = opaque::<T>(move _1) -> [return: bb1, unwind unreachable];
+      }
+
+      bb1: {
+-         _4 = move _2;
+-         StorageDead(_2);
+-         StorageLive(_2);
+-         _0 = opaque::<T>(move _4) -> [return: bb2, unwind unreachable];
++         _0 = opaque::<T>(move _2) -> [return: bb2, unwind unreachable];
+      }
+
+      bb2: {
+-         StorageDead(_2);
+          return;
+      }
+  }
+

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

@@ -0,0 +1,38 @@
+// skip-filecheck
+// EMIT_MIR_FOR_EACH_PANIC_STRATEGY
+//@ test-mir-pass: CopyProp
+
+#![feature(custom_mir, core_intrinsics)]
+
+// Check that we remove the storage statements if the head
+// becomes uninitialized before it is used again.
+
+use std::intrinsics::mir::*;
+
+// EMIT_MIR copy_prop_storage_dead_twice.dead_twice.CopyProp.diff
+#[custom_mir(dialect = "runtime")]
+pub fn dead_twice<T: Copy>(_1: T) -> T {
+    mir! {
+        let _2: T;
+        let _3: T;
+        {
+            StorageLive(_2);
+            Call(_2 = opaque(Move(_1)), ReturnTo(bb1), UnwindUnreachable())
+        }
+        bb1 = {
+            let _3 = Move(_2);
+            StorageDead(_2);
+            StorageLive(_2);
+            Call(RET = opaque(Move(_3)), ReturnTo(bb2), UnwindUnreachable())
+        }
+        bb2 = {
+            StorageDead(_2);
+            Return()
+        }
+    }
+}
+
+#[inline(never)]
+fn opaque<T>(a: T) -> T {
+    a
+}

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;
     }
 }