From f763f0e707b86dc7fee33ac230dee46254b36a71 Mon Sep 17 00:00:00 2001
From: Nick Fitzgerald <fitzgen@gmail.com>
Date: Fri, 2 Aug 2024 12:08:30 -0700
Subject: [PATCH] Cranelift: Add a helper for getting a block's successors
 (#9067)

Co-authored-by: Trevor Elliott <Telliott@fastly.com>
---
 cranelift/codegen/src/ir/function.rs | 10 ++++++
 cranelift/codegen/src/traversals.rs  | 46 +++++++++++++---------------
 2 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/cranelift/codegen/src/ir/function.rs b/cranelift/codegen/src/ir/function.rs
index f5bccf19f7..d15be02fee 100644
--- a/cranelift/codegen/src/ir/function.rs
+++ b/cranelift/codegen/src/ir/function.rs
@@ -313,6 +313,16 @@ impl FunctionStencil {
         Ok(())
     }
 
+    /// Returns an iterator over the blocks succeeding the given block.
+    pub fn block_successors(&self, block: Block) -> impl DoubleEndedIterator<Item = Block> + '_ {
+        self.layout.last_inst(block).into_iter().flat_map(|inst| {
+            self.dfg.insts[inst]
+                .branch_destination(&self.dfg.jump_tables)
+                .iter()
+                .map(|block| block.block(&self.dfg.value_lists))
+        })
+    }
+
     /// Returns true if the function is function that doesn't call any other functions. This is not
     /// to be confused with a "leaf function" in Windows terminology.
     pub fn is_leaf(&self) -> bool {
diff --git a/cranelift/codegen/src/traversals.rs b/cranelift/codegen/src/traversals.rs
index f329c2a6c2..5efb703e8a 100644
--- a/cranelift/codegen/src/traversals.rs
+++ b/cranelift/codegen/src/traversals.rs
@@ -96,31 +96,27 @@ impl Iterator for DfsIter<'_> {
 
         if event == Event::Enter && self.dfs.seen.insert(block) {
             self.dfs.stack.push((Event::Exit, block));
-            if let Some(inst) = self.func.layout.last_inst(block) {
-                self.dfs.stack.extend(
-                    self.func.dfg.insts[inst]
-                        .branch_destination(&self.func.dfg.jump_tables)
-                        .iter()
-                        // Heuristic: chase the children in reverse. This puts
-                        // the first successor block first in the postorder, all
-                        // other things being equal, which tends to prioritize
-                        // loop backedges over out-edges, putting the edge-block
-                        // closer to the loop body and minimizing live-ranges in
-                        // linear instruction space. This heuristic doesn't have
-                        // any effect on the computation of dominators, and is
-                        // purely for other consumers of the postorder we cache
-                        // here.
-                        .rev()
-                        .map(|block| block.block(&self.func.dfg.value_lists))
-                        // This is purely an optimization to avoid additional
-                        // iterations of the loop, and is not required; it's
-                        // merely inlining the check from the outer conditional
-                        // of this case to avoid the extra loop iteration. This
-                        // also avoids potential excess stack growth.
-                        .filter(|block| !self.dfs.seen.contains(*block))
-                        .map(|block| (Event::Enter, block)),
-                );
-            }
+            self.dfs.stack.extend(
+                self.func
+                    .block_successors(block)
+                    // Heuristic: chase the children in reverse. This puts
+                    // the first successor block first in the postorder, all
+                    // other things being equal, which tends to prioritize
+                    // loop backedges over out-edges, putting the edge-block
+                    // closer to the loop body and minimizing live-ranges in
+                    // linear instruction space. This heuristic doesn't have
+                    // any effect on the computation of dominators, and is
+                    // purely for other consumers of the postorder we cache
+                    // here.
+                    .rev()
+                    // This is purely an optimization to avoid additional
+                    // iterations of the loop, and is not required; it's
+                    // merely inlining the check from the outer conditional
+                    // of this case to avoid the extra loop iteration. This
+                    // also avoids potential excess stack growth.
+                    .filter(|block| !self.dfs.seen.contains(*block))
+                    .map(|block| (Event::Enter, block)),
+            );
         }
 
         Some((event, block))