Auto merge of rust-lang#128299 - DianQK:clone-copy, r=<try>

Simplify the canonical clone method and the copy-like forms to copy

Fixes rust-lang#128081.

r? `@cjgillot`

compiler/rustc_mir_transform/src/gvn.rs

@@ -876,6 +876,95 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
         None
     }
 
+    fn try_as_place_elem(
+        &mut self,
+        proj: ProjectionElem<VnIndex, Ty<'tcx>>,
+        loc: Location,
+    ) -> Option<PlaceElem<'tcx>> {
+        Some(match proj {
+            ProjectionElem::Deref => ProjectionElem::Deref,
+            ProjectionElem::Field(idx, ty) => ProjectionElem::Field(idx, ty),
+            ProjectionElem::Index(idx) => {
+                let Some(local) = self.try_as_local(idx, loc) else {
+                    return None;
+                };
+                self.reused_locals.insert(local);
+                ProjectionElem::Index(local)
+            }
+            ProjectionElem::ConstantIndex { offset, min_length, from_end } => {
+                ProjectionElem::ConstantIndex { offset, min_length, from_end }
+            }
+            ProjectionElem::Subslice { from, to, from_end } => {
+                ProjectionElem::Subslice { from, to, from_end }
+            }
+            ProjectionElem::Downcast(symbol, idx) => ProjectionElem::Downcast(symbol, idx),
+            ProjectionElem::OpaqueCast(idx) => ProjectionElem::OpaqueCast(idx),
+            ProjectionElem::Subtype(idx) => ProjectionElem::Subtype(idx),
+        })
+    }
+
+    fn simplify_aggregate_to_copy(
+        &mut self,
+        rvalue: &mut Rvalue<'tcx>,
+        location: Location,
+        fields: &[VnIndex],
+        variant_index: VariantIdx,
+    ) -> Option<VnIndex> {
+        let Some(&first_field) = fields.first() else {
+            return None;
+        };
+        let Value::Projection(copy_from_value, _) = *self.get(first_field) else {
+            return None;
+        };
+        // All fields must correspond one-to-one and come from the same aggregate value.
+        if fields.iter().enumerate().any(|(index, &v)| {
+            if let Value::Projection(pointer, ProjectionElem::Field(from_index, _)) = *self.get(v)
+                && copy_from_value == pointer
+                && from_index.index() == index
+            {
+                return false;
+            }
+            true
+        }) {
+            return None;
+        }
+
+        let mut copy_from_local_value = copy_from_value;
+        if let Value::Projection(pointer, proj) = *self.get(copy_from_value)
+            && let ProjectionElem::Downcast(_, read_variant) = proj
+        {
+            if variant_index == read_variant {
+                // When copying a variant, there is no need to downcast.
+                copy_from_local_value = pointer;
+            } else {
+                // The copied variant must be identical.
+                return None;
+            }
+        }
+
+        let tcx = self.tcx;
+        let mut projection = SmallVec::<[PlaceElem<'tcx>; 1]>::new();
+        loop {
+            if let Some(local) = self.try_as_local(copy_from_local_value, location) {
+                projection.reverse();
+                let place = Place { local, projection: tcx.mk_place_elems(projection.as_slice()) };
+                if rvalue.ty(self.local_decls, tcx) == place.ty(self.local_decls, tcx).ty {
+                    self.reused_locals.insert(local);
+                    *rvalue = Rvalue::Use(Operand::Copy(place));
+                    return Some(copy_from_value);
+                }
+                return None;
+            } else if let Value::Projection(pointer, proj) = *self.get(copy_from_local_value)
+                && let Some(proj) = self.try_as_place_elem(proj, location)
+            {
+                projection.push(proj);
+                copy_from_local_value = pointer;
+            } else {
+                return None;
+            }
+        }
+    }
+
     fn simplify_aggregate(
         &mut self,
         rvalue: &mut Rvalue<'tcx>,
@@ -972,6 +1061,13 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
             }
         }
 
+        if let AggregateTy::Def(_, _) = ty
+            && let Some(value) =
+                self.simplify_aggregate_to_copy(rvalue, location, &fields, variant_index)
+        {
+            return Some(value);
+        }
+
         Some(self.insert(Value::Aggregate(ty, variant_index, fields)))
     }
 
@@ -1485,7 +1581,7 @@ impl<'tcx> VnState<'_, 'tcx> {
     }
 
     /// If there is a local which is assigned `index`, and its assignment strictly dominates `loc`,
