feat(ssa): Post dominator tree

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

@@ -115,20 +115,42 @@ impl ControlFlowGraph {
     pub(crate) fn successors(
         &self,
         basic_block_id: BasicBlockId,
-    ) -> impl ExactSizeIterator<Item = BasicBlockId> + '_ {
+    ) -> impl ExactSizeIterator<Item = BasicBlockId> + DoubleEndedIterator + '_ {
         self.data
             .get(&basic_block_id)
             .expect("ICE: Attempted to iterate successors of block not found within cfg.")
             .successors
             .iter()
             .copied()
     }
+
+    /// Reverse the control flow graph
+    #[cfg(test)]
+    pub(crate) fn reverse(&self) -> Self {
+        let mut reversed_cfg = ControlFlowGraph::default();
+
+        for (block_id, node) in &self.data {
+            // For each block, reverse the edges
+            // In the reversed CFG, successors becomes predecessors
+            for &successor in &node.successors {
+                reversed_cfg.add_edge(successor, *block_id);
+            }
+        }
+
+        reversed_cfg
+    }
+
+    /// Returns the entry blocks for a CFG. This is all nodes without any predecessors.
+    pub(crate) fn compute_entry_blocks(&self) -> Vec<BasicBlockId> {
+        self.data.keys().filter(|&&block| self.predecessors(block).len() == 0).copied().collect()
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use crate::ssa::ir::{
-        call_stack::CallStackId, instruction::TerminatorInstruction, map::Id, types::Type,
+        basic_block::BasicBlockId, call_stack::CallStackId, instruction::TerminatorInstruction,
+        map::Id, types::Type,
     };
 
     use super::{super::function::Function, ControlFlowGraph};
@@ -146,8 +168,7 @@ mod tests {
         ControlFlowGraph::with_function(&func);
     }
 
-    #[test]
-    fn jumps() {
+    fn build_test_function() -> (Function, BasicBlockId, BasicBlockId, BasicBlockId) {
         // Build function of form
         // fn func {
         //   block0(cond: u1):
@@ -181,6 +202,50 @@ mod tests {
             call_stack: CallStackId::root(),
         });
 
+        (func, block0_id, block1_id, block2_id)
+    }
+
+    fn modify_test_function(
+        func: &mut Function,
+        block0_id: BasicBlockId,
+        block1_id: BasicBlockId,
+        block2_id: BasicBlockId,
+    ) -> BasicBlockId {
+        // Modify function to form:
+        // fn func {
+        //   block0(cond: u1):
+        //     jmpif cond, then: block1, else: ret_block
+        //   block1():
+        //     jmpif cond, then: block1, else: block2
+        //   block2():
+        //     jmp ret_block()
+        //   ret_block():
+        //     return ()
+        // }
+        let ret_block_id = func.dfg.make_block();
+        func.dfg[ret_block_id].set_terminator(TerminatorInstruction::Return {
+            return_values: vec![],
+            call_stack: CallStackId::root(),
+        });
+        func.dfg[block2_id].set_terminator(TerminatorInstruction::Jmp {
+            destination: ret_block_id,
+            arguments: vec![],
+            call_stack: CallStackId::root(),
+        });
+        let cond = func.dfg[block0_id].parameters()[0];
+        func.dfg[block0_id].set_terminator(TerminatorInstruction::JmpIf {
+            condition: cond,
+            then_destination: block1_id,
+            else_destination: ret_block_id,
+            call_stack: CallStackId::root(),
+        });
+        ret_block_id
+    }
+
+    #[test]
+    fn jumps() {
+        let (mut func, block0_id, block1_id, block2_id) = build_test_function();
+
         let mut cfg = ControlFlowGraph::with_function(&func);
 
         #[allow(clippy::needless_collect)]
@@ -212,33 +277,7 @@ mod tests {
             assert!(block1_successors.contains(&block2_id));
         }
 
-        // Modify function to form:
-        // fn func {
-        //   block0(cond: u1):
-        //     jmpif cond, then: block1, else: ret_block
-        //   block1():
-        //     jmpif cond, then: block1, else: block2
-        //   block2():
-        //     jmp ret_block()
-        //   ret_block():
-        //     return ()
-        // }
-        let ret_block_id = func.dfg.make_block();
-        func.dfg[ret_block_id].set_terminator(TerminatorInstruction::Return {
-            return_values: vec![],
-            call_stack: CallStackId::root(),
-        });
-        func.dfg[block2_id].set_terminator(TerminatorInstruction::Jmp {
-            destination: ret_block_id,
-            arguments: vec![],
-            call_stack: CallStackId::root(),
-        });
-        func.dfg[block0_id].set_terminator(TerminatorInstruction::JmpIf {
-            condition: cond,
-            then_destination: block1_id,
-            else_destination: ret_block_id,
-            call_stack: CallStackId::root(),
-        });
+        let ret_block_id = modify_test_function(&mut func, block0_id, block1_id, block2_id);
 
         // Recompute new and changed blocks
         cfg.recompute_block(&func, block0_id);
@@ -275,4 +314,86 @@ mod tests {
             assert!(block2_successors.contains(&ret_block_id));
         }
     }
+
+    #[test]
+    fn reversed_cfg_jumps() {
+        let (mut func, block0_id, block1_id, block2_id) = build_test_function();
+
+        let mut cfg = ControlFlowGraph::with_function(&func);
+        let reversed_cfg = cfg.reverse();
+
+        #[allow(clippy::needless_collect)]
+        {
+            let block0_predecessors: Vec<_> = reversed_cfg.predecessors(block0_id).collect();
+            let block1_predecessors: Vec<_> = reversed_cfg.predecessors(block1_id).collect();
+            let block2_predecessors: Vec<_> = reversed_cfg.predecessors(block2_id).collect();
+
+            let block0_successors: Vec<_> = reversed_cfg.successors(block0_id).collect();
+            let block1_successors: Vec<_> = reversed_cfg.successors(block1_id).collect();
+            let block2_successors: Vec<_> = reversed_cfg.successors(block2_id).collect();
+
+            assert_eq!(block0_predecessors.len(), 2);
+            assert_eq!(block1_predecessors.len(), 2);
+            assert_eq!(block2_predecessors.len(), 0);
+
+            assert!(block0_predecessors.contains(&block1_id));
+            assert!(block0_predecessors.contains(&block2_id));
+            assert!(block1_predecessors.contains(&block1_id));
+            assert!(block1_predecessors.contains(&block2_id));
+
+            assert_eq!(block0_successors.len(), 0);
+            assert_eq!(block1_successors.len(), 2);
+            assert_eq!(block2_successors.len(), 2);
+
+            assert!(block1_successors.contains(&block0_id));
+            assert!(block1_successors.contains(&block1_id));
+            assert!(block2_successors.contains(&block0_id));
+            assert!(block2_successors.contains(&block1_id));
+        }
+
+        let ret_block_id = modify_test_function(&mut func, block0_id, block1_id, block2_id);
+
+        // Recompute new and changed blocks
+        cfg.recompute_block(&func, block0_id);
+        cfg.recompute_block(&func, block2_id);
+        cfg.recompute_block(&func, ret_block_id);
+
+        let reversed_cfg = cfg.reverse();
+
+        #[allow(clippy::needless_collect)]
+        {
+            let block0_predecessors: Vec<_> = reversed_cfg.predecessors(block0_id).collect();
+            let block1_predecessors: Vec<_> = reversed_cfg.predecessors(block1_id).collect();
+            let block2_predecessors: Vec<_> = reversed_cfg.predecessors(block2_id).collect();
+            let ret_block_predecessors: Vec<_> = reversed_cfg.predecessors(ret_block_id).collect();
+
+            let block0_successors: Vec<_> = reversed_cfg.successors(block0_id).collect();
+            let block1_successors: Vec<_> = reversed_cfg.successors(block1_id).collect();
+            let block2_successors: Vec<_> = reversed_cfg.successors(block2_id).collect();
+            let ret_block_successors: Vec<_> = reversed_cfg.successors(ret_block_id).collect();
+
+            assert_eq!(block0_predecessors.len(), 2);
+            assert_eq!(block1_predecessors.len(), 2);
+            assert_eq!(block2_predecessors.len(), 1);
+            assert_eq!(ret_block_predecessors.len(), 0);
+
+            assert!(block0_predecessors.contains(&block1_id));
+            assert!(block0_predecessors.contains(&ret_block_id));
+            assert!(block1_predecessors.contains(&block1_id));
+            assert!(block1_predecessors.contains(&block2_id));
+            assert!(!block2_predecessors.contains(&block0_id));
+            assert!(block2_predecessors.contains(&ret_block_id));
+
+            assert_eq!(block0_successors.len(), 0);
+            assert_eq!(block1_successors.len(), 2);
+            assert_eq!(block2_successors.len(), 1);
+            assert_eq!(ret_block_successors.len(), 2);
+
+            assert!(block1_successors.contains(&block0_id));
+            assert!(block1_successors.contains(&block1_id));
+            assert!(block2_successors.contains(&block1_id));
+            assert!(ret_block_successors.contains(&block0_id));
+            assert!(ret_block_successors.contains(&block2_id));
+        }
+    }
 }

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

