loop analysis simplification

cranelift/codegen/src/dominator_tree.rs

@@ -60,6 +60,15 @@ impl DominatorTree {
         &self.postorder
     }
 
+    /// Get an iterator over CFG reverse post-order of blocks used to compute the dominator tree.
+    ///
+    /// Note that the post-order is not updated automatically when the CFG is modified. It is
+    /// computed from scratch and cached by `compute()`.
+    pub fn cfg_rpo(&self) -> impl Iterator<Item = &Block> {
+        debug_assert!(self.is_valid());
+        self.postorder.iter().rev()
+    }
+
     /// Returns the immediate dominator of `block`.
     ///
     /// `block_a` is said to *dominate* `block_b` if all control flow paths from the function

cranelift/codegen/src/loop_analysis.rs

@@ -5,12 +5,12 @@ use crate::dominator_tree::DominatorTree;
 use crate::entity::entity_impl;
 use crate::entity::SecondaryMap;
 use crate::entity::{Keys, PrimaryMap};
-use crate::flowgraph::{BlockPredecessor, ControlFlowGraph};
+use crate::flowgraph::ControlFlowGraph;
 use crate::ir::{Block, Function, Layout};
 use crate::packed_option::PackedOption;
 use crate::timing;
 use alloc::vec::Vec;
-use smallvec::{smallvec, SmallVec};
+use smallvec::SmallVec;
 
 /// A opaque reference to a code loop.
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
@@ -190,6 +190,19 @@ impl LoopAnalysis {
         self.valid = false;
     }
 
+    // Determines if a block dominates any predecessor
+    // and thus is a loop header.
+    fn is_block_loop_header(
+        block: Block,
+        cfg: &ControlFlowGraph,
+        domtree: &DominatorTree,
+        layout: &Layout,
+    ) -> bool {
+        // A block is a loop header if it dominates any of its predecessors.
+        cfg.pred_iter(block)
+            .any(|pred| domtree.dominates(block, pred.inst, layout))
+    }
+
     // Traverses the CFG in reverse postorder and create a loop object for every block having a
     // back edge.
     fn find_loop_headers(
@@ -198,21 +211,13 @@ impl LoopAnalysis {
         domtree: &DominatorTree,
         layout: &Layout,
     ) {
-        // We traverse the CFG in reverse postorder
-        for &block in domtree.cfg_postorder().iter().rev() {
-            for BlockPredecessor {
-                inst: pred_inst, ..
-            } in cfg.pred_iter(block)
-            {
-                // If the block dominates one of its predecessors it is a back edge
-                if domtree.dominates(block, pred_inst, layout) {
-                    // This block is a loop header, so we create its associated loop
-                    let lp = self.loops.push(LoopData::new(block, None));
-                    self.block_loop_map[block] = lp.into();
-                    break;
-                    // We break because we only need one back edge to identify a loop header.
-                }
-            }
+        for &block in domtree
+            .cfg_rpo()
+            .filter(|&&block| Self::is_block_loop_header(block, cfg, domtree, layout))
+        {
+            // This block is a loop header, so we create its associated loop
+            let lp = self.loops.push(LoopData::new(block, None));
+            self.block_loop_map[block] = lp.into();
         }
     }
 
@@ -229,16 +234,15 @@ impl LoopAnalysis {
         // We handle each loop header in reverse order, corresponding to a pseudo postorder
         // traversal of the graph.
         for lp in self.loops().rev() {
-            for BlockPredecessor {
-                block: pred,
-                inst: pred_inst,
-            } in cfg.pred_iter(self.loops[lp].header)
-            {
-                // We follow the back edges
-                if domtree.dominates(self.loops[lp].header, pred_inst, layout) {
-                    stack.push(pred);
-                }
-            }
+            // Push all predecessors of this header that it dominates onto the stack.
+            stack.extend(
+                cfg.pred_iter(self.loops[lp].header)
+                    .filter(|pred| {
+                        // We follow the back edges
+                        domtree.dominates(self.loops[lp].header, pred.inst, layout)
+                    })
+                    .map(|pred| pred.block),
+            );
             while let Some(node) = stack.pop() {
                 let continue_dfs: Option<Block>;
                 match self.block_loop_map[node].expand() {
@@ -283,16 +287,14 @@ impl LoopAnalysis {
                 // Now we have handled the popped node and need to continue the DFS by adding the
                 // predecessors of that node
                 if let Some(continue_dfs) = continue_dfs {
-                    for BlockPredecessor { block: pred, .. } in cfg.pred_iter(continue_dfs) {
-                        stack.push(pred)
-                    }
+                    stack.extend(cfg.pred_iter(continue_dfs).map(|pred| pred.block));
                 }
             }
         }
     }
 
     fn assign_loop_levels(&mut self) {
-        let mut stack: SmallVec<[Loop; 8]> = smallvec![];
+        let mut stack: SmallVec<[Loop; 8]> = SmallVec::new();
         for lp in self.loops.keys() {
             if self.loops[lp].level == LoopLevel::invalid() {
                 stack.push(lp);

cranelift/codegen/src/machinst/blockorder.rs

@@ -192,7 +192,7 @@ impl BlockLoweringOrder {
 
         let mut lowered_order = Vec::new();
 
-        for &block in domtree.cfg_postorder().iter().rev() {
+        for &block in domtree.cfg_rpo() {
             lowered_order.push(LoweredBlock::Orig { block });
 
             if block_out_count[block] > 1 {

cranelift/codegen/src/remove_constant_phis.rs

@@ -233,7 +233,7 @@ pub fn do_remove_constant_phis(func: &mut Function, domtree: &mut DominatorTree)
     let mut summaries =
         SecondaryMap::<Block, BlockSummary>::with_capacity(domtree.cfg_postorder().len());
 
-    for b in domtree.cfg_postorder().iter().rev().copied() {
+    for b in domtree.cfg_rpo().copied() {
         let formals = func.dfg.block_params(b);
         let mut summary = BlockSummary::new(&bump, formals);
 
@@ -269,7 +269,7 @@ pub fn do_remove_constant_phis(func: &mut Function, domtree: &mut DominatorTree)
     // Set up initial solver state
     let mut state = SolverState::new();
 
-    for b in domtree.cfg_postorder().iter().rev().copied() {
+    for b in domtree.cfg_rpo().copied() {
         // For each block, get the formals
         if b == entry_block {
             continue;
@@ -288,7 +288,7 @@ pub fn do_remove_constant_phis(func: &mut Function, domtree: &mut DominatorTree)
         iter_no += 1;
         let mut changed = false;
 
-        for src in domtree.cfg_postorder().iter().rev().copied() {
+        for src in domtree.cfg_rpo().copied() {
             let src_summary = &summaries[src];
             for edge in &src_summary.dests {
                 assert!(edge.block != entry_block);