From 599e8db8389111d5c27934c5294879d9caf262e1 Mon Sep 17 00:00:00 2001
From: Camille Gillot <gillot.camille@gmail.com>
Date: Sat, 27 Sep 2025 00:14:32 +0000
Subject: [PATCH] Use MirPatch in simplify_branches.

---
 compiler/rustc_mir_transform/src/patch.rs     | 30 +++++++++++++++----
 .../src/simplify_branches.rs                  | 19 +++++++-----
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/compiler/rustc_mir_transform/src/patch.rs b/compiler/rustc_mir_transform/src/patch.rs
index d831ab50b1ac0..cc8ea76011bf6 100644
--- a/compiler/rustc_mir_transform/src/patch.rs
+++ b/compiler/rustc_mir_transform/src/patch.rs
@@ -11,6 +11,8 @@ use tracing::debug;
 /// once with `apply`. This is useful for MIR transformation passes.
 pub(crate) struct MirPatch<'tcx> {
     term_patch_map: FxHashMap<BasicBlock, TerminatorKind<'tcx>>,
+    /// Set of statements that should be replaced by `Nop`.
+    nop_statements: Vec<Location>,
     new_blocks: Vec<BasicBlockData<'tcx>>,
     new_statements: Vec<(Location, StatementKind<'tcx>)>,
     new_locals: Vec<LocalDecl<'tcx>>,
@@ -33,6 +35,7 @@ impl<'tcx> MirPatch<'tcx> {
     pub(crate) fn new(body: &Body<'tcx>) -> Self {
         let mut result = MirPatch {
             term_patch_map: Default::default(),
+            nop_statements: vec![],
             new_blocks: vec![],
             new_statements: vec![],
             new_locals: vec![],
@@ -212,6 +215,15 @@ impl<'tcx> MirPatch<'tcx> {
         self.term_patch_map.insert(block, new);
     }
 
+    /// Mark given statement to be replaced by a `Nop`.
+    ///
+    /// This method only works on statements from the initial body, and cannot be used to remove
+    /// statements from `add_statement` or `add_assign`.
+    #[tracing::instrument(level = "debug", skip(self))]
+    pub(crate) fn nop_statement(&mut self, loc: Location) {
+        self.nop_statements.push(loc);
+    }
+
     /// Queues the insertion of a statement at a given location. The statement
     /// currently at that location, and all statements that follow, are shifted
     /// down. If multiple statements are queued for addition at the same
@@ -257,11 +269,8 @@ impl<'tcx> MirPatch<'tcx> {
         bbs.extend(self.new_blocks);
         body.local_decls.extend(self.new_locals);
 
-        // The order in which we patch terminators does not change the result.
-        #[allow(rustc::potential_query_instability)]
-        for (src, patch) in self.term_patch_map {
-            debug!("MirPatch: patching block {:?}", src);
-            bbs[src].terminator_mut().kind = patch;
+        for loc in self.nop_statements {
+            bbs[loc.block].statements[loc.statement_index].make_nop();
         }
 
         let mut new_statements = self.new_statements;
@@ -285,6 +294,17 @@ impl<'tcx> MirPatch<'tcx> {
                 .insert(loc.statement_index, Statement::new(source_info, stmt));
             delta += 1;
         }
+
+        // The order in which we patch terminators does not change the result.
+        #[allow(rustc::potential_query_instability)]
+        for (src, patch) in self.term_patch_map {
+            debug!("MirPatch: patching block {:?}", src);
+            let bb = &mut bbs[src];
+            if let TerminatorKind::Unreachable = patch {
+                bb.statements.clear();
+            }
+            bb.terminator_mut().kind = patch;
+        }
     }
 
     fn source_info_for_index(data: &BasicBlockData<'_>, loc: Location) -> SourceInfo {
diff --git a/compiler/rustc_mir_transform/src/simplify_branches.rs b/compiler/rustc_mir_transform/src/simplify_branches.rs
index 886f4d6e50900..ed94a058ec6d9 100644
--- a/compiler/rustc_mir_transform/src/simplify_branches.rs
+++ b/compiler/rustc_mir_transform/src/simplify_branches.rs
@@ -2,6 +2,8 @@ use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
 use tracing::trace;
 
+use crate::patch::MirPatch;
+
 pub(super) enum SimplifyConstCondition {
     AfterConstProp,
     Final,
@@ -19,8 +21,10 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyConstCondition {
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         trace!("Running SimplifyConstCondition on {:?}", body.source);
         let typing_env = body.typing_env(tcx);
-        'blocks: for block in body.basic_blocks_mut() {
-            for stmt in block.statements.iter_mut() {
+        let mut patch = MirPatch::new(body);
+
+        'blocks: for (bb, block) in body.basic_blocks.iter_enumerated() {
+            for (statement_index, stmt) in block.statements.iter().enumerate() {
                 // Simplify `assume` of a known value: either a NOP or unreachable.
                 if let StatementKind::Intrinsic(box ref intrinsic) = stmt.kind
                     && let NonDivergingIntrinsic::Assume(discr) = intrinsic
@@ -28,17 +32,16 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyConstCondition {
                     && let Some(constant) = c.const_.try_eval_bool(tcx, typing_env)
                 {
                     if constant {
-                        stmt.make_nop();
+                        patch.nop_statement(Location { block: bb, statement_index });
                     } else {
-                        block.statements.clear();
-                        block.terminator_mut().kind = TerminatorKind::Unreachable;
+                        patch.patch_terminator(bb, TerminatorKind::Unreachable);
                         continue 'blocks;
                     }
                 }
             }
 
-            let terminator = block.terminator_mut();
-            terminator.kind = match terminator.kind {
+            let terminator = block.terminator();
+            let terminator = match terminator.kind {
                 TerminatorKind::SwitchInt {
                     discr: Operand::Constant(ref c), ref targets, ..
                 } => {
@@ -58,7 +61,9 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyConstCondition {
                 },
                 _ => continue,
             };
+            patch.patch_terminator(bb, terminator);
         }
+        patch.apply(body);
     }
 
     fn is_required(&self) -> bool {