Dest prop: Support removing writes when this unblocks optimizations

Author: JakobDegen

compiler/rustc_mir_transform/src/dest_prop.rs

@@ -217,11 +217,12 @@ 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();
 
-            for (src, candidates) in candidates.c.iter() {
-                if merged_locals.contains(*src) {
+            for (src, candidates) in candidates.c.drain(..) {
+                if merged_locals.contains(src) {
                     continue;
                 }
-                let Some(dest) = candidates.iter().find(|dest| !merged_locals.contains(**dest))
+                let Some(dest) =
+                    candidates.into_iter().find(|(dest, _)| !merged_locals.contains(*dest))
                 else {
                     continue;
                 };
@@ -231,14 +232,15 @@ impl<'tcx> MirPass<'tcx> for DestinationPropagation {
                     break;
                 }
 
-                // Replace `src` by `dest` everywhere.
-                merges.insert(*src, *dest);
-                merged_locals.insert(*src);
-                merged_locals.insert(*dest);
-
                 // Update liveness information based on the merge we just performed.
                 // Every location where `src` was live, `dest` will be live.
-                live.union_rows(*src, *dest);
+                live.union_rows(src, dest.0);
+
+                // Replace `src` by `dest` everywhere.
+                merged_locals.insert(src);
+                merged_locals.insert(dest.0);
+                merges.insert(src, dest.clone());
+                merges.insert(dest.0, dest);
             }
             trace!(merging = ?merges);
 
@@ -260,7 +262,7 @@ impl<'tcx> MirPass<'tcx> for DestinationPropagation {
 /// frequently. Everything with a `&'alloc` lifetime points into here.
 #[derive(Default)]
 struct Allocations {
-    candidates: FxIndexMap<Local, Vec<Local>>,
+    candidates: FxIndexMap<Local, Vec<(Local, Vec<Location>)>>,
     candidates_reverse: FxIndexMap<Local, Vec<Local>>,
     write_info: WriteInfo,
     // PERF: Do this for `MaybeLiveLocals` allocations too.
@@ -282,7 +284,7 @@ struct Candidates<'alloc> {
     ///
     /// We will still report that we would like to merge `_1` and `_2` in an attempt to allow us to
     /// remove that assignment.
-    c: &'alloc mut FxIndexMap<Local, Vec<Local>>,
+    c: &'alloc mut FxIndexMap<Local, Vec<(Local, Vec<Location>)>>,
     /// A reverse index of the `c` set; if the `c` set contains `a => Place { local: b, proj }`,
     /// then this contains `b => a`.
     // PERF: Possibly these should be `SmallVec`s?
@@ -297,7 +299,7 @@ struct Candidates<'alloc> {
 fn apply_merges<'tcx>(
     body: &mut Body<'tcx>,
     tcx: TyCtxt<'tcx>,
-    merges: &FxIndexMap<Local, Local>,
+    merges: &FxIndexMap<Local, (Local, Vec<Location>)>,
     merged_locals: &BitSet<Local>,
 ) {
     let mut merger = Merger { tcx, merges, merged_locals };
@@ -306,18 +308,27 @@ fn apply_merges<'tcx>(
 
 struct Merger<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
-    merges: &'a FxIndexMap<Local, Local>,
+    merges: &'a FxIndexMap<Local, (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 {
+            return false;
+        };
+        to_remove.contains(&location)
+    }
+}
+
 impl<'a, 'tcx> MutVisitor<'tcx> for Merger<'a, 'tcx> {
     fn tcx(&self) -> TyCtxt<'tcx> {
         self.tcx
     }
 
     fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _location: Location) {
         if let Some(dest) = self.merges.get(local) {
-            *local = *dest;
+            *local = dest.0;
         }
     }
 
@@ -348,10 +359,27 @@ impl<'a, 'tcx> MutVisitor<'tcx> for Merger<'a, 'tcx> {
                     _ => {}
                 }
             }
+            StatementKind::Deinit(place) => {
+                if self.should_remove_write_at(place.local, location) {
+                    statement.make_nop();
+                }
+            }
 
             _ => {}
         }
     }
+
+    fn visit_operand(&mut self, op: &mut Operand<'tcx>, location: Location) {
+        self.super_operand(op, location);
+        match op {
+            Operand::Move(place) => {
+                if self.should_remove_write_at(place.local, location) {
+                    *op = Operand::Copy(*place);
+                }
+            }
+            _ => (),
+        }
+    }
 }
 
 //////////////////////////////////////////////////////////