@@ -142,6 +142,9 @@ impl DominatorTree {
 
     /// Allocate and compute a dominator tree from a pre-computed control flow graph and
     /// post-order counterpart.
+    ///
+    /// This method should be used for when we want to compute a post-dominator tree.
+    /// A post-dominator tree just expects the control flow graph to be reversed.
     pub(crate) fn with_cfg_and_post_order(cfg: &ControlFlowGraph, post_order: &PostOrder) -> Self {
         let mut dom_tree = DominatorTree { nodes: HashMap::default(), cache: HashMap::default() };
         dom_tree.compute_dominator_tree(cfg, post_order);
@@ -159,6 +162,18 @@ impl DominatorTree {
         Self::with_cfg_and_post_order(&cfg, &post_order)
     }
 
+    /// Allocate and compute a post-dominator tree for the given function.
+    ///
+    /// This approach computes the reversed control flow graph and post-order internally and then
+    /// discards them. If either should be retained for reuse, it is better to instead pre-compute them
+    /// and build the dominator tree with `DominatorTree::with_cfg_and_post_order`.
+    #[cfg(test)]
+    pub(crate) fn with_function_post_dom(func: &Function) -> Self {
+        let reversed_cfg = ControlFlowGraph::with_function(func).reverse();
+        let post_order = PostOrder::with_cfg(&reversed_cfg);
+        Self::with_cfg_and_post_order(&reversed_cfg, &post_order)
+    }
+
     /// Build a dominator tree from a control flow graph using Keith D. Cooper's
     /// "Simple, Fast Dominator Algorithm."
     fn compute_dominator_tree(&mut self, cfg: &ControlFlowGraph, post_order: &PostOrder) {
@@ -403,8 +418,7 @@ mod tests {
         dt.dominates(b2, b1);
     }
 
-    #[test]
-    fn backwards_layout() {
+    fn backwards_layout_setup() -> (Function, BasicBlockId, BasicBlockId, BasicBlockId) {
         // func {
         //   block0():
         //     jmp block2()
@@ -425,10 +439,16 @@ mod tests {
         builder.terminate_with_jmp(block1_id, vec![]);
 
         let ssa = builder.finish();
-        let func = ssa.main();
+        let func = ssa.main().clone();
         let block0_id = func.entry_block();
 
-        let mut dt = DominatorTree::with_function(func);
+        (func, block0_id, block1_id, block2_id)
+    }
+
+    #[test]
+    fn backwards_layout() {
+        let (func, block0_id, block1_id, block2_id) = backwards_layout_setup();
+        let mut dt = DominatorTree::with_function(&func);
 
         // Expected dominance tree:
         // block0 {
@@ -473,6 +493,55 @@ mod tests {
         assert!(dt.dominates(block2_id, block2_id));
     }
 
+    #[test]
+    fn post_dom_backwards_layout() {
+        let (func, block0_id, block1_id, block2_id) = backwards_layout_setup();
+
+        let mut post_dom = DominatorTree::with_function_post_dom(&func);
+
+        // Expected post-dominator tree:
+        // block1 {
+        //   block2 {
+        //     block0
+        //   }
+        // }
+
+        assert_eq!(post_dom.immediate_dominator(block1_id), None);
+        assert_eq!(post_dom.immediate_dominator(block2_id), Some(block1_id));
+        assert_eq!(post_dom.immediate_dominator(block0_id), Some(block2_id));
+
+        assert_eq!(post_dom.reverse_post_order_cmp(block0_id, block0_id), Ordering::Equal);
+        assert_eq!(post_dom.reverse_post_order_cmp(block0_id, block1_id), Ordering::Greater);
+        assert_eq!(post_dom.reverse_post_order_cmp(block0_id, block2_id), Ordering::Greater);
+
+        assert_eq!(post_dom.reverse_post_order_cmp(block1_id, block0_id), Ordering::Less);
+        assert_eq!(post_dom.reverse_post_order_cmp(block1_id, block1_id), Ordering::Equal);
+        assert_eq!(post_dom.reverse_post_order_cmp(block1_id, block2_id), Ordering::Less);
+
+        assert_eq!(post_dom.reverse_post_order_cmp(block2_id, block0_id), Ordering::Less);
+        assert_eq!(post_dom.reverse_post_order_cmp(block2_id, block1_id), Ordering::Greater);
+        assert_eq!(post_dom.reverse_post_order_cmp(block2_id, block2_id), Ordering::Equal);
+
+        // Post-dominance matrix:
+        // ✓: Row item post-dominates column item
+        //    b0  b1  b2
+        // b0 ✓
+        // b1 ✓   ✓   ✓
+        // b2 ✓       ✓
+
+        assert!(post_dom.dominates(block0_id, block0_id));
+        assert!(!post_dom.dominates(block0_id, block1_id));
+        assert!(!post_dom.dominates(block0_id, block2_id));
+
+        assert!(post_dom.dominates(block1_id, block0_id));
+        assert!(post_dom.dominates(block1_id, block1_id));
+        assert!(post_dom.dominates(block1_id, block2_id));
+
+        assert!(post_dom.dominates(block2_id, block0_id));
+        assert!(!post_dom.dominates(block2_id, block1_id));
+        assert!(post_dom.dominates(block2_id, block2_id));
+    }
+
     #[test]
     fn test_find_map_dominator() {
         let (dt, b0, b1, b2, _b3) = unreachable_node_setup();

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

@@ -7,6 +7,8 @@ use std::collections::HashSet;
 
 use crate::ssa::ir::{basic_block::BasicBlockId, function::Function};
 
+use super::cfg::ControlFlowGraph;
+
 /// Depth-first traversal stack state marker for computing the cfg post-order.
 enum Visit {
     First,
@@ -25,21 +27,30 @@ impl PostOrder {
 impl PostOrder {
     /// Allocate and compute a function's block post-order.
     pub(crate) fn with_function(func: &Function) -> Self {
-        PostOrder(Self::compute_post_order(func))
+        let cfg = ControlFlowGraph::with_function(func);
+        Self::with_cfg(&cfg)
+    }
+
+    /// Allocate and compute a function's block post-order.
+    pub(crate) fn with_cfg(cfg: &ControlFlowGraph) -> Self {
+        PostOrder(Self::compute_post_order(cfg))
     }
 
     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
+    // Computes the post-order of the CFG 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())];
+    fn compute_post_order(cfg: &ControlFlowGraph) -> Vec<BasicBlockId> {
+        let mut stack = vec![];
         let mut visited: HashSet<BasicBlockId> = HashSet::new();
         let mut post_order: Vec<BasicBlockId> = Vec::new();
 
+        // Set root blocks
+        stack.extend(cfg.compute_entry_blocks().into_iter().map(|root| (Visit::First, root)));
+
         while let Some((visit, block_id)) = stack.pop() {
             match visit {
                 Visit::First => {
@@ -50,10 +61,10 @@ impl PostOrder {
                         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() {
+                        for successor_id in cfg.successors(block_id).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.
+                                // This not visited check would also be covered by the next
+                                // iteration, but checking here too saves an iteration per successor.
                                 stack.push((Visit::First, successor_id));
                             }
                         }