diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir.rs b/crates/noirc_evaluator/src/ssa_refactor/ir.rs index 1f6cca9157d..d52f380d3d4 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir.rs @@ -2,9 +2,11 @@ pub(crate) mod basic_block; pub(crate) mod cfg; pub(crate) mod constant; pub(crate) mod dfg; +pub(crate) mod dom; pub(crate) mod function; pub(crate) mod instruction; pub(crate) mod map; +pub(crate) mod post_order; pub(crate) mod printer; pub(crate) mod types; pub(crate) mod value; diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs index 8a3f74c4a64..e8b09f518d8 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block.rs @@ -76,7 +76,9 @@ impl BasicBlock { /// Iterate over all the successors of the currently block, as determined by /// the blocks jumped to in the terminator instruction. If there is no terminator /// instruction yet, this will iterate 0 times. - pub(crate) fn successors(&self) -> impl ExactSizeIterator { + pub(crate) fn successors( + &self, + ) -> impl ExactSizeIterator + DoubleEndedIterator { match &self.terminator { Some(TerminatorInstruction::Jmp { destination, .. }) => vec![*destination].into_iter(), Some(TerminatorInstruction::JmpIf { then_destination, else_destination, .. }) => { diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/dom.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/dom.rs new file mode 100644 index 00000000000..9a0916f62c8 --- /dev/null +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/dom.rs @@ -0,0 +1,433 @@ +//! The dominator tree of a function, represented as a hash map of each reachable block id to its +//! immediate dominator. +//! +//! Dominator trees are useful for tasks such as identifying back-edges in loop analysis or +//! calculating dominance frontiers. + +use std::{cmp::Ordering, collections::HashMap}; + +use super::{ + basic_block::BasicBlockId, cfg::ControlFlowGraph, function::Function, post_order::PostOrder, +}; + +/// Dominator tree node. We keep one of these per reachable block. +#[derive(Clone, Default)] +struct DominatorTreeNode { + /// The block's idx in the control flow graph's reverse post-order + reverse_post_order_idx: u32, + + /// The block that immediately dominated that of the node in question. + /// + /// This will be None for the entry block, which has no immediate dominator. + immediate_dominator: Option, +} + +impl DominatorTreeNode { + /// Updates the immediate dominator estimate, returning true if it has changed. + /// + /// This is used internally as a shorthand during `compute_dominator_tree`. + pub(self) fn update_estimate(&mut self, immediate_dominator: BasicBlockId) -> bool { + let immediate_dominator = Some(immediate_dominator); + if self.immediate_dominator == immediate_dominator { + false + } else { + self.immediate_dominator = immediate_dominator; + true + } + } +} + +/// The dominator tree for a single function. +pub(crate) struct DominatorTree { + /// The nodes of the dominator tree + /// + /// After dominator tree computation has complete, this will contain a node for every + /// reachable block, and no nodes for unreachable blocks. + nodes: HashMap, +} + +/// Methods for querying the dominator tree. +impl DominatorTree { + /// Is `block_id` reachable from the entry block? + pub(crate) fn is_reachable(&self, block_id: BasicBlockId) -> bool { + self.nodes.contains_key(&block_id) + } + + /// Returns the immediate dominator of `block_id`. + /// + /// A block is said to *dominate* `block_id` if all control flow paths from the function + /// entry to `block_id` must go through the block. + /// + /// The *immediate dominator* is the dominator that is closest to `block_id`. All other + /// dominators also dominate the immediate dominator. + /// + /// This returns `None` if `block_id` is not reachable from the entry block, or if it is the + /// entry block which has no dominators. + pub(crate) fn immediate_dominator(&self, block_id: BasicBlockId) -> Option { + self.nodes.get(&block_id).and_then(|node| node.immediate_dominator) + } + + /// Compare two blocks relative to the reverse post-order. + pub(crate) fn reverse_post_order_cmp(&self, a: BasicBlockId, b: BasicBlockId) -> Ordering { + match (self.nodes.get(&a), self.nodes.get(&b)) { + (Some(a), Some(b)) => a.reverse_post_order_idx.cmp(&b.reverse_post_order_idx), + _ => unreachable!("Post order for unreachable block is undefined"), + } + } + + /// Returns `true` if `block_a_id` dominates `block_b_id`. + /// + /// This means that every control-flow path from the function entry to `block_b_id` must go + /// through `block_a_id`. + /// + /// This function panics if either of the blocks are unreachable. + /// + /// An instruction is considered to dominate itself. + pub(crate) fn dominates(&self, block_a_id: BasicBlockId, mut block_b_id: BasicBlockId) -> bool { + // Walk up the dominator tree from "b" until we encounter or pass "a". Doing the + // comparison on the reverse post-order may allows to test whether we have passed "a" + // without waiting until we reach the root of the tree. + loop { + match self.reverse_post_order_cmp(block_a_id, block_b_id) { + Ordering::Less => { + block_b_id = match self.immediate_dominator(block_b_id) { + Some(immediate_dominator) => immediate_dominator, + None => return false, // a is unreachable, so we climbed past the entry + } + } + Ordering::Greater => return false, + Ordering::Equal => return true, + } + } + } + + /// Allocate and compute a dominator tree from a pre-computed control flow graph and + /// post-order counterpart. + pub(crate) fn with_cfg_and_post_order(cfg: &ControlFlowGraph, post_order: &PostOrder) -> Self { + let mut dom_tree = DominatorTree { nodes: HashMap::new() }; + dom_tree.compute_dominator_tree(cfg, post_order); + dom_tree + } + + /// Allocate and compute a dominator tree for the given function. + /// + /// This approach computes the control flow graph and post-order internally and then + /// discards them. If either should be retained reuse it is better to instead pre-compute them + /// and build the dominator tree with `DominatorTree::with_cfg_and_post_order`. + pub(crate) fn with_function(func: &Function) -> Self { + let cfg = ControlFlowGraph::with_function(func); + let post_order = PostOrder::with_function(func); + Self::with_cfg_and_post_order(&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) { + // We'll be iterating over a reverse post-order of the CFG, skipping the entry block. + let (entry_block_id, entry_free_post_order) = post_order + .as_slice() + .split_last() + .expect("ICE: functions always have at least one block"); + + // Do a first pass where we assign reverse post-order indices to all reachable nodes. The + // entry block will be the only node with no immediate dominator. + self.nodes.insert( + *entry_block_id, + DominatorTreeNode { reverse_post_order_idx: 0, immediate_dominator: None }, + ); + for (i, &block_id) in entry_free_post_order.iter().rev().enumerate() { + // Indices have been displaced by 1 by the removal of the entry node + let reverse_post_order_idx = i as u32 + 1; + + // Due to the nature of the post-order traversal, every node we visit will have at + // least one predecessor that has previously been assigned during this loop. + let immediate_dominator = self.compute_immediate_dominator(block_id, cfg); + self.nodes.insert( + block_id, + DominatorTreeNode { + immediate_dominator: Some(immediate_dominator), + reverse_post_order_idx, + }, + ); + } + + // Now that we have reverse post-order indices for everything and initial immediate + // dominator estimates, iterate until convergence. + // + // If the function is free of irreducible control flow, this will exit after one iteration. + let mut changed = true; + while changed { + changed = false; + for &block_id in entry_free_post_order.iter().rev() { + let immediate_dominator = self.compute_immediate_dominator(block_id, cfg); + changed = self + .nodes + .get_mut(&block_id) + .expect("Assigned in first pass") + .update_estimate(immediate_dominator); + } + } + } + + // Compute the immediate dominator for `block_id` using the pre-calculate immediate dominators + // of reachable nodes. + fn compute_immediate_dominator( + &self, + block_id: BasicBlockId, + cfg: &ControlFlowGraph, + ) -> BasicBlockId { + // Get an iterator with just the reachable, already visited predecessors to `block_id`. + // Note that during the first pass `node` was pre-populated with all reachable blocks. + let mut reachable_predecessors = + cfg.predecessors(block_id).filter(|pred_id| self.nodes.contains_key(pred_id)); + + // This function isn't called on unreachable blocks or the entry block, so the reverse + // post-order will contain at least one predecessor to this block. + let mut immediate_dominator = + reachable_predecessors.next().expect("block node must have one reachable predecessor"); + + for predecessor in reachable_predecessors { + immediate_dominator = self.common_dominator(immediate_dominator, predecessor); + } + + immediate_dominator + } + + /// Compute the common dominator of two basic blocks. + /// + /// Both basic blocks are assumed to be reachable. + fn common_dominator( + &self, + mut block_a_id: BasicBlockId, + mut block_b_id: BasicBlockId, + ) -> BasicBlockId { + loop { + match self.reverse_post_order_cmp(block_a_id, block_b_id) { + Ordering::Less => { + // "a" comes before "b" in the reverse post-order. Move "b" up. + block_b_id = self.nodes[&block_b_id] + .immediate_dominator + .expect("Unreachable basic block?"); + } + Ordering::Greater => { + // "b" comes before "a" in the reverse post-order. Move "a" up. + block_a_id = self.nodes[&block_a_id] + .immediate_dominator + .expect("Unreachable basic block?"); + } + Ordering::Equal => break, + } + } + + debug_assert_eq!(block_a_id, block_b_id, "Unreachable block passed to common_dominator?"); + block_a_id + } +} + +#[cfg(test)] +mod tests { + use std::cmp::Ordering; + + use crate::ssa_refactor::{ + ir::{ + basic_block::BasicBlockId, dom::DominatorTree, function::Function, + instruction::TerminatorInstruction, map::Id, types::Type, + }, + ssa_builder::FunctionBuilder, + }; + + #[test] + fn empty() { + let func_id = Id::test_new(0); + let mut func = Function::new("func".into(), func_id); + let block0_id = func.entry_block(); + func.dfg.set_block_terminator( + block0_id, + TerminatorInstruction::Return { return_values: vec![] }, + ); + let dom_tree = DominatorTree::with_function(&func); + assert!(dom_tree.dominates(block0_id, block0_id)); + } + + // Testing setup for a function with an unreachable block2 + fn unreachable_node_setup( + ) -> (DominatorTree, BasicBlockId, BasicBlockId, BasicBlockId, BasicBlockId) { + // func() { + // block0(cond: u1): + // jmpif v0 block2() block3() + // block1(): + // jmp block2() + // block2(): + // jmp block3() + // block3(): + // return () + // } + let func_id = Id::test_new(0); + let mut builder = FunctionBuilder::new("func".into(), func_id); + + let cond = builder.add_parameter(Type::unsigned(1)); + let block1_id = builder.insert_block(); + let block2_id = builder.insert_block(); + let block3_id = builder.insert_block(); + + builder.terminate_with_jmpif(cond, block2_id, block3_id); + builder.switch_to_block(block1_id); + builder.terminate_with_jmp(block2_id, vec![]); + builder.switch_to_block(block2_id); + builder.terminate_with_jmp(block3_id, vec![]); + builder.switch_to_block(block3_id); + builder.terminate_with_return(vec![]); + + let ssa = builder.finish(); + let func = ssa.functions.first().unwrap(); + let block0_id = func.entry_block(); + + let dt = DominatorTree::with_function(func); + (dt, block0_id, block1_id, block2_id, block3_id) + } + + // Expected dominator tree + // block0 { + // block2 + // block3 + // } + + // Dominance matrix + // ✓: Row item dominates column item + // !: Querying row item's dominance of column item panics. (i.e. invalid) + // b0 b1 b2 b3 + // b0 ✓ ! ✓ ✓ + // b1 ! ! ! ! + // b2 ! ✓ + // b3 ! ✓ + // Note that from a local view block 1 dominates blocks 1,2 & 3, but since this block is + // unreachable, performing this query indicates an internal compiler error. + #[test] + fn unreachable_node_asserts() { + let (dt, b0, _b1, b2, b3) = unreachable_node_setup(); + + assert!(dt.dominates(b0, b0)); + assert!(dt.dominates(b0, b2)); + assert!(dt.dominates(b0, b3)); + + assert!(!dt.dominates(b2, b0)); + assert!(dt.dominates(b2, b2)); + assert!(!dt.dominates(b2, b3)); + + assert!(!dt.dominates(b3, b0)); + assert!(!dt.dominates(b3, b2)); + assert!(dt.dominates(b3, b3)); + } + + #[test] + #[should_panic] + fn unreachable_node_panic_b0_b1() { + let (dt, b0, b1, _b2, _b3) = unreachable_node_setup(); + dt.dominates(b0, b1); + } + + #[test] + #[should_panic] + fn unreachable_node_panic_b1_b0() { + let (dt, b0, b1, _b2, _b3) = unreachable_node_setup(); + dt.dominates(b1, b0); + } + + #[test] + #[should_panic] + fn unreachable_node_panic_b1_b1() { + let (dt, _b0, b1, _b2, _b3) = unreachable_node_setup(); + dt.dominates(b1, b1); + } + + #[test] + #[should_panic] + fn unreachable_node_panic_b1_b2() { + let (dt, _b0, b1, b2, _b3) = unreachable_node_setup(); + dt.dominates(b1, b2); + } + + #[test] + #[should_panic] + fn unreachable_node_panic_b1_b3() { + let (dt, _b0, b1, _b2, b3) = unreachable_node_setup(); + dt.dominates(b1, b3); + } + + #[test] + #[should_panic] + fn unreachable_node_panic_b3_b1() { + let (dt, _b0, b1, b2, _b3) = unreachable_node_setup(); + dt.dominates(b2, b1); + } + + #[test] + fn backwards_layout() { + // func { + // block0(): + // jmp block2() + // block1(): + // return () + // block2(): + // jump block1() + // } + let func_id = Id::test_new(0); + let mut builder = FunctionBuilder::new("func".into(), func_id); + let block1_id = builder.insert_block(); + let block2_id = builder.insert_block(); + + builder.terminate_with_jmp(block2_id, vec![]); + builder.switch_to_block(block1_id); + builder.terminate_with_return(vec![]); + builder.switch_to_block(block2_id); + builder.terminate_with_jmp(block1_id, vec![]); + + let ssa = builder.finish(); + let func = ssa.functions.first().unwrap(); + let block0_id = func.entry_block(); + + let dt = DominatorTree::with_function(func); + + // Expected dominance tree: + // block0 { + // block2 { + // block1 + // } + // } + + assert_eq!(dt.immediate_dominator(block0_id), None); + assert_eq!(dt.immediate_dominator(block1_id), Some(block2_id)); + assert_eq!(dt.immediate_dominator(block2_id), Some(block0_id)); + + assert_eq!(dt.reverse_post_order_cmp(block0_id, block0_id), Ordering::Equal); + assert_eq!(dt.reverse_post_order_cmp(block0_id, block1_id), Ordering::Less); + assert_eq!(dt.reverse_post_order_cmp(block0_id, block2_id), Ordering::Less); + + assert_eq!(dt.reverse_post_order_cmp(block1_id, block0_id), Ordering::Greater); + assert_eq!(dt.reverse_post_order_cmp(block1_id, block1_id), Ordering::Equal); + assert_eq!(dt.reverse_post_order_cmp(block1_id, block2_id), Ordering::Greater); + + assert_eq!(dt.reverse_post_order_cmp(block2_id, block0_id), Ordering::Greater); + assert_eq!(dt.reverse_post_order_cmp(block2_id, block1_id), Ordering::Less); + assert_eq!(dt.reverse_post_order_cmp(block2_id, block2_id), Ordering::Equal); + + // Dominance matrix: + // ✓: Row item dominates column item + // b0 b1 b2 + // b0 ✓ ✓ ✓ + // b1 ✓ + // b2 ✓ ✓ + + assert!(dt.dominates(block0_id, block0_id)); + assert!(dt.dominates(block0_id, block1_id)); + assert!(dt.dominates(block0_id, block2_id)); + + assert!(!dt.dominates(block1_id, block0_id)); + assert!(dt.dominates(block1_id, block1_id)); + assert!(!dt.dominates(block1_id, block2_id)); + + assert!(!dt.dominates(block2_id, block0_id)); + assert!(dt.dominates(block2_id, block1_id)); + assert!(dt.dominates(block2_id, block2_id)); + } +} diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/post_order.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/post_order.rs new file mode 100644 index 00000000000..984f10a64af --- /dev/null +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/post_order.rs @@ -0,0 +1,163 @@ +//! The post-order for a given function represented as a vector of basic block ids. +//! +//! This ordering is beneficial to the efficiency of various algorithms, such as those for dead +//! code elimination and calculating dominance trees. + +use std::collections::HashSet; + +use crate::ssa_refactor::ir::{basic_block::BasicBlockId, function::Function}; + +/// Depth-first traversal stack state marker for computing the cfg post-order. +enum Visit { + First, + Last, +} + +pub(crate) struct PostOrder(Vec); + +impl PostOrder { + pub(crate) fn as_slice(&self) -> &[BasicBlockId] { + self.0.as_slice() + } +} + +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)) + } + + // 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 { + let mut stack = vec![(Visit::First, func.entry_block())]; + let mut visited: HashSet = HashSet::new(); + let mut post_order: Vec = 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 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); + } + } + } + post_order + } +} + +#[cfg(test)] +mod tests { + use crate::ssa_refactor::ir::{ + function::Function, instruction::TerminatorInstruction, map::Id, post_order::PostOrder, + types::Type, + }; + + #[test] + fn single_block() { + let func_id = Id::test_new(0); + let func = Function::new("func".into(), func_id); + let post_order = PostOrder::with_function(&func); + assert_eq!(post_order.0, [func.entry_block()]); + } + + #[test] + fn arb_graph_with_unreachable() { + // A → B C + // ↓ ↗ ↓ ↓ + // D ← E → F + // (`A` is entry block) + // Expected post-order working: + // A { + // B { + // E { + // D { + // B (seen) + // } -> push(D) + // F { + // } -> push(F) + // } -> push(E) + // } -> push(B) + // D (seen) + // } -> push(A) + // Result: + // D, F, E, B, A, (C dropped as unreachable) + + let func_id = Id::test_new(0); + let mut func = Function::new("func".into(), func_id); + let block_a_id = func.entry_block(); + let block_b_id = func.dfg.make_block(); + let block_c_id = func.dfg.make_block(); + let block_d_id = func.dfg.make_block(); + let block_e_id = func.dfg.make_block(); + let block_f_id = func.dfg.make_block(); + + // A → B • + // ↓ + // D • • + let cond_a = func.dfg.add_block_parameter(block_a_id, Type::unsigned(1)); + func.dfg.set_block_terminator( + block_a_id, + TerminatorInstruction::JmpIf { + condition: cond_a, + then_destination: block_b_id, + else_destination: block_d_id, + }, + ); + // • B • + // • ↓ • + // • E • + func.dfg.set_block_terminator( + block_b_id, + TerminatorInstruction::Jmp { destination: block_e_id, arguments: vec![] }, + ); + // • • • + // + // D ← E → F + let cond_e = func.dfg.add_block_parameter(block_e_id, Type::unsigned(1)); + func.dfg.set_block_terminator( + block_e_id, + TerminatorInstruction::JmpIf { + condition: cond_e, + then_destination: block_d_id, + else_destination: block_f_id, + }, + ); + // • B • + // ↗ + // D • • + func.dfg.set_block_terminator( + block_d_id, + TerminatorInstruction::Jmp { destination: block_b_id, arguments: vec![] }, + ); + // • • C + // • • ↓ + // • • F + func.dfg.set_block_terminator( + block_c_id, + TerminatorInstruction::Jmp { destination: block_f_id, arguments: vec![] }, + ); + + let post_order = PostOrder::with_function(&func); + assert_eq!(post_order.0, [block_d_id, block_f_id, block_e_id, block_b_id, block_a_id]); + } +} diff --git a/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/program.rs b/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/program.rs index 99d49456210..de4f01fc613 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/program.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ssa_gen/program.rs @@ -4,7 +4,7 @@ use crate::ssa_refactor::ir::function::Function; /// Contains the entire SSA representation of the program. pub struct Ssa { - functions: Vec, + pub functions: Vec, } impl Ssa {