JumpThreading: Avoid computing dominators to identify loop headers.

compiler/rustc_mir_transform/src/jump_threading.rs

@@ -84,7 +84,7 @@ impl<'tcx> crate::MirPass<'tcx> for JumpThreading {
             body,
             arena,
             map: Map::new(tcx, body, Some(MAX_PLACES)),
-            loop_headers: loop_headers(body),
+            maybe_loop_headers: maybe_loop_headers(body),
             opportunities: Vec::new(),
         };
 
@@ -100,7 +100,7 @@ impl<'tcx> crate::MirPass<'tcx> for JumpThreading {
 
         // Verify that we do not thread through a loop header.
         for to in opportunities.iter() {
-            assert!(to.chain.iter().all(|&block| !finder.loop_headers.contains(block)));
+            assert!(to.chain.iter().all(|&block| !finder.maybe_loop_headers.contains(block)));
         }
         OpportunitySet::new(body, opportunities).apply(body);
     }
@@ -124,7 +124,7 @@ struct TOFinder<'a, 'tcx> {
     ecx: InterpCx<'tcx, DummyMachine>,
     body: &'a Body<'tcx>,
     map: Map<'tcx>,
-    loop_headers: DenseBitSet<BasicBlock>,
+    maybe_loop_headers: DenseBitSet<BasicBlock>,
     /// We use an arena to avoid cloning the slices when cloning `state`.
     arena: &'a DroplessArena,
     opportunities: Vec<ThreadingOpportunity>,
@@ -190,7 +190,7 @@ impl<'a, 'tcx> TOFinder<'a, 'tcx> {
     #[instrument(level = "trace", skip(self))]
     fn start_from_switch(&mut self, bb: BasicBlock) {
         let bbdata = &self.body[bb];
-        if bbdata.is_cleanup || self.loop_headers.contains(bb) {
+        if bbdata.is_cleanup || self.maybe_loop_headers.contains(bb) {
             return;
         }
         let Some((discr, targets)) = bbdata.terminator().kind.as_switch() else { return };
@@ -235,7 +235,7 @@ impl<'a, 'tcx> TOFinder<'a, 'tcx> {
         depth: usize,
     ) {
         // Do not thread through loop headers.
-        if self.loop_headers.contains(bb) {
+        if self.maybe_loop_headers.contains(bb) {
             return;
         }
 
@@ -833,20 +833,28 @@ enum Update {
     Decr,
 }
 
-/// Compute the set of loop headers in the given body. We define a loop header as a block which has
-/// at least a predecessor which it dominates. This definition is only correct for reducible CFGs.
-/// But if the CFG is already irreducible, there is no point in trying much harder.
-/// is already irreducible.
-fn loop_headers(body: &Body<'_>) -> DenseBitSet<BasicBlock> {
-    let mut loop_headers = DenseBitSet::new_empty(body.basic_blocks.len());
-    let dominators = body.basic_blocks.dominators();
-    // Only visit reachable blocks.
-    for (bb, bbdata) in traversal::preorder(body) {
+/// Compute the set of loop headers in the given body. A loop header is usually defined as a block
+/// which dominates one of its predecessors. This definition is only correct for reducible CFGs.
+/// However, computing dominators is expensive, so we approximate according to the post-order
+/// traversal order. A loop header for us is a block which is visited after its predecessor in
+/// post-order. This is ok as we mostly need a heuristic.
+fn maybe_loop_headers(body: &Body<'_>) -> DenseBitSet<BasicBlock> {
+    let mut maybe_loop_headers = DenseBitSet::new_empty(body.basic_blocks.len());
+    let mut visited = DenseBitSet::new_empty(body.basic_blocks.len());
+    for (bb, bbdata) in traversal::postorder(body) {
+        // Post-order means we visit successors before the block for acyclic CFGs.
+        // If the successor is not visited yet, consider it a loop header.
         for succ in bbdata.terminator().successors() {
-            if dominators.dominates(succ, bb) {
-                loop_headers.insert(succ);
+            if !visited.contains(succ) {
+                maybe_loop_headers.insert(succ);
             }
         }
+
+        // Only mark `bb` as visited after we checked the successors, in case we have a self-loop.
+        //     bb1: goto -> bb1;
+        let _new = visited.insert(bb);
+        debug_assert!(_new);
     }
-    loop_headers
+
+    maybe_loop_headers
 }