Auto merge of #140464 - oli-obk:successors-mut-perf, r=petrochenkov

Use a closure instead of three chained iterators

Fixes the perf regression from #123948

That PR had chained a third option to the iterator which apparently didn't optimize well

7 files changed, 61 insertions, 76 deletions

diff --git a/compiler/rustc_middle/src/mir/terminator.rs b/compiler/rustc_middle/src/mir/terminator.rs
index 8a1ead7d19d..0834fa8844c 100644
--- a/compiler/rustc_middle/src/mir/terminator.rs
+++ b/compiler/rustc_middle/src/mir/terminator.rs
@@ -437,8 +437,8 @@ impl<'tcx> Terminator<'tcx> {
     }
 
     #[inline]
-    pub fn successors_mut(&mut self) -> SuccessorsMut<'_> {
-        self.kind.successors_mut()
+    pub fn successors_mut<'a>(&'a mut self, f: impl FnMut(&'a mut BasicBlock)) {
+        self.kind.successors_mut(f)
     }
 
     #[inline]
@@ -486,7 +486,6 @@ pub use helper::*;
 mod helper {
     use super::*;
     pub type Successors<'a> = impl DoubleEndedIterator<Item = BasicBlock> + 'a;
-    pub type SuccessorsMut<'a> = impl DoubleEndedIterator<Item = &'a mut BasicBlock> + 'a;
 
     impl SwitchTargets {
         /// Like [`SwitchTargets::target_for_value`], but returning the same type as
@@ -560,69 +559,63 @@ mod helper {
         }
 
         #[inline]
-        #[define_opaque(SuccessorsMut)]
-        pub fn successors_mut(&mut self) -> SuccessorsMut<'_> {
+        pub fn successors_mut<'a>(&'a mut self, mut f: impl FnMut(&'a mut BasicBlock)) {
             use self::TerminatorKind::*;
-            match *self {
-                // 3-successors for async drop: target, unwind, dropline (parent coroutine drop)
-                Drop {
-                    target: ref mut t,
-                    unwind: UnwindAction::Cleanup(ref mut u),
-                    drop: Some(ref mut d),
-                    ..
-                } => slice::from_mut(t).into_iter().chain(Some(u).into_iter().chain(Some(d))),
-                // 2-successors
-                Call {
-                    target: Some(ref mut t), unwind: UnwindAction::Cleanup(ref mut u), ..
+            match self {
+                Drop { target, unwind, drop, .. } => {
+                    f(target);
+                    if let UnwindAction::Cleanup(u) = unwind {
+                        f(u)
+                    }
+                    if let Some(d) = drop {
+                        f(d)
+                    }
+                }
+                Call { target, unwind, .. } => {
+                    if let Some(target) = target {
+                        f(target);
+                    }
+                    if let UnwindAction::Cleanup(u) = unwind {
+                        f(u)
+                    }
                 }
-                | Yield { resume: ref mut t, drop: Some(ref mut u), .. }
-                | Drop {
-                    target: ref mut t,
-                    unwind: UnwindAction::Cleanup(ref mut u),
-                    drop: None,
-                    ..
+                Yield { resume, drop, .. } => {
+                    f(resume);
+                    if let Some(d) = drop {
+                        f(d)
+                    }
                 }
-                | Drop { target: ref mut t, unwind: _, drop: Some(ref mut u), .. }
-                | Assert { target: ref mut t, unwind: UnwindAction::Cleanup(ref mut u), .. }
-                | FalseUnwind {
-                    real_target: ref mut t,
-                    unwind: UnwindAction::Cleanup(ref mut u),
-                } => slice::from_mut(t).into_iter().chain(Some(u).into_iter().chain(None)),
-                // single successor
-                Goto { target: ref mut t }
-                | Call { target: None, unwind: UnwindAction::Cleanup(ref mut t), .. }
-                | Call { target: Some(ref mut t), unwind: _, .. }
-                | Yield { resume: ref mut t, drop: None, .. }
-                | Drop { target: ref mut t, unwind: _, .. }
-                | Assert { target: ref mut t, unwind: _, .. }
-                | FalseUnwind { real_target: ref mut t, unwind: _ } => {
-                    slice::from_mut(t).into_iter().chain(None.into_iter().chain(None))
+                Assert { target, unwind, .. } | FalseUnwind { real_target: target, unwind } => {
+                    f(target);
+                    if let UnwindAction::Cleanup(u) = unwind {
+                        f(u)
+                    }
+                }
+                Goto { target } => {
+                    f(target);
                 }
-                // No successors
                 UnwindResume
                 | UnwindTerminate(_)
                 | CoroutineDrop
                 | Return
                 | Unreachable
-                | TailCall { .. }
-                | Call { target: None, unwind: _, .. } => {
-                    (&mut []).into_iter().chain(None.into_iter().chain(None))
-                }
-                // Multiple successors
-                InlineAsm { ref mut targets, unwind: UnwindAction::Cleanup(ref mut u), .. } => {
-                    targets.iter_mut().chain(Some(u).into_iter().chain(None))
-                }
-                InlineAsm { ref mut targets, unwind: _, .. } => {
-                    targets.iter_mut().chain(None.into_iter().chain(None))
+                | TailCall { .. } => {}
+                InlineAsm { targets, unwind, .. } => {
+                    for target in targets {
+                        f(target);
+                    }
+                    if let UnwindAction::Cleanup(u) = unwind {
+                        f(u)
+                    }
                 }
-                SwitchInt { ref mut targets, .. } => {
-                    targets.targets.iter_mut().chain(None.into_iter().chain(None))
+                SwitchInt { targets, .. } => {
+                    for target in &mut targets.targets {
+                        f(target);
+                    }
                 }
-                // FalseEdge
-                FalseEdge { ref mut real_target, ref mut imaginary_target } => {
-                    slice::from_mut(real_target)
-                        .into_iter()
-                        .chain(Some(imaginary_target).into_iter().chain(None))
+                FalseEdge { real_target, imaginary_target } => {
+                    f(real_target);
+                    f(imaginary_target);
                 }
             }
         }
diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs
index 263f0c40f5a..66f106bec6f 100644
--- a/compiler/rustc_mir_transform/src/coroutine.rs
+++ b/compiler/rustc_mir_transform/src/coroutine.rs
@@ -1064,10 +1064,8 @@ fn insert_switch<'tcx>(
         },
     );
 
-    let blocks = body.basic_blocks_mut().iter_mut();
-
-    for target in blocks.flat_map(|b| b.terminator_mut().successors_mut()) {
-        *target += 1;
+    for b in body.basic_blocks_mut().iter_mut() {
+        b.terminator_mut().successors_mut(|target| *target += 1);
     }
 }
 
diff --git a/compiler/rustc_mir_transform/src/jump_threading.rs b/compiler/rustc_mir_transform/src/jump_threading.rs
index ada2c0b76cf..31b361ec1a9 100644
--- a/compiler/rustc_mir_transform/src/jump_threading.rs
+++ b/compiler/rustc_mir_transform/src/jump_threading.rs
@@ -757,12 +757,12 @@ impl OpportunitySet {
 
             // Replace `succ` by `new_succ` where it appears.
             let mut num_edges = 0;
-            for s in basic_blocks[current].terminator_mut().successors_mut() {
+            basic_blocks[current].terminator_mut().successors_mut(|s| {
                 if *s == succ {
                     *s = new_succ;
                     num_edges += 1;
                 }
-            }
+            });
 
             // Update predecessors with the new block.
             let _new_succ = self.predecessors.push(num_edges);
diff --git a/compiler/rustc_mir_transform/src/prettify.rs b/compiler/rustc_mir_transform/src/prettify.rs
index 8ccfbe2f194..8217feff24e 100644
--- a/compiler/rustc_mir_transform/src/prettify.rs
+++ b/compiler/rustc_mir_transform/src/prettify.rs
@@ -115,9 +115,7 @@ impl<'tcx> MutVisitor<'tcx> for BasicBlockUpdater<'tcx> {
     }
 
     fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, _location: Location) {
-        for succ in terminator.successors_mut() {
-            *succ = self.map[*succ];
-        }
+        terminator.successors_mut(|succ| *succ = self.map[*succ]);
     }
 }
 
diff --git a/compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs b/compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs
index 1dd34005d66..797056ad52d 100644
--- a/compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs
+++ b/compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs
@@ -58,13 +58,13 @@ impl<'tcx> crate::MirPass<'tcx> for RemoveNoopLandingPads {
                 }
             }
 
-            for target in body[bb].terminator_mut().successors_mut() {
+            body[bb].terminator_mut().successors_mut(|target| {
                 if *target != resume_block && nop_landing_pads.contains(*target) {
                     debug!("    folding noop jump to {:?} to resume block", target);
                     *target = resume_block;
                     jumps_folded += 1;
                 }
-            }
+            });
 
             let is_nop_landing_pad = self.is_nop_landing_pad(bb, body, &nop_landing_pads);
             if is_nop_landing_pad {
diff --git a/compiler/rustc_mir_transform/src/shim.rs b/compiler/rustc_mir_transform/src/shim.rs
index 0d9a04b760a..9688ac8ed2e 100644
--- a/compiler/rustc_mir_transform/src/shim.rs
+++ b/compiler/rustc_mir_transform/src/shim.rs
@@ -298,10 +298,9 @@ fn local_decls_for_sig<'tcx>(
 fn dropee_emit_retag<'tcx>(
     tcx: TyCtxt<'tcx>,
     body: &mut Body<'tcx>,
-    dropee_ptr: Place<'tcx>,
+    mut dropee_ptr: Place<'tcx>,
     span: Span,
 ) -> Place<'tcx> {
-    let mut dropee_ptr = dropee_ptr;
     if tcx.sess.opts.unstable_opts.mir_emit_retag {
         let source_info = SourceInfo::outermost(span);
         // We want to treat the function argument as if it was passed by `&mut`. As such, we
@@ -365,8 +364,8 @@ fn build_drop_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, ty: Option<Ty<'tcx>>)
         new_body(source, blocks, local_decls_for_sig(&sig, span), sig.inputs().len(), span);
 
     // The first argument (index 0), but add 1 for the return value.
-    let mut dropee_ptr = Place::from(Local::new(1 + 0));
-    dropee_ptr = dropee_emit_retag(tcx, &mut body, dropee_ptr, span);
+    let dropee_ptr = Place::from(Local::new(1 + 0));
+    let dropee_ptr = dropee_emit_retag(tcx, &mut body, dropee_ptr, span);
 
     if ty.is_some() {
         let patch = {
diff --git a/compiler/rustc_mir_transform/src/simplify.rs b/compiler/rustc_mir_transform/src/simplify.rs
index 4f2cce8ac10..8f88228d9bb 100644
--- a/compiler/rustc_mir_transform/src/simplify.rs
+++ b/compiler/rustc_mir_transform/src/simplify.rs
@@ -147,9 +147,8 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
                 let mut terminator =
                     self.basic_blocks[bb].terminator.take().expect("invalid terminator state");
 
-                for successor in terminator.successors_mut() {
-                    self.collapse_goto_chain(successor, &mut changed);
-                }
+                terminator
+                    .successors_mut(|successor| self.collapse_goto_chain(successor, &mut changed));
 
                 let mut inner_changed = true;
                 merged_blocks.clear();
@@ -375,9 +374,7 @@ pub(super) fn remove_dead_blocks(body: &mut Body<'_>) {
     }
 
     for block in basic_blocks {
-        for target in block.terminator_mut().successors_mut() {
-            *target = replacements[target.index()];
-        }
+        block.terminator_mut().successors_mut(|target| *target = replacements[target.index()]);
     }
 }