Unify a source with all possible destinations.

cjgillot · cjgillot · commit 99f6bcf38065 · 2025-09-07T16:45:00.000Z
diff --git a/compiler/rustc_mir_transform/src/dest_prop.rs b/compiler/rustc_mir_transform/src/dest_prop.rs
@@ -190,29 +190,51 @@ impl<'tcx> crate::MirPass<'tcx> for DestinationPropagation {
         for (src, candidates) in candidates.c.into_iter() {
             trace!(?src, ?candidates);
 
-            let Some(src) = relevant.find(src) else { continue };
-            let Some(src_live_ranges) = &live.row(src) else { continue };
-            trace!(?src, ?src_live_ranges);
-
-            let dst = candidates.into_iter().find_map(|dst| {
-                let dst = relevant.find(dst)?;
-                let dst_live_ranges = &live.row(dst)?;
+            for dst in candidates {
+                // We call `relevant.find(src)` inside the loop, as a previous iteration may have
+                // renamed `src` to one of the locals in `dst`.
+                let Some(mut src) = relevant.find(src) else { continue };
+                let Some(src_live_ranges) = live.row(src) else { continue };
+                trace!(?src, ?src_live_ranges);
+
+                let Some(mut dst) = relevant.find(dst) else { continue };
+                let Some(dst_live_ranges) = live.row(dst) else { continue };
                 trace!(?dst, ?dst_live_ranges);
 
-                let disjoint = src_live_ranges.disjoint(dst_live_ranges);
-                disjoint.then_some(dst)
-            });
-            let Some(dst) = dst else { continue };
+                if src_live_ranges.disjoint(dst_live_ranges) {
+                    // We want to replace `src` by `dst`.
+                    let mut orig_src = relevant.original[src];
+                    let mut orig_dst = relevant.original[dst];
+
+                    // The return place and function arguments are required and cannot be renamed.
+                    // This check cannot be made during candidate collection, as we may want to
+                    // unify the same non-required local with several required locals.
+                    match (is_local_required(orig_src, body), is_local_required(orig_dst, body)) {
+                        // Renaming `src` is ok.
+                        (false, _) => {}
+                        // Renaming `src` is wrong, but renaming `dst` is ok.
+                        (true, false) => {
+                            std::mem::swap(&mut src, &mut dst);
+                            std::mem::swap(&mut orig_src, &mut orig_dst);
+                        }
+                        // Neither local can be renamed, so skip this case.
+                        (true, true) => continue,
+                    }
 
-            merged_locals.insert(relevant.original[src]);
-            merged_locals.insert(relevant.original[dst]);
+                    trace!(?src, ?dst, "merge");
+                    merged_locals.insert(orig_src);
+                    merged_locals.insert(orig_dst);
 
-            relevant.union(src, dst);
-            live.union_rows(src, dst);
+                    // Replace `src` by `dst`.
+                    relevant.union(src, dst);
+                    live.union_rows(/* read */ src, /* write */ dst);
+                }
+            }
         }
         trace!(?merged_locals);
 
         relevant.make_idempotent();
+        trace!(?relevant.renames);
 
         if merged_locals.is_empty() {
             return;
@@ -402,41 +424,6 @@ impl Candidates {
     }
 }
 
-/// If the pair of places is being considered for merging, returns the candidate which would be
-/// merged in order to accomplish this.
-///
-/// The contract here is in one direction - there is a guarantee that merging the locals that are
-/// outputted by this function would result in an assignment between the inputs becoming a
-/// self-assignment. However, there is no guarantee that the returned pair is actually suitable for
-/// merging - candidate collection must still check this independently.
-///
-/// This output is unique for each unordered pair of input places.
-fn places_to_candidate_pair<'tcx>(
-    a: Place<'tcx>,
-    b: Place<'tcx>,
-    body: &Body<'tcx>,
-) -> Option<(Local, Local)> {
-    let (mut a, mut b) = if a.projection.len() == 0 && b.projection.len() == 0 {
-        (a.local, b.local)
-    } else {
-        return None;
-    };
-
-    // By sorting, we make sure we're input order independent
-    if a > b {
-        std::mem::swap(&mut a, &mut b);
-    }
-
-    // We could now return `(a, b)`, but then we miss some candidates in the case where `a` can't be
-    // used as a `src`.
-    if is_local_required(a, body) {
-        std::mem::swap(&mut a, &mut b);
-    }
-    // We could check `is_local_required` again here, but there's no need - after all, we make no
-    // promise that the candidate pair is actually valid
-    Some((a, b))
-}
-
 struct FindAssignments<'a, 'tcx> {
     body: &'a Body<'tcx>,
     candidates: FxIndexMap<Local, Vec<Local>>,
@@ -449,11 +436,9 @@ impl<'tcx> Visitor<'tcx> for FindAssignments<'_, 'tcx> {
             lhs,
             Rvalue::CopyForDeref(rhs) | Rvalue::Use(Operand::Copy(rhs) | Operand::Move(rhs)),
         )) = &statement.kind
