Dest prop: Separate writes to remove into separate data structure in merger

Author: JakobDegen

compiler/rustc_mir_transform/src/dest_prop.rs

@@ -132,7 +132,7 @@
 //! [attempt 3]: https://github.com/rust-lang/rust/pull/72632
 
 use crate::MirPass;
-use rustc_data_structures::fx::{FxIndexMap, IndexEntry, IndexOccupiedEntry};
+use rustc_data_structures::fx::{FxHashMap, FxIndexMap, IndexEntry, IndexOccupiedEntry};
 use rustc_index::bit_set::BitSet;
 use rustc_index::interval::SparseIntervalMatrix;
 use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
@@ -216,6 +216,7 @@ impl<'tcx> MirPass<'tcx> for DestinationPropagation {
 
             // This is the set of merges we will apply this round. It is a subset of the candidates.
             let mut merges = FxIndexMap::default();
+            let mut remove_writes = FxHashMap::default();
 
             for (src, candidates) in candidates.c.drain(..) {
                 if merged_locals.contains(src) {
@@ -239,8 +240,10 @@ impl<'tcx> MirPass<'tcx> for DestinationPropagation {
                 // Replace `src` by `dest` everywhere.
                 merged_locals.insert(src);
                 merged_locals.insert(dest.0);
-                merges.insert(src, dest.clone());
-                merges.insert(dest.0, dest);
+                merges.insert(src, dest.0);
+                if !dest.1.is_empty() {
+                    remove_writes.insert(dest.0, dest.1);
+                }
             }
             trace!(merging = ?merges);
 
@@ -249,7 +252,7 @@ impl<'tcx> MirPass<'tcx> for DestinationPropagation {
             }
             round_count += 1;
 
-            apply_merges(body, tcx, &merges, &merged_locals);
+            apply_merges(body, tcx, &merges, &remove_writes, &merged_locals);
         }
 
         trace!(round_count);
@@ -299,22 +302,24 @@ struct Candidates<'alloc> {
 fn apply_merges<'tcx>(
     body: &mut Body<'tcx>,
     tcx: TyCtxt<'tcx>,
-    merges: &FxIndexMap<Local, (Local, Vec<Location>)>,
+    merges: &FxIndexMap<Local, Local>,
+    remove_writes: &FxHashMap<Local, Vec<Location>>,
     merged_locals: &BitSet<Local>,
 ) {
-    let mut merger = Merger { tcx, merges, merged_locals };
+    let mut merger = Merger { tcx, merges, remove_writes, merged_locals };
     merger.visit_body_preserves_cfg(body);
 }
 
 struct Merger<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
-    merges: &'a FxIndexMap<Local, (Local, Vec<Location>)>,
+    merges: &'a FxIndexMap<Local, Local>,
+    remove_writes: &'a FxHashMap<Local, Vec<Location>>,
     merged_locals: &'a BitSet<Local>,
 }
 
 impl<'a, 'tcx> Merger<'a, 'tcx> {
     fn should_remove_write_at(&self, local: Local, location: Location) -> bool {
-        let Some((_, to_remove)) = self.merges.get(&local) else {
+        let Some(to_remove) = self.remove_writes.get(&local) else {
             return false;
         };
         to_remove.contains(&location)
@@ -328,7 +333,7 @@ impl<'a, 'tcx> MutVisitor<'tcx> for Merger<'a, 'tcx> {
 
     fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _location: Location) {
         if let Some(dest) = self.merges.get(local) {
-            *local = dest.0;
+            *local = *dest;
         }
     }