From b959a23eb524f59733fafac7bed6d1005eb9ff77 Mon Sep 17 00:00:00 2001
From: Camille GILLOT <gillot.camille@gmail.com>
Date: Wed, 26 Jun 2024 16:18:03 +0000
Subject: [PATCH] Move entry point to a method.

---
 .../rustc_mir_transform/src/jump_threading.rs | 78 ++++++++++---------
 1 file changed, 41 insertions(+), 37 deletions(-)

diff --git a/compiler/rustc_mir_transform/src/jump_threading.rs b/compiler/rustc_mir_transform/src/jump_threading.rs
index 0dad003ee5f42..6874568e38360 100644
--- a/compiler/rustc_mir_transform/src/jump_threading.rs
+++ b/compiler/rustc_mir_transform/src/jump_threading.rs
@@ -91,43 +91,8 @@ impl<'tcx> MirPass<'tcx> for JumpThreading {
             opportunities: Vec::new(),
         };
 
-        for (bb, bbdata) in body.basic_blocks.iter_enumerated() {
-            debug!(?bb, term = ?bbdata.terminator());
-            if bbdata.is_cleanup || loop_headers.contains(bb) {
-                continue;
-            }
-            let Some((discr, targets)) = bbdata.terminator().kind.as_switch() else { continue };
-            let Some(discr) = discr.place() else { continue };
-            debug!(?discr, ?bb);
-
-            let discr_ty = discr.ty(body, tcx).ty;
-            let Ok(discr_layout) = finder.ecx.layout_of(discr_ty) else { continue };
-
-            let Some(discr) = finder.map.find(discr.as_ref()) else { continue };
-            debug!(?discr);
-
-            let cost = CostChecker::new(tcx, param_env, None, body);
-
-            let mut state = State::new(ConditionSet::default(), finder.map);
-
-            let conds = if let Some((value, then, else_)) = targets.as_static_if() {
-                let Some(value) = ScalarInt::try_from_uint(value, discr_layout.size) else {
-                    continue;
-                };
-                arena.alloc_from_iter([
-                    Condition { value, polarity: Polarity::Eq, target: then },
-                    Condition { value, polarity: Polarity::Ne, target: else_ },
-                ])
-            } else {
-                arena.alloc_from_iter(targets.iter().filter_map(|(value, target)| {
-                    let value = ScalarInt::try_from_uint(value, discr_layout.size)?;
-                    Some(Condition { value, polarity: Polarity::Eq, target })
-                }))
-            };
-            let conds = ConditionSet(conds);
-            state.insert_value_idx(discr, conds, finder.map);
-
-            finder.find_opportunity(bb, state, cost, 0);
+        for bb in body.basic_blocks.indices() {
+            finder.start_from_switch(bb);
         }
 
         let opportunities = finder.opportunities;
@@ -215,6 +180,45 @@ impl<'tcx, 'a> TOFinder<'tcx, 'a> {
         state.all(|cs| cs.0.is_empty())
     }
 
+    /// Recursion entry point to find threading opportunities.
+    #[instrument(level = "trace", skip(self))]
+    fn start_from_switch(&mut self, bb: BasicBlock) -> Option<!> {
+        let bbdata = &self.body[bb];
+        if bbdata.is_cleanup || self.loop_headers.contains(bb) {
+            return None;
+        }
+        let (discr, targets) = bbdata.terminator().kind.as_switch()?;
+        let discr = discr.place()?;
+        debug!(?discr, ?bb);
+
+        let discr_ty = discr.ty(self.body, self.tcx).ty;
+        let discr_layout = self.ecx.layout_of(discr_ty).ok()?;
+
+        let discr = self.map.find(discr.as_ref())?;
+        debug!(?discr);
+
+        let cost = CostChecker::new(self.tcx, self.param_env, None, self.body);
+        let mut state = State::new(ConditionSet::default(), self.map);
+
+        let conds = if let Some((value, then, else_)) = targets.as_static_if() {
+            let value = ScalarInt::try_from_uint(value, discr_layout.size) ?;
+            self.arena.alloc_from_iter([
+                Condition { value, polarity: Polarity::Eq, target: then },
+                Condition { value, polarity: Polarity::Ne, target: else_ },
+            ])
+        } else {
+            self.arena.alloc_from_iter(targets.iter().filter_map(|(value, target)| {
+                let value = ScalarInt::try_from_uint(value, discr_layout.size)?;
+                Some(Condition { value, polarity: Polarity::Eq, target })
+            }))
+        };
+        let conds = ConditionSet(conds);
+        state.insert_value_idx(discr, conds, self.map);
+
+        self.find_opportunity(bb, state, cost, 0);
+        None
+    }
+
     /// Recursion entry point to find threading opportunities.
     #[instrument(level = "trace", skip(self, cost), ret)]
     fn find_opportunity(