Auto merge of #106908 - cjgillot:copyprop-ssa, r=oli-obk

Implement simple CopyPropagation based on SSA analysis

This PR extracts the "copy propagation" logic from #106285.

MIR may produce chains of assignment between locals, like `_x = move? _y`.
This PR attempts to remove such chains by unifying locals.

The current implementation is a bit overzealous in turning moves into copies, and in removing storage statements.

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -0,0 +1,178 @@
+use rustc_index::bit_set::BitSet;
+use rustc_index::vec::IndexVec;
+use rustc_middle::mir::visit::*;
+use rustc_middle::mir::*;
+use rustc_middle::ty::TyCtxt;
+use rustc_mir_dataflow::impls::borrowed_locals;
+
+use crate::ssa::SsaLocals;
+use crate::MirPass;
+
+/// Unify locals that copy each other.
+///
+/// We consider patterns of the form
+///   _a = rvalue
+///   _b = move? _a
+///   _c = move? _a
+///   _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.
+pub struct CopyProp;
+
+impl<'tcx> MirPass<'tcx> for CopyProp {
+    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
+        sess.mir_opt_level() >= 4
+    }
+
+    #[instrument(level = "trace", skip(self, tcx, body))]
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        debug!(def_id = ?body.source.def_id());
+        propagate_ssa(tcx, body);
+    }
+}
+
+fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+    let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id());
+    let borrowed_locals = borrowed_locals(body);
+    let ssa = SsaLocals::new(tcx, param_env, body, &borrowed_locals);
+
+    let fully_moved = fully_moved_locals(&ssa, body);
+    debug!(?fully_moved);
+
+    let mut storage_to_remove = BitSet::new_empty(fully_moved.domain_size());
+    for (local, &head) in ssa.copy_classes().iter_enumerated() {
+        if local != head {
+            storage_to_remove.insert(head);
+        }
+    }
+
+    let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h);
+
+    Replacer {
+        tcx,
+        copy_classes: &ssa.copy_classes(),
+        fully_moved,
+        borrowed_locals,
+        storage_to_remove,
+    }
+    .visit_body_preserves_cfg(body);
+
+    if any_replacement {
+        crate::simplify::remove_unused_definitions(body);
+    }
+}
+
+/// `SsaLocals` computed equivalence classes between locals considering copy/move assignments.
+///
+/// This function also returns whether all the `move?` in the pattern are `move` and not copies.
+/// A local which is in the bitset can be replaced by `move _a`. Otherwise, it must be
+/// replaced by `copy _a`, as we cannot move multiple times from `_a`.
+///
+/// If an operand copies `_c`, it must happen before the assignment `_d = _c`, otherwise it is UB.
+/// This means that replacing it by a copy of `_a` if ok, since this copy happens before `_c` is
+/// moved, and therefore that `_d` is moved.
+#[instrument(level = "trace", skip(ssa, body))]
+fn fully_moved_locals(ssa: &SsaLocals, body: &Body<'_>) -> BitSet<Local> {
+    let mut fully_moved = BitSet::new_filled(body.local_decls.len());
+
+    for (_, rvalue) in ssa.assignments(body) {
+        let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) | Rvalue::CopyForDeref(place))
+            = rvalue
+        else { continue };
+
+        let Some(rhs) = place.as_local() else { continue };
+        if !ssa.is_ssa(rhs) {
+            continue;
+        }
+
+        if let Rvalue::Use(Operand::Copy(_)) | Rvalue::CopyForDeref(_) = rvalue {
+            fully_moved.remove(rhs);
+        }
+    }
+
+    ssa.meet_copy_equivalence(&mut fully_moved);
+
+    fully_moved
+}
+
+/// Utility to help performing subtitution of `*pattern` by `target`.
+struct Replacer<'a, 'tcx> {
+    tcx: TyCtxt<'tcx>,
+    fully_moved: BitSet<Local>,
+    storage_to_remove: BitSet<Local>,
+    borrowed_locals: BitSet<Local>,
+    copy_classes: &'a IndexVec<Local, Local>,
+}
+
+impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
+    fn tcx(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+
+    fn visit_local(&mut self, local: &mut Local, ctxt: PlaceContext, _: Location) {
+        let new_local = self.copy_classes[*local];
+        match ctxt {
+            // Do not modify the local in storage statements.
+            PlaceContext::NonUse(NonUseContext::StorageLive | NonUseContext::StorageDead) => {}
+            // The local should have been marked as non-SSA.
+            PlaceContext::MutatingUse(_) => assert_eq!(*local, new_local),
+            // We access the value.
+            _ => *local = new_local,
+        }
+    }
+
+    fn visit_place(&mut self, place: &mut Place<'tcx>, ctxt: PlaceContext, loc: Location) {
+        if let Some(new_projection) = self.process_projection(&place.projection, loc) {
+            place.projection = self.tcx().intern_place_elems(&new_projection);
+        }
+
+        let observes_address = match ctxt {
+            PlaceContext::NonMutatingUse(
+                NonMutatingUseContext::SharedBorrow
+                | NonMutatingUseContext::ShallowBorrow
+                | NonMutatingUseContext::UniqueBorrow
+                | NonMutatingUseContext::AddressOf,
+            ) => true,
+            // For debuginfo, merging locals is ok.
+            PlaceContext::NonUse(NonUseContext::VarDebugInfo) => {
+                self.borrowed_locals.contains(place.local)
+            }
+            _ => false,
+        };
+        if observes_address && !place.is_indirect() {
+            // We observe the address of `place.local`. Do not replace it.
+        } else {
+            self.visit_local(
+                &mut place.local,
+                PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy),
+                loc,
+            )
+        }
+    }
+
+    fn visit_operand(&mut self, operand: &mut Operand<'tcx>, loc: Location) {
+        if let Operand::Move(place) = *operand
+            && let Some(local) = place.as_local()
+            && !self.fully_moved.contains(local)
+        {
+            *operand = Operand::Copy(place);
+        }
+        self.super_operand(operand, loc);
+    }
+
+    fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) {
+        if let StatementKind::StorageDead(l) = stmt.kind
+            && self.storage_to_remove.contains(l)
+        {
+            stmt.make_nop();
+        } else if let StatementKind::Assign(box (ref place, ref mut rvalue)) = stmt.kind
+            && place.as_local().is_some()
+        {
+            // Do not replace assignments.
+            self.visit_rvalue(rvalue, loc)
+        } else {
+            self.super_statement(stmt, loc);
+        }
+    }
+}