@@ -373,30 +401,35 @@ struct FilterInformation<'a, 'body, 'alloc, 'tcx> {
 // through these methods, and not directly.
 impl<'alloc> Candidates<'alloc> {
     /// Just `Vec::retain`, but the condition is inverted and we add debugging output
-    fn vec_filter_candidates(
+    fn vec_modify_candidates(
         src: Local,
-        v: &mut Vec<Local>,
-        mut f: impl FnMut(Local) -> CandidateFilter,
+        v: &mut Vec<(Local, Vec<Location>)>,
+        mut f: impl FnMut(Local) -> CandidateModification,
         at: Location,
     ) {
-        v.retain(|dest| {
-            let remove = f(*dest);
-            if remove == CandidateFilter::Remove {
+        v.retain_mut(|(dest, remove_writes)| match f(*dest) {
+            CandidateModification::Remove => {
                 trace!("eliminating {:?} => {:?} due to conflict at {:?}", src, dest, at);
+                false
+            }
+            CandidateModification::RemoveWrite => {
+                trace!("marking write for {:?} => {:?} as needing removing at {:?}", src, dest, at);
+                remove_writes.push(at);
+                true
             }
-            remove == CandidateFilter::Keep
+            CandidateModification::Keep => true,
         });
     }
 
     /// `vec_filter_candidates` but for an `Entry`
     fn entry_filter_candidates(
-        mut entry: IndexOccupiedEntry<'_, Local, Vec<Local>>,
+        mut entry: IndexOccupiedEntry<'_, Local, Vec<(Local, Vec<Location>)>>,
         p: Local,
-        f: impl FnMut(Local) -> CandidateFilter,
+        f: impl FnMut(Local) -> CandidateModification,
         at: Location,
     ) {
         let candidates = entry.get_mut();
-        Self::vec_filter_candidates(p, candidates, f, at);
+        Self::vec_modify_candidates(p, candidates, f, at);
         if candidates.len() == 0 {
             entry.remove();
         }
@@ -406,7 +439,7 @@ impl<'alloc> Candidates<'alloc> {
     fn filter_candidates_by(
         &mut self,
         p: Local,
-        mut f: impl FnMut(Local) -> CandidateFilter,
+        mut f: impl FnMut(Local) -> CandidateModification,
         at: Location,
     ) {
         // Cover the cases where `p` appears as a `src`
@@ -420,7 +453,8 @@ impl<'alloc> Candidates<'alloc> {
         // We use `retain` here to remove the elements from the reverse set if we've removed the
         // matching candidate in the forward set.
         srcs.retain(|src| {
-            if f(*src) == CandidateFilter::Keep {
+            let modification = f(*src);
+            if modification == CandidateModification::Keep {
                 return true;
             }
             let IndexEntry::Occupied(entry) = self.c.entry(*src) else {
@@ -430,18 +464,20 @@ impl<'alloc> Candidates<'alloc> {
                 entry,
                 *src,
                 |dest| {
-                    if dest == p { CandidateFilter::Remove } else { CandidateFilter::Keep }
+                    if dest == p { modification } else { CandidateModification::Keep }
                 },
                 at,
             );
-            false
+            // Remove the src from the reverse set if we removed the candidate pair
+            modification == CandidateModification::RemoveWrite
         });
     }
 }
 
 #[derive(Copy, Clone, PartialEq, Eq)]
-enum CandidateFilter {
+enum CandidateModification {
     Keep,
+    RemoveWrite,
     Remove,
 }
 
@@ -500,32 +536,37 @@ impl<'a, 'body, 'alloc, 'tcx> FilterInformation<'a, 'body, 'alloc, 'tcx> {
 
     fn apply_conflicts(&mut self) {
         let writes = &self.write_info.writes;
-        for p in writes {
+        for &(p, is_removable) in writes {
+            let modification = if is_removable {
+                CandidateModification::RemoveWrite
+            } else {
+                CandidateModification::Remove
+            };
             let other_skip = self.write_info.skip_pair.and_then(|(a, b)| {
-                if a == *p {
+                if a == p {
                     Some(b)
-                } else if b == *p {
+                } else if b == p {
                     Some(a)
                 } else {
                     None
                 }
             });
             let at = self.points.point_from_location(self.at);
             self.candidates.filter_candidates_by(
-                *p,
+                p,
                 |q| {
                     if Some(q) == other_skip {
-                        return CandidateFilter::Keep;
+                        return CandidateModification::Keep;
                     }
                     // It is possible that a local may be live for less than the
                     // duration of a statement This happens in the case of function
                     // calls or inline asm. Because of this, we also mark locals as
                     // conflicting when both of them are written to in the same
                     // statement.
-                    if self.live.contains(q, at) || writes.contains(&q) {
-                        CandidateFilter::Remove
+                    if self.live.contains(q, at) || writes.iter().any(|&(x, _)| x == q) {
+                        modification
                     } else {
-                        CandidateFilter::Keep
+                        CandidateModification::Keep
                     }
                 },
                 self.at,
@@ -537,7 +578,9 @@ impl<'a, 'body, 'alloc, 'tcx> FilterInformation<'a, 'body, 'alloc, 'tcx> {
 /// Describes where a statement/terminator writes to
 #[derive(Default, Debug)]
 struct WriteInfo {
-    writes: Vec<Local>,
+    /// Which locals are written to. The `bool` is true if the write is "removable," ie if it comes
+    /// from a `Operand::Move` or `Deinit`.
+    writes: Vec<(Local, bool)>,
     /// If this pair of locals is a candidate pair, completely skip processing it during this
     /// statement. All other candidates are unaffected.
     skip_pair: Option<(Local, Local)>,
@@ -581,10 +624,11 @@ impl WriteInfo {
                     | Rvalue::CopyForDeref(_) => (),
                 }
             }
+            StatementKind::Deinit(p) => {
+                self.writes.push((p.local, true));
+            }
             // Retags are technically also reads, but reporting them as a write suffices
-            StatementKind::SetDiscriminant { place, .. }
-            | StatementKind::Deinit(place)
-            | StatementKind::Retag(_, place) => {
+            StatementKind::SetDiscriminant { place, .. } | StatementKind::Retag(_, place) => {
                 self.add_place(**place);
             }
             StatementKind::Intrinsic(_)
@@ -669,16 +713,12 @@ impl WriteInfo {
     }
 
     fn add_place(&mut self, place: Place<'_>) {
-        self.writes.push(place.local);
+        self.writes.push((place.local, false));
     }
 
     fn add_operand<'tcx>(&mut self, op: &Operand<'tcx>) {
         match op {
-            // FIXME(JakobDegen): In a previous version, the `Move` case was incorrectly treated as
-            // being a read only. This was unsound, however we cannot add a regression test because
-            // it is not possible to set this off with current MIR. Once we have that ability, a
-            // regression test should be added.
-            Operand::Move(p) => self.add_place(*p),
+            Operand::Move(p) => self.writes.push((p.local, true)),
             Operand::Copy(_) | Operand::Constant(_) => (),
         }
     }
@@ -733,7 +773,7 @@ fn places_to_candidate_pair<'tcx>(
 fn find_candidates<'alloc, 'tcx>(
     body: &Body<'tcx>,
     borrowed: &BitSet<Local>,
-    candidates: &'alloc mut FxIndexMap<Local, Vec<Local>>,
+    candidates: &'alloc mut FxIndexMap<Local, Vec<(Local, Vec<Location>)>>,
     candidates_reverse: &'alloc mut FxIndexMap<Local, Vec<Local>>,
 ) -> Candidates<'alloc> {
     candidates.clear();
@@ -747,16 +787,16 @@ fn find_candidates<'alloc, 'tcx>(
     }
     // Generate the reverse map
     for (src, cands) in candidates.iter() {
-        for dest in cands.iter().copied() {
-            candidates_reverse.entry(dest).or_default().push(*src);
+        for (dest, _) in cands.iter() {
+            candidates_reverse.entry(*dest).or_default().push(*src);
         }
     }
     Candidates { c: candidates, reverse: candidates_reverse }
 }
 
 struct FindAssignments<'a, 'alloc, 'tcx> {
     body: &'a Body<'tcx>,
-    candidates: &'alloc mut FxIndexMap<Local, Vec<Local>>,
+    candidates: &'alloc mut FxIndexMap<Local, Vec<(Local, Vec<Location>)>>,
     borrowed: &'a BitSet<Local>,
 }
 
@@ -793,7 +833,7 @@ impl<'tcx> Visitor<'tcx> for FindAssignments<'_, '_, 'tcx> {
             }
 
             // We may insert duplicates here, but that's fine
-            self.candidates.entry(src).or_default().push(dest);
+            self.candidates.entry(src).or_default().push((dest, Vec::new()));
         }
     }
 }

tests/mir-opt/dest-prop/move.move_simple.DestinationPropagation.panic-abort.diff

@@ -0,0 +1,35 @@
+- // MIR for `move_simple` before DestinationPropagation
++ // MIR for `move_simple` after DestinationPropagation
+  
+  fn move_simple(_1: i32) -> () {
+      debug x => _1;
+      let mut _0: ();
+      let _2: ();
+      let mut _3: i32;
+      let mut _4: i32;
+  
+      bb0: {
+          StorageLive(_2);
+-         StorageLive(_3);
+-         _3 = _1;
+-         StorageLive(_4);
+-         _4 = _1;
+-         _2 = use_both(move _3, move _4) -> [return: bb1, unwind unreachable];
++         nop;
++         nop;
++         nop;
++         nop;
++         _2 = use_both(_1, _1) -> [return: bb1, unwind unreachable];
+      }
+  
+      bb1: {
+-         StorageDead(_4);
+-         StorageDead(_3);
++         nop;
++         nop;
+          StorageDead(_2);
+          _0 = const ();
+          return;
+      }
+  }
+