From 11cfc16a60dcca27a1872fa3e83df35d797e8521 Mon Sep 17 00:00:00 2001
From: dianqk <dianqk@dianqk.net>
Date: Sun, 20 Apr 2025 20:56:41 +0800
Subject: [PATCH] mir-opt: Use one MirPatch in MatchBranchSimplification

---
 .../rustc_mir_transform/src/match_branches.rs | 37 +++++++------------
 1 file changed, 14 insertions(+), 23 deletions(-)

diff --git a/compiler/rustc_mir_transform/src/match_branches.rs b/compiler/rustc_mir_transform/src/match_branches.rs
index b37241185c935..8c0c309689902 100644
--- a/compiler/rustc_mir_transform/src/match_branches.rs
+++ b/compiler/rustc_mir_transform/src/match_branches.rs
@@ -19,30 +19,32 @@ impl<'tcx> crate::MirPass<'tcx> for MatchBranchSimplification {
 
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         let typing_env = body.typing_env(tcx);
-        let mut should_cleanup = false;
-        for bb_idx in body.basic_blocks.indices() {
-            match &body.basic_blocks[bb_idx].terminator().kind {
+        let mut apply_patch = false;
+        let mut patch = MirPatch::new(body);
+        for (bb, bb_data) in body.basic_blocks.iter_enumerated() {
+            match &bb_data.terminator().kind {
                 TerminatorKind::SwitchInt {
                     discr: Operand::Copy(_) | Operand::Move(_),
                     targets,
                     ..
                     // We require that the possible target blocks don't contain this block.
-                } if !targets.all_targets().contains(&bb_idx) => {}
+                } if !targets.all_targets().contains(&bb) => {}
                 // Only optimize switch int statements
                 _ => continue,
             };
 
-            if SimplifyToIf.simplify(tcx, body, bb_idx, typing_env).is_some() {
-                should_cleanup = true;
+            if SimplifyToIf.simplify(tcx, body, &mut patch, bb, typing_env).is_some() {
+                apply_patch = true;
                 continue;
             }
-            if SimplifyToExp::default().simplify(tcx, body, bb_idx, typing_env).is_some() {
-                should_cleanup = true;
+            if SimplifyToExp::default().simplify(tcx, body, &mut patch, bb, typing_env).is_some() {
+                apply_patch = true;
                 continue;
             }
         }
 
-        if should_cleanup {
+        if apply_patch {
+            patch.apply(body);
             simplify_cfg(tcx, body);
         }
     }
@@ -59,7 +61,8 @@ trait SimplifyMatch<'tcx> {
     fn simplify(
         &mut self,
         tcx: TyCtxt<'tcx>,
-        body: &mut Body<'tcx>,
+        body: &Body<'tcx>,
+        patch: &mut MirPatch<'tcx>,
         switch_bb_idx: BasicBlock,
         typing_env: ty::TypingEnv<'tcx>,
     ) -> Option<()> {
@@ -73,8 +76,6 @@ trait SimplifyMatch<'tcx> {
         let discr_ty = discr.ty(body.local_decls(), tcx);
         self.can_simplify(tcx, targets, typing_env, bbs, discr_ty)?;
 
-        let mut patch = MirPatch::new(body);
-
         // Take ownership of items now that we know we can optimize.
         let discr = discr.clone();
 
@@ -87,19 +88,9 @@ trait SimplifyMatch<'tcx> {
         let parent_end = Location { block: switch_bb_idx, statement_index };
         patch.add_statement(parent_end, StatementKind::StorageLive(discr_local));
         patch.add_assign(parent_end, Place::from(discr_local), Rvalue::Use(discr));
-        self.new_stmts(
-            tcx,
-            targets,
-            typing_env,
-            &mut patch,
-            parent_end,
-            bbs,
-            discr_local,
-            discr_ty,
-        );
+        self.new_stmts(tcx, targets, typing_env, patch, parent_end, bbs, discr_local, discr_ty);
         patch.add_statement(parent_end, StatementKind::StorageDead(discr_local));
         patch.patch_terminator(switch_bb_idx, bbs[first].terminator().kind.clone());
-        patch.apply(body);
         Some(())
     }