compiler/rustc_mir_transform/src/dest_prop.rs

@@ -328,7 +328,8 @@ impl<'a, 'tcx> MutVisitor<'tcx> for Merger<'a, 'tcx> {
         match &statement.kind {
             StatementKind::Assign(box (dest, rvalue)) => {
                 match rvalue {
-                    Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) => {
+                    Rvalue::CopyForDeref(place)
+                    | Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) => {
                         // These might've been turned into self-assignments by the replacement
                         // (this includes the original statement we wanted to eliminate).
                         if dest == place {
@@ -755,7 +756,7 @@ impl<'tcx> Visitor<'tcx> for FindAssignments<'_, '_, 'tcx> {
     fn visit_statement(&mut self, statement: &Statement<'tcx>, _: Location) {
         if let StatementKind::Assign(box (
             lhs,
-            Rvalue::Use(Operand::Copy(rhs) | Operand::Move(rhs)),
+            Rvalue::CopyForDeref(rhs) | Rvalue::Use(Operand::Copy(rhs) | Operand::Move(rhs)),
         )) = &statement.kind
         {
             let Some((src, dest)) = places_to_candidate_pair(*lhs, *rhs, self.body) else {

compiler/rustc_mir_transform/src/lib.rs

@@ -54,6 +54,7 @@ mod const_debuginfo;
 mod const_goto;
 mod const_prop;
 mod const_prop_lint;
+mod copy_prop;
 mod coverage;
 mod ctfe_limit;
 mod dataflow_const_prop;
@@ -87,6 +88,7 @@ mod required_consts;
 mod reveal_all;
 mod separate_const_switch;
 mod shim;
+mod ssa;
 // This pass is public to allow external drivers to perform MIR cleanup
 pub mod simplify;
 mod simplify_branches;
@@ -563,6 +565,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &instcombine::InstCombine,
             &separate_const_switch::SeparateConstSwitch,
             &simplify::SimplifyLocals::new("before-const-prop"),
+            &copy_prop::CopyProp,
             //
             // FIXME(#70073): This pass is responsible for both optimization as well as some lints.
             &const_prop::ConstProp,

compiler/rustc_mir_transform/src/simplify.rs

@@ -404,6 +404,18 @@ impl<'tcx> MirPass<'tcx> for SimplifyLocals {
     }
 }
 
+pub fn remove_unused_definitions<'tcx>(body: &mut Body<'tcx>) {
+    // First, we're going to get a count of *actual* uses for every `Local`.
+    let mut used_locals = UsedLocals::new(body);
+
+    // Next, we're going to remove any `Local` with zero actual uses. When we remove those
+    // `Locals`, we're also going to subtract any uses of other `Locals` from the `used_locals`
+    // count. For example, if we removed `_2 = discriminant(_1)`, then we'll subtract one from
+    // `use_counts[_1]`. That in turn might make `_1` unused, so we loop until we hit a
+    // fixedpoint where there are no more unused locals.
+    remove_unused_definitions_helper(&mut used_locals, body);
+}
+
 pub fn simplify_locals<'tcx>(body: &mut Body<'tcx>, tcx: TyCtxt<'tcx>) {
     // First, we're going to get a count of *actual* uses for every `Local`.
     let mut used_locals = UsedLocals::new(body);
@@ -413,7 +425,7 @@ pub fn simplify_locals<'tcx>(body: &mut Body<'tcx>, tcx: TyCtxt<'tcx>) {
     // count. For example, if we removed `_2 = discriminant(_1)`, then we'll subtract one from
     // `use_counts[_1]`. That in turn might make `_1` unused, so we loop until we hit a
     // fixedpoint where there are no more unused locals.
-    remove_unused_definitions(&mut used_locals, body);
+    remove_unused_definitions_helper(&mut used_locals, body);
 
     // Finally, we'll actually do the work of shrinking `body.local_decls` and remapping the `Local`s.
     let map = make_local_map(&mut body.local_decls, &used_locals);
@@ -548,7 +560,7 @@ impl<'tcx> Visitor<'tcx> for UsedLocals {
 }
 
 /// Removes unused definitions. Updates the used locals to reflect the changes made.
-fn remove_unused_definitions(used_locals: &mut UsedLocals, body: &mut Body<'_>) {
+fn remove_unused_definitions_helper(used_locals: &mut UsedLocals, body: &mut Body<'_>) {
     // The use counts are updated as we remove the statements. A local might become unused
     // during the retain operation, leading to a temporary inconsistency (storage statements or
     // definitions referencing the local might remain). For correctness it is crucial that this