chore: Rework DIE pass to operate on a postorder SCC CFG

Cargo.toml

@@ -163,6 +163,7 @@ tracing = "0.1.40"
 tracing-web = "0.1.3"
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 rust-embed = "6.6.0"
+petgraph = "0.6"
 
 [profile.dev]
 # This is required to be able to run `cargo test` in acvm_js due to the `locals exceeds maximum` error.

compiler/noirc_evaluator/Cargo.toml

@@ -29,6 +29,7 @@ tracing.workspace = true
 chrono = "0.4.37"
 rayon.workspace = true
 cfg-if.workspace = true
+petgraph.workspace = true
 
 [dev-dependencies]
 proptest.workspace = true

compiler/noirc_evaluator/src/ssa/ir/function.rs

@@ -2,8 +2,12 @@
 
 use iter_extended::vecmap;
 use noirc_frontend::monomorphization::ast::InlineType;
+use petgraph::graph::NodeIndex;
+use petgraph::prelude::DiGraph;
 use serde::{Deserialize, Serialize};
 
+use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet};
+
 use super::basic_block::BasicBlockId;
 use super::dfg::DataFlowGraph;
 use super::instruction::TerminatorInstruction;
@@ -187,6 +191,46 @@
 
         unreachable!("SSA Function {} has no reachable return instruction!", self.id())
     }
+
+    /// Return each SCC (strongly-connected component) of the CFG of this function
+    /// in postorder. In practice this is identical to a normal CFG except that
+    /// blocks within loops will all be grouped together in the same SCC.
+    pub(crate) fn postorder_scc_cfg(&self) -> Vec<Vec<BasicBlockId>> {
+        let mut cfg = DiGraph::new();
+        let mut stack = vec![self.entry_block];
+
+        let mut visited = HashSet::default();
+        let mut block_to_index = HashMap::default();
+        let mut index_to_block = HashMap::default();
+
+        // Add or create a block node in the cfg
+        let mut get_block_index = |cfg: &mut DiGraph<_, _>, block| {
+            block_to_index.get(&block).copied().unwrap_or_else(|| {
+                let index = cfg.add_node(block);
+                block_to_index.insert(block, index);
+                index_to_block.insert(index, block);
+                index
+            })
+        };
+
+        // Populate each reachable block & edges between them
+        while let Some(block) = stack.pop() {
+            if visited.insert(block) {
+                let block_index = get_block_index(&mut cfg, block);
+
+                for successor in self.dfg[block].successors() {
+                    stack.push(successor);
+                    let successor_index = get_block_index(&mut cfg, successor);
+                    cfg.add_edge(block_index, successor_index, ());
+                }
+            }
+        }
+
+        // Perform tarjan_scc to get strongly connected components.
+        // Lucky for us, this already returns SCCs in postorder.
+        let postorder: Vec<Vec<NodeIndex>> = petgraph::algo::tarjan_scc(&cfg);
+        vecmap(postorder, |indices| vecmap(indices, |index| index_to_block[&index]))
+    }
 }
 
 impl std::fmt::Display for RuntimeType {

compiler/noirc_evaluator/src/ssa/ir/post_order.rs

@@ -15,6 +15,12 @@ enum Visit {
 
 pub(crate) struct PostOrder(Vec<BasicBlockId>);
 
+/// An SCC Post Order is a post-order visit of each strongly-connected
+/// component in a function's control flow graph. In practice this is
+/// identical to a normal post order except that blocks within loops
+/// will all be grouped together within the same visit step.
+pub(crate) struct SccPostOrder(Vec<Vec<BasicBlockId>>);
+
 impl PostOrder {
     pub(crate) fn as_slice(&self) -> &[BasicBlockId] {
         self.0.as_slice()
@@ -24,49 +30,49 @@ impl PostOrder {
 impl PostOrder {
     /// Allocate and compute a function's block post-order. Pos
     pub(crate) fn with_function(func: &Function) -> Self {
-        PostOrder(Self::compute_post_order(func))
+        PostOrder(compute_post_order(func))
     }
 
     pub(crate) fn into_vec(self) -> Vec<BasicBlockId> {
         self.0
     }
+}
 
-    // Computes the post-order of the function by doing a depth-first traversal of the
-    // function's entry block's previously unvisited children. Each block is sequenced according
-    // to when the traversal exits it.
-    fn compute_post_order(func: &Function) -> Vec<BasicBlockId> {
-        let mut stack = vec![(Visit::First, func.entry_block())];
-        let mut visited: HashSet<BasicBlockId> = HashSet::new();
-        let mut post_order: Vec<BasicBlockId> = Vec::new();
-
-        while let Some((visit, block_id)) = stack.pop() {
-            match visit {
-                Visit::First => {
-                    if !visited.contains(&block_id) {
-                        // This is the first time we pop the block, so we need to scan its
-                        // successors and then revisit it.
-                        visited.insert(block_id);
-                        stack.push((Visit::Last, block_id));
-                        // Stack successors for visiting. Because items are taken from the top of the
-                        // stack, we push the item that's due for a visit first to the top.
-                        for successor_id in func.dfg[block_id].successors().rev() {
-                            if !visited.contains(&successor_id) {
-                                // This not visited check would also be cover by the next
-                                // iteration, but checking here two saves an iteration per successor.
-                                stack.push((Visit::First, successor_id));
-                            }
+// Computes the post-order of the function by doing a depth-first traversal of the
+// function's entry block's previously unvisited children. Each block is sequenced according
+// to when the traversal exits it.
+fn compute_post_order(func: &Function) -> Vec<BasicBlockId> {
+    let mut stack = vec![(Visit::First, func.entry_block())];
+    let mut visited: HashSet<BasicBlockId> = HashSet::new();
+    let mut post_order: Vec<BasicBlockId> = Vec::new();
+
+    while let Some((visit, block_id)) = stack.pop() {
+        match visit {
+            Visit::First => {
+                if !visited.contains(&block_id) {
+                    // This is the first time we pop the block, so we need to scan its
+                    // successors and then revisit it.
+                    visited.insert(block_id);
+                    stack.push((Visit::Last, block_id));
+                    // Stack successors for visiting. Because items are taken from the top of the
+                    // stack, we push the item that's due for a visit first to the top.
+                    for successor_id in func.dfg[block_id].successors() {
+                        if !visited.contains(&successor_id) {
+                            // This not visited check would also be cover by the next
+                            // iteration, but checking here two saves an iteration per successor.
+                            stack.push((Visit::First, successor_id));
                         }
                     }
                 }
+            }
 
-                Visit::Last => {
-                    // We've finished all this node's successors.
-                    post_order.push(block_id);
-                }
+            Visit::Last => {
+                // We've finished all this node's successors.
+                post_order.push(block_id);
             }
         }
-        post_order
     }
+    post_order
 }
 
 #[cfg(test)]
@@ -104,7 +110,7 @@ mod tests {
         //   D (seen)
         // } -> push(A)
         // Result:
-        // D, F, E, B, A, (C dropped as unreachable)
+        // F, E, B, D, A, (C dropped as unreachable)
 
         let func_id = Id::test_new(0);
         let mut builder = FunctionBuilder::new("func".into(), func_id);
@@ -145,6 +151,6 @@ mod tests {
         let func = ssa.main();
         let post_order = PostOrder::with_function(func);
         let block_a_id = func.entry_block();
-        assert_eq!(post_order.0, [block_d_id, block_f_id, block_e_id, block_b_id, block_a_id]);
+        assert_eq!(post_order.0, [block_f_id, block_e_id, block_b_id, block_d_id, block_a_id]);
     }
 }