-    /// return it.
+    /// return it. If you used this local, add it to `reused_locals` to remove storage statements.
     fn try_as_local(&mut self, index: VnIndex, loc: Location) -> Option<Local> {
         let other = self.rev_locals.get(index)?;
         other

compiler/rustc_mir_transform/src/lib.rs

@@ -588,15 +588,14 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             // Now, we need to shrink the generated MIR.
             &ref_prop::ReferencePropagation,
             &sroa::ScalarReplacementOfAggregates,
-            &match_branches::MatchBranchSimplification,
-            // inst combine is after MatchBranchSimplification to clean up Ne(_1, false)
             &multiple_return_terminators::MultipleReturnTerminators,
             // After simplifycfg, it allows us to discover new opportunities for peephole optimizations.
             &instsimplify::InstSimplify::AfterSimplifyCfg,
             &simplify::SimplifyLocals::BeforeConstProp,
             &dead_store_elimination::DeadStoreElimination::Initial,
             &gvn::GVN,
             &simplify::SimplifyLocals::AfterGVN,
+            &match_branches::MatchBranchSimplification,
             &dataflow_const_prop::DataflowConstProp,
             &single_use_consts::SingleUseConsts,
             &o1(simplify_branches::SimplifyConstCondition::AfterConstProp),

compiler/rustc_mir_transform/src/match_branches.rs

@@ -1,10 +1,16 @@
-use std::iter;
+use std::{iter, usize};
 
+use rustc_const_eval::const_eval::mk_eval_cx_for_const_val;
+use rustc_index::bit_set::BitSet;
 use rustc_index::IndexSlice;
 use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::*;
+use rustc_middle::ty;
 use rustc_middle::ty::layout::{IntegerExt, TyAndLayout};
+use rustc_middle::ty::util::Discr;
 use rustc_middle::ty::{ParamEnv, ScalarInt, Ty, TyCtxt};
+use rustc_mir_dataflow::impls::{borrowed_locals, MaybeTransitiveLiveLocals};
+use rustc_mir_dataflow::Analysis;
 use rustc_target::abi::Integer;
 use rustc_type_ir::TyKind::*;
 
@@ -48,6 +54,10 @@ impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
                 should_cleanup = true;
                 continue;
             }
+            if simplify_to_copy(tcx, body, bb_idx, param_env).is_some() {
+                should_cleanup = true;
+                continue;
+            }
         }
 
         if should_cleanup {
@@ -515,3 +525,202 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
         }
     }
 }
