Implement MaybeUninitializedLocals analysis for copy_prop mir-opt to remove fewer storage statements

Author: ohadravid

compiler/rustc_mir_dataflow/src/impls/initialized.rs

@@ -558,6 +558,80 @@ impl<'tcx> Analysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
     }
 }
 
+/// A dataflow analysis that tracks locals that are maybe uninitialized.
+///
+/// This is a simpler analysis than `MaybeUninitializedPlaces`, because it does not track
+/// individual fields.
+pub struct MaybeUninitializedLocals;
+
+impl MaybeUninitializedLocals {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl<'tcx> Analysis<'tcx> for MaybeUninitializedLocals {
+    type Domain = DenseBitSet<mir::Local>;
+
+    const NAME: &'static str = "maybe_uninit_locals";
+
+    fn bottom_value(&self, body: &Body<'tcx>) -> Self::Domain {
+        // bottom = all locals are initialized.
+        DenseBitSet::new_empty(body.local_decls.len())
+    }
+
+    fn initialize_start_block(&self, body: &Body<'tcx>, state: &mut Self::Domain) {
+        // All locals start as uninitialized...
+        state.insert_all();
+        // ...except for arguments, which are definitely initialized.
+        for arg in body.args_iter() {
+            state.remove(arg);
+        }
+    }
+
+    fn apply_primary_statement_effect(
+        &mut self,
+        state: &mut Self::Domain,
+        statement: &mir::Statement<'tcx>,
+        _location: Location,
+    ) {
+        match statement.kind {
+            // An assignment makes a local initialized.
+            mir::StatementKind::Assign(box (place, _)) => {
+                if let Some(local) = place.as_local() {
+                    state.remove(local);
+                }
+            }
+            // Deinit makes the local uninitialized.
+            mir::StatementKind::Deinit(box place) => {
+                // A deinit makes a local uninitialized.
+                if let Some(local) = place.as_local() {
+                    state.insert(local);
+                }
+            }
+            // Storage{Live,Dead} makes a local uninitialized.
+            mir::StatementKind::StorageLive(local) | mir::StatementKind::StorageDead(local) => {
+                state.insert(local);
+            }
+            _ => {}
+        }
+    }
+
+    fn apply_call_return_effect(
+        &mut self,
+        state: &mut Self::Domain,
+        _block: mir::BasicBlock,
+        return_places: CallReturnPlaces<'_, 'tcx>,
+    ) {
+        // The return place of a call is initialized.
+        return_places.for_each(|place| {
+            if let Some(local) = place.as_local() {
+                state.remove(local);
+            }
+        });
+    }
+}
+
 /// There can be many more `InitIndex` than there are locals in a MIR body.
 /// We use a mixed bitset to avoid paying too high a memory footprint.
 pub type EverInitializedPlacesDomain = MixedBitSet<InitIndex>;

compiler/rustc_mir_dataflow/src/impls/mod.rs

@@ -6,7 +6,7 @@ mod storage_liveness;
 pub use self::borrowed_locals::{MaybeBorrowedLocals, borrowed_locals};
 pub use self::initialized::{
     EverInitializedPlaces, EverInitializedPlacesDomain, MaybeInitializedPlaces,
-    MaybeUninitializedPlaces, MaybeUninitializedPlacesDomain,
+    MaybeUninitializedLocals, MaybeUninitializedPlaces, MaybeUninitializedPlacesDomain,
 };
 pub use self::liveness::{
     MaybeLiveLocals, MaybeTransitiveLiveLocals, TransferFunction as LivenessTransferFunction,

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -3,6 +3,8 @@ 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::MaybeUninitializedLocals;
+use rustc_mir_dataflow::{Analysis, ResultsCursor};
 use tracing::{debug, instrument};
 
 use crate::ssa::SsaLocals;
@@ -16,7 +18,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 {
@@ -30,21 +32,56 @@ impl<'tcx> crate::MirPass<'tcx> for CopyProp {
 
         let typing_env = body.typing_env(tcx);
         let ssa = SsaLocals::new(tcx, body, typing_env);
-        debug!(borrowed_locals = ?ssa.borrowed_locals());
+        let borrowed_locals = ssa.borrowed_locals().clone();
+
+        debug!(?borrowed_locals);
         debug!(copy_classes = ?ssa.copy_classes());
 
         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 {
-                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);
+
+                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);
+                }
             }
         }
 
         let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h);
 
