Switch back to replace->uninit->storage remove

Author: ohadravid

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -49,12 +49,13 @@ impl<'tcx> crate::MirPass<'tcx> for CopyProp {
 
         let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h);
 
-        let storage_to_remove = if any_replacement {
-            // The replacer will remove tautological moves from the head to the local,
-            // so we remove them before running `MaybeUninitializedPlaces`.
-            TautologicalMoveAssignmentRemover {
+        if any_replacement {
+            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(),
             }
             .visit_body_preserves_cfg(body);
@@ -65,29 +66,18 @@ impl<'tcx> crate::MirPass<'tcx> for CopyProp {
                 .iterate_to_fixpoint(tcx, body, Some("mir_opt::copy_prop"))
                 .into_results_cursor(body);
 
-            let mut storage_checker = StorageChecker {
-                copy_classes: ssa.copy_classes(),
-                maybe_uninit,
-                head_storage_to_check,
-                storage_to_remove: DenseBitSet::new_empty(fully_moved.domain_size()),
-            };
+            debug!(?head_storage_to_check);
 
-            storage_checker.visit_body(body);
+            let storage_to_remove = {
+                let mut storage_checker =
+                    StorageChecker { maybe_uninit, head_storage_to_check, storage_to_remove };
 
-            storage_checker.storage_to_remove
-        } else {
-            // Will be empty anyway.
-            head_storage_to_check
-        };
+                storage_checker.visit_body(body);
 
-        Replacer {
-            tcx,
-            copy_classes: ssa.copy_classes(),
-            fully_moved,
-            borrowed_locals: ssa.borrowed_locals(),
-            storage_to_remove,
+                storage_checker.storage_to_remove
+            };
+            StorageRemover { tcx, storage_to_remove }.visit_body_preserves_cfg(body);
         }
-        .visit_body_preserves_cfg(body);
 
         if any_replacement {
             crate::simplify::remove_unused_definitions(body);
@@ -134,46 +124,10 @@ fn fully_moved_locals(ssa: &SsaLocals, body: &Body<'_>) -> DenseBitSet<Local> {
     fully_moved
 }
 
-// We don't want to consider tautological moves in `MaybeUninitializedPlaces`, so we remove them before running it.
-struct TautologicalMoveAssignmentRemover<'a, 'tcx> {
-    tcx: TyCtxt<'tcx>,
-    copy_classes: &'a IndexSlice<Local, Local>,
-    borrowed_locals: &'a DenseBitSet<Local>,
-}
-
-impl<'tcx> MutVisitor<'tcx> for TautologicalMoveAssignmentRemover<'_, 'tcx> {
-    fn tcx(&self) -> TyCtxt<'tcx> {
-        self.tcx
-    }
-
-    fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) {
-        self.super_statement(stmt, loc);
-
-        // Similar to `Replacer::visit_statement`, but only handle moves.
-        if let StatementKind::Assign(box (lhs, ref rhs)) = stmt.kind
-            && let Rvalue::Use(Operand::Move(rhs)) = *rhs
-        {
-            if self.borrowed_locals.contains(lhs.local) || self.borrowed_locals.contains(rhs.local)
-            {
-                return;
-            }
-
-            let new_lhs_local = self.copy_classes[lhs.local];
-            let new_rhs_local = self.copy_classes[rhs.local];
-
-            if new_lhs_local == new_rhs_local && lhs.projection == rhs.projection {
-                debug!(?loc, ?lhs, ?rhs, "removing a to-be-tautological assignment");
-                stmt.make_nop();
-            }
-        }
-    }
-}
-
 /// Utility to help performing substitution of `*pattern` by `target`.
 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>,
 }
@@ -225,14 +179,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.
@@ -241,6 +187,7 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
                 *rhs
             && lhs == rhs
         {
+            debug!(?stmt, ?loc, ?lhs, ?rhs, "removing tautological assignment");
             stmt.make_nop();
         }
     }
@@ -250,42 +197,63 @@ struct StorageChecker<'a, 'tcx> {
     storage_to_remove: DenseBitSet<Local>,
     head_storage_to_check: DenseBitSet<Local>,
     maybe_uninit: ResultsCursor<'a, 'tcx, MaybeUninitializedPlaces<'a, 'tcx>>,
-    copy_classes: &'a IndexSlice<Local, Local>,
 }
 
 impl<'a, 'tcx> Visitor<'tcx> for StorageChecker<'a, 'tcx> {
     fn visit_local(&mut self, local: Local, context: PlaceContext, loc: Location) {
-        if !context.is_use() {
-            return;
+        // We don't need to check storage statements and stores (for which the local doesn't need to be initialized).
+        match context {
+            PlaceContext::MutatingUse(MutatingUseContext::Store) | 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);
+                }
+            }
         }
+    }
+}
+
+/// Utility to help performing substitution of `*pattern` by `target`.
+struct StorageRemover<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    storage_to_remove: DenseBitSet<Local>,
+}
 
-        let head = self.copy_classes[local];
+impl<'tcx> MutVisitor<'tcx> for StorageRemover<'tcx> {
+    fn tcx(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
 
-        if local == head {
-            // Every original use of the head is known to be initialized, so we don't need to check it.
+    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;
         }
 
-        // 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)
-            && let Some(move_idx) =
-                self.maybe_uninit.analysis().move_data().rev_lookup.find_local(head)
-        {
-            self.maybe_uninit.seek_after_primary_effect(loc);
-
-            if self.maybe_uninit.get().contains(move_idx) {
-                debug!(
-                    ?loc,
-                    ?local,
-                    ?head,
-                    ?move_idx,
-                    "found use of local with 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);
-            }
-        }
+        self.super_statement(stmt, loc);
     }
 }