+
+/// This is primarily used to merge these copy statements that simplified the canonical enum clone method by GVN.
+/// The GVN simplified
+/// ```ignore (syntax-highlighting-only)
+/// match a {
+///     Foo::A(x) => Foo::A(*x),
+///     Foo::B => Foo::B
+/// }
+/// ```
+/// to
+/// ```ignore (syntax-highlighting-only)
+/// match a {
+///     Foo::A(_x) => a, // copy a
+///     Foo::B => Foo::B
+/// }
+/// ```
+/// This function will simplify into a copy statement.
+fn simplify_to_copy<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    body: &mut Body<'tcx>,
+    switch_bb_idx: BasicBlock,
+    param_env: ParamEnv<'tcx>,
+) -> Option<()> {
+    let bbs = &body.basic_blocks;
+    // Check if the copy source matches the following pattern.
+    // _2 = discriminant(*_1); // "*_1" is the expected the copy source.
+    // switchInt(move _2) -> [0: bb3, 1: bb2, otherwise: bb1];
+    let &Statement {
+        kind: StatementKind::Assign(box (discr_place, Rvalue::Discriminant(expected_src_place))),
+        ..
+    } = bbs[switch_bb_idx].statements.last()?
+    else {
+        return None;
+    };
+    let expected_src_ty = expected_src_place.ty(body.local_decls(), tcx);
+    if !expected_src_ty.ty.is_enum() || expected_src_ty.variant_index.is_some() {
+        return None;
+    }
+    let targets = match bbs[switch_bb_idx].terminator().kind {
+        TerminatorKind::SwitchInt { ref discr, ref targets, .. }
+            if discr.place() == Some(discr_place) =>
+        {
+            targets
+        }
+        _ => return None,
+    };
+    // We require that the possible target blocks all be distinct.
+    if !targets.is_distinct() {
+        return None;
+    }
+    if !bbs[targets.otherwise()].is_empty_unreachable() {
+        return None;
+    }
+    // Check that destinations are identical, and if not, then don't optimize this block.
+    let mut target_iter = targets.iter();
+    let first_terminator_kind = &bbs[target_iter.next().unwrap().1].terminator().kind;
+    if !target_iter
+        .all(|(_, other_target)| first_terminator_kind == &bbs[other_target].terminator().kind)
+    {
+        return None;
+    }
+
+    let borrowed_locals = borrowed_locals(body);
+    let mut live = None;
+    let mut expected_dest_place = None;
+    for (index, target_bb) in targets.iter() {
+        let stmts = &bbs[target_bb].statements;
+        if stmts.is_empty() {
+            return None;
+        }
+        if let [Statement { kind: StatementKind::Assign(box (place, rvalue)), .. }] =
+            bbs[target_bb].statements.as_slice()
+        {
+            let dest_ty = place.ty(body.local_decls(), tcx);
+            if dest_ty.ty != expected_src_ty.ty || dest_ty.variant_index.is_some() {
+                return None;
+            }
+            let ty::Adt(def, _) = dest_ty.ty.kind() else {
+                return None;
+            };
+            if *expected_dest_place.get_or_insert(*place) != *place {
+                return None;
+            }
+            match rvalue {
+                // Check if `_3 = const Foo::B` can be transformed to `_3 = copy *_1`.
+                Rvalue::Use(Operand::Constant(box constant))
+                    if let Const::Val(const_, ty) = constant.const_ =>
+                {
+                    let (ecx, op) =
+                        mk_eval_cx_for_const_val(tcx.at(constant.span), param_env, const_, ty)?;
+                    let variant = ecx.read_discriminant(&op).ok()?;
+                    if !def.variants()[variant].fields.is_empty() {
+                        return None;
+                    }
+                    let Discr { val, .. } = ty.discriminant_for_variant(tcx, variant)?;
+                    if val != index {
+                        return None;
+                    }
+                }
+                Rvalue::Use(Operand::Copy(src_place)) if *src_place == expected_src_place => {}
+                // Check if `_3 = Foo::B` can be transformed to `_3 = copy *_1`.
+                Rvalue::Aggregate(box AggregateKind::Adt(_, variant_index, _, _, None), fields)
+                    if fields.is_empty()
+                        && let Some(Discr { val, .. }) =
+                            expected_src_ty.ty.discriminant_for_variant(tcx, *variant_index)
+                        && val == index => {}
+                _ => return None,
+            }
+        } else {
+            // If the BB contains more than one statement, we have to check if these statements can be ignored.
+            let mut lived_stmts: BitSet<usize> =
+                BitSet::new_filled(bbs[target_bb].statements.len());
+            let mut dest_place = None;
+            for (statement_index, statement) in bbs[target_bb].statements.iter().enumerate().rev() {
+                let loc = Location { block: target_bb, statement_index };
+                if let StatementKind::Assign(assign) = &statement.kind {
+                    if !assign.1.is_safe_to_remove() {
+                        return None;
+                    }
+                }
+                match &statement.kind {
+                    StatementKind::Assign(box (place, _))
+                    | StatementKind::SetDiscriminant { place: box place, .. }
+                    | StatementKind::Deinit(box place) => {
+                        if place.is_indirect() || borrowed_locals.contains(place.local) {
+                            return None;
+                        }
+                        let live = live.get_or_insert_with(|| {
+                            MaybeTransitiveLiveLocals::new(&borrowed_locals)
+                                .into_engine(tcx, body)
+                                .iterate_to_fixpoint()
+                                .into_results_cursor(body)
+                        });
+                        live.seek_before_primary_effect(loc);
+                        if !live.get().contains(place.local) {
+                            lived_stmts.remove(statement_index);
+                        } else if matches!(
+                            &statement.kind,
+                            StatementKind::Assign(box (_, Rvalue::Use(Operand::Copy(_))))
+                        ) && dest_place.is_none()
+                        {
+                            // There is only one statement that cannot be ignored that can be used as an expected copy statement.
+                            dest_place = Some(*place);
+                        } else {
+                            return None;
+                        }
+                    }
+                    StatementKind::StorageLive(_)
+                    | StatementKind::StorageDead(_)
+                    | StatementKind::Nop => (),
+
+                    StatementKind::Retag(_, _)
+                    | StatementKind::Coverage(_)
+                    | StatementKind::Intrinsic(_)
+                    | StatementKind::ConstEvalCounter
+                    | StatementKind::PlaceMention(_)
+                    | StatementKind::FakeRead(_)
+                    | StatementKind::AscribeUserType(_, _) => {
+                        return None;
+                    }
+                }
+            }
+            let dest_place = dest_place?;
+            if *expected_dest_place.get_or_insert(dest_place) != dest_place {
+                return None;
+            }
+            // We can ignore the paired StorageLive and StorageDead.
+            let mut storage_live_locals: BitSet<Local> = BitSet::new_empty(body.local_decls.len());
+            for stmt_index in lived_stmts.iter() {
+                let statement = &bbs[target_bb].statements[stmt_index];
+                match &statement.kind {
+                    StatementKind::Assign(box (place, Rvalue::Use(Operand::Copy(src_place))))
+                        if *place == dest_place && *src_place == expected_src_place => {}
+                    StatementKind::StorageLive(local)
+                        if *local != dest_place.local && storage_live_locals.insert(*local) => {}
+                    StatementKind::StorageDead(local)
+                        if *local != dest_place.local && storage_live_locals.remove(*local) => {}
+                    StatementKind::Nop => {}
+                    _ => return None,
+                }
+            }
+            if !storage_live_locals.is_empty() {
+                return None;
+            }
+        }
+    }
+    let expected_dest_place = expected_dest_place?;
+    let statement_index = bbs[switch_bb_idx].statements.len();
+    let parent_end = Location { block: switch_bb_idx, statement_index };
+    let mut patch = MirPatch::new(body);
+    patch.add_assign(
+        parent_end,
+        expected_dest_place,
+        Rvalue::Use(Operand::Copy(expected_src_place)),
+    );
+    patch.patch_terminator(switch_bb_idx, first_terminator_kind.clone());
+    patch.apply(body);
+    Some(())
+}