+        // 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() {
+            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,
+                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
+        };
+
+        debug!(?storage_to_remove);
+
         Replacer { tcx, copy_classes: ssa.copy_classes(), fully_moved, storage_to_remove }
             .visit_body_preserves_cfg(body);
 
@@ -152,3 +189,51 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
         }
     }
 }
+
+// Marks heads of copy classes that are maybe uninitialized at the location of a local
+// as needing storage statement removal.
+struct StorageChecker<'a, 'tcx> {
+    maybe_uninit: ResultsCursor<'a, 'tcx, MaybeUninitializedLocals>,
+    copy_classes: &'a IndexSlice<Local, Local>,
+    head_storage_to_check: DenseBitSet<Local>,
+    storage_to_remove: DenseBitSet<Local>,
+}
+
+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::Yield
+                | MutatingUseContext::AsmOutput,
+            )
+            | PlaceContext::NonUse(_) => {
+                return;
+            }
+            _ => {}
+        };
+
+        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);
+            }
+        }
+    }
+}

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

@@ -0,0 +1,21 @@
+- // MIR for `f` before CopyProp
++ // MIR for `f` after CopyProp
+  
+  fn f(_1: (T, T)) -> T {
+      let mut _0: T;
+      let mut _2: T;
+      let mut _3: T;
+      let mut _4: &T;
+  
+      bb0: {
+-         StorageLive(_2);
+          _2 = copy (_1.0: T);
+-         _3 = copy _2;
+-         _4 = &_3;
+-         StorageDead(_2);
++         _4 = &_2;
+          _0 = copy (*_4);
+          return;
+      }
+  }
+  

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

@@ -0,0 +1,30 @@
+// skip-filecheck
+//@ test-mir-pass: CopyProp
+
+#![feature(custom_mir, core_intrinsics, freeze)]
+
+// Check that we remove the storage statements if one of the locals is borrowed,
+// and the head isn't borrowed.
+
+use std::intrinsics::mir::*;
+use std::marker::Freeze;
+
+// EMIT_MIR copy_prop_borrowed_storage_not_removed.f.CopyProp.diff
+
+#[custom_mir(dialect = "runtime")]
+pub fn f<T: Copy + Freeze>(_1: (T, T)) -> T {
+    mir! {
+        let _2: T;
+        let _3: T;
+        let _4: &T;
+        {
+            StorageLive(_2);
+            _2 = _1.0;
+            _3 = _2;
+            _4 = &_3;
+            StorageDead(_2);
+            RET = *_4;
+            Return()
+        }
+    }
+}

tests/mir-opt/copy-prop/copy_prop_storage_twice.dead_twice.CopyProp.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_twice.live_twice.CopyProp.diff

@@ -0,0 +1,27 @@
+- // MIR for `live_twice` before CopyProp
++ // MIR for `live_twice` after CopyProp
+  
+  fn live_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;
+-         StorageLive(_2);
+-         _0 = opaque::<T>(copy _4) -> [return: bb2, unwind unreachable];
++         _0 = opaque::<T>(copy _2) -> [return: bb2, unwind unreachable];
+      }
+  
+      bb2: {
+-         StorageDead(_2);
+          return;
+      }
+  }
+  

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

@@ -0,0 +1,61 @@
+// skip-filecheck
+//@ test-mir-pass: CopyProp
+//@ compile-flags: -Zlint-mir=false
+
+#![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_twice.dead_twice.CopyProp.diff
+// EMIT_MIR copy_prop_storage_twice.live_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()
+        }
+    }
+}
+
+#[custom_mir(dialect = "runtime")]
+pub fn live_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);
+            StorageLive(_2);
+            Call(RET = opaque(_3), ReturnTo(bb2), UnwindUnreachable())
+        }
+        bb2 = {
+            StorageDead(_2);
+            Return()
+        }
+    }
+}
+
+#[inline(never)]
+fn opaque<T>(a: T) -> T {
+    a
+}