Improve efficiency of has_back_edge(...)

Author: bryangarza

compiler/rustc_data_structures/src/graph/dominators/mod.rs

@@ -304,13 +304,18 @@ fn compress(
     }
 }
 
+/// Tracks the list of dominators for each node.
 #[derive(Clone, Debug)]
 pub struct Dominators<N: Idx> {
     post_order_rank: IndexVec<N, usize>,
+    // Even though we track only the immediate dominator of each node, it's
+    // possible to get its full list of dominators by looking up the dominator
+    // of each dominator. (See the `impl Iterator for Iter` definition).
     immediate_dominators: IndexVec<N, Option<N>>,
 }
 
 impl<Node: Idx> Dominators<Node> {
+    /// Whether the given Node has an immediate dominator.
     pub fn is_reachable(&self, node: Node) -> bool {
         self.immediate_dominators[node].is_some()
     }
@@ -320,6 +325,8 @@ impl<Node: Idx> Dominators<Node> {
         self.immediate_dominators[node].unwrap()
     }
 
+    /// Provides an iterator over each dominator up the CFG, for the given Node.
+    /// See the `impl Iterator for Iter` definition to understand how this works.
     pub fn dominators(&self, node: Node) -> Iter<'_, Node> {
         assert!(self.is_reachable(node), "node {node:?} is not reachable");
         Iter { dominators: self, node: Some(node) }

compiler/rustc_mir_transform/src/ctfe_limit.rs

@@ -2,8 +2,9 @@
 //! (thus indicating there is a loop in the CFG), or whose terminator is a function call.
 use crate::MirPass;
 
+use rustc_data_structures::graph::dominators::Dominators;
 use rustc_middle::mir::{
-    BasicBlock, BasicBlockData, BasicBlocks, Body, Statement, StatementKind, TerminatorKind,
+    BasicBlock, BasicBlockData, Body, Statement, StatementKind, TerminatorKind,
 };
 use rustc_middle::ty::TyCtxt;
 
@@ -12,13 +13,14 @@ pub struct CtfeLimit;
 impl<'tcx> MirPass<'tcx> for CtfeLimit {
     #[instrument(skip(self, _tcx, body))]
     fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        let doms = body.basic_blocks.dominators();
         let indices: Vec<BasicBlock> = body
             .basic_blocks
             .iter_enumerated()
             .filter_map(|(node, node_data)| {
                 if matches!(node_data.terminator().kind, TerminatorKind::Call { .. })
                     // Back edges in a CFG indicate loops
-                    || has_back_edge(&body.basic_blocks, node, &node_data)
+                    || has_back_edge(&doms, node, &node_data)
                 {
                     Some(node)
                 } else {
@@ -37,17 +39,16 @@ impl<'tcx> MirPass<'tcx> for CtfeLimit {
 }
 
 fn has_back_edge(
-    basic_blocks: &BasicBlocks<'_>,
+    doms: &Dominators<BasicBlock>,
     node: BasicBlock,
     node_data: &BasicBlockData<'_>,
 ) -> bool {
-    let doms = basic_blocks.dominators();
-    basic_blocks.indices().any(|potential_dom| {
-        doms.is_reachable(potential_dom)
-            && doms.is_reachable(node)
-            && doms.is_dominated_by(node, potential_dom)
-            && node_data.terminator().successors().into_iter().any(|succ| succ == potential_dom)
-    })
+    if !doms.is_reachable(node) {
+        return false;
+    }
+    // Check if any of the dominators of the node are also the node's successor.
+    doms.dominators(node)
+        .any(|dom| node_data.terminator().successors().into_iter().any(|succ| succ == dom))
 }
 
 fn insert_counter(basic_block_data: &mut BasicBlockData<'_>) {