+            && let Some(src) = lhs.as_local()
+            && let Some(dest) = rhs.as_local()
         {
-            let Some((src, dest)) = places_to_candidate_pair(*lhs, *rhs, self.body) else {
-                return;
-            };
-
             // As described at the top of the file, we do not go near things that have
             // their address taken.
             if self.borrowed.contains(src) || self.borrowed.contains(dest) {
@@ -470,11 +455,6 @@ impl<'tcx> Visitor<'tcx> for FindAssignments<'_, 'tcx> {
                 return;
             }
 
-            // Also, we need to make sure that MIR actually allows the `src` to be removed
-            if is_local_required(src, self.body) {
-                return;
-            }
-
             // We may insert duplicates here, but that's fine
             self.candidates.entry(src).or_default().push(dest);
         }
diff --git a/tests/mir-opt/dest-prop/cycle.main.DestinationPropagation.panic-abort.diff b/tests/mir-opt/dest-prop/cycle.main.DestinationPropagation.panic-abort.diff
@@ -8,25 +8,23 @@
       let _5: ();
       let mut _6: i32;
       scope 1 {
--         debug x => _1;
-+         debug x => _4;
+          debug x => _1;
           let _2: i32;
           scope 2 {
 -             debug y => _2;
-+             debug y => _4;
++             debug y => _1;
               let _3: i32;
               scope 3 {
 -                 debug z => _3;
-+                 debug z => _4;
++                 debug z => _1;
               }
           }
       }
   
       bb0: {
 -         StorageLive(_1);
--         _1 = val() -> [return: bb1, unwind unreachable];
 +         nop;
-+         _4 = val() -> [return: bb1, unwind unreachable];
+          _1 = val() -> [return: bb1, unwind unreachable];
       }
   
       bb1: {
@@ -47,14 +45,17 @@
 +         nop;
 +         nop;
           StorageLive(_5);
-          StorageLive(_6);
+-         StorageLive(_6);
 -         _6 = copy _1;
-+         _6 = copy _4;
-          _5 = std::mem::drop::<i32>(move _6) -> [return: bb2, unwind unreachable];
+-         _5 = std::mem::drop::<i32>(move _6) -> [return: bb2, unwind unreachable];
++         nop;
++         nop;
++         _5 = std::mem::drop::<i32>(move _1) -> [return: bb2, unwind unreachable];
       }
   
       bb2: {
-          StorageDead(_6);
+-         StorageDead(_6);
++         nop;
           StorageDead(_5);
           _0 = const ();
 -         StorageDead(_3);
diff --git a/tests/mir-opt/dest-prop/cycle.main.DestinationPropagation.panic-unwind.diff b/tests/mir-opt/dest-prop/cycle.main.DestinationPropagation.panic-unwind.diff
@@ -8,25 +8,23 @@
       let _5: ();
       let mut _6: i32;
       scope 1 {
--         debug x => _1;
-+         debug x => _4;
+          debug x => _1;
           let _2: i32;
           scope 2 {
 -             debug y => _2;
-+             debug y => _4;
++             debug y => _1;
               let _3: i32;
               scope 3 {
 -                 debug z => _3;
-+                 debug z => _4;
++                 debug z => _1;
               }
           }
       }
   
       bb0: {
 -         StorageLive(_1);
--         _1 = val() -> [return: bb1, unwind continue];
 +         nop;
-+         _4 = val() -> [return: bb1, unwind continue];
+          _1 = val() -> [return: bb1, unwind continue];
       }
   
       bb1: {
@@ -47,14 +45,17 @@
 +         nop;
 +         nop;
           StorageLive(_5);
-          StorageLive(_6);
+-         StorageLive(_6);
 -         _6 = copy _1;
-+         _6 = copy _4;
-          _5 = std::mem::drop::<i32>(move _6) -> [return: bb2, unwind continue];
+-         _5 = std::mem::drop::<i32>(move _6) -> [return: bb2, unwind continue];
++         nop;
++         nop;
++         _5 = std::mem::drop::<i32>(move _1) -> [return: bb2, unwind continue];
       }
   
       bb2: {
-          StorageDead(_6);
+-         StorageDead(_6);
++         nop;
           StorageDead(_5);
           _0 = const ();
 -         StorageDead(_3);
diff --git a/tests/mir-opt/dest-prop/dead_stores_79191.f.DestinationPropagation.after.panic-abort.mir b/tests/mir-opt/dest-prop/dead_stores_79191.f.DestinationPropagation.after.panic-abort.mir
@@ -7,16 +7,16 @@ fn f(_1: usize) -> usize {
     let mut _3: usize;
     let mut _4: usize;
     scope 1 {
-        debug b => _3;
+        debug b => _2;
     }
 
     bb0: {
         nop;
-        _3 = copy _1;
+        _2 = copy _1;
         _1 = const 5_usize;
         nop;
         nop;
-        _1 = move _3;
+        _1 = move _2;
         nop;
         nop;
         nop;
diff --git a/tests/mir-opt/dest-prop/dead_stores_79191.f.DestinationPropagation.after.panic-unwind.mir b/tests/mir-opt/dest-prop/dead_stores_79191.f.DestinationPropagation.after.panic-unwind.mir
@@ -7,16 +7,16 @@ fn f(_1: usize) -> usize {
     let mut _3: usize;
     let mut _4: usize;
     scope 1 {
-        debug b => _3;
+        debug b => _2;
     }
 
     bb0: {
         nop;
-        _3 = copy _1;
+        _2 = copy _1;
         _1 = const 5_usize;
         nop;
         nop;
-        _1 = move _3;
+        _1 = move _2;
         nop;
         nop;
         nop;
diff --git a/tests/mir-opt/dest-prop/dead_stores_better.f.DestinationPropagation.after.panic-abort.mir b/tests/mir-opt/dest-prop/dead_stores_better.f.DestinationPropagation.after.panic-abort.mir
@@ -7,16 +7,16 @@ fn f(_1: usize) -> usize {
     let mut _3: usize;
     let mut _4: usize;
     scope 1 {
-        debug b => _3;
+        debug b => _2;
     }
 
     bb0: {
         nop;
-        _3 = copy _1;
+        _2 = copy _1;
         _1 = const 5_usize;
         nop;
         nop;
-        _1 = move _3;
+        _1 = move _2;
         nop;
         nop;
         nop;
diff --git a/tests/mir-opt/dest-prop/dead_stores_better.f.DestinationPropagation.after.panic-unwind.mir b/tests/mir-opt/dest-prop/dead_stores_better.f.DestinationPropagation.after.panic-unwind.mir
@@ -7,16 +7,16 @@ fn f(_1: usize) -> usize {
     let mut _3: usize;
     let mut _4: usize;
     scope 1 {
-        debug b => _3;
+        debug b => _2;
     }
 
     bb0: {
         nop;
-        _3 = copy _1;
+        _2 = copy _1;
         _1 = const 5_usize;
         nop;
         nop;
-        _1 = move _3;
+        _1 = move _2;
         nop;
         nop;
         nop;
diff --git a/tests/mir-opt/pre-codegen/clone_as_copy.enum_clone_as_copy.PreCodegen.after.mir b/tests/mir-opt/pre-codegen/clone_as_copy.enum_clone_as_copy.PreCodegen.after.mir
@@ -6,8 +6,8 @@ fn enum_clone_as_copy(_1: &Enum1) -> Enum1 {
     scope 1 (inlined <Enum1 as Clone>::clone) {
         debug self => _1;
         let mut _2: isize;
-        let mut _3: &AllCopy;
-        let mut _4: &NestCopy;
+        let _3: &AllCopy;
+        let _4: &NestCopy;
         scope 2 {
             debug __self_0 => _3;
             scope 6 (inlined <AllCopy as Clone>::clone) {
