diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir.rs b/crates/noirc_evaluator/src/ssa_refactor/ir.rs index 851b86e511f..1a1ca9eab89 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir.rs @@ -1,4 +1,6 @@ pub(crate) mod basic_block; +pub(crate) mod basic_block_visitors; +pub(crate) mod cfg; pub(crate) mod constant; pub(crate) mod dfg; pub(crate) mod function; diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block_visitors.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block_visitors.rs new file mode 100644 index 00000000000..e0d5dc1b3df --- /dev/null +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/basic_block_visitors.rs @@ -0,0 +1,23 @@ +use super::{ + basic_block::{BasicBlock, BasicBlockId}, + instruction::TerminatorInstruction, +}; + +/// Visit all successors of a block with a given visitor closure. The closure +/// arguments are the branch instruction that is used to reach the successor, +/// and the id of the successor block itself. +pub(crate) fn visit_block_succs(basic_block: &BasicBlock, mut visit: F) { + match basic_block + .terminator() + .expect("ICE: No terminator indicates block is still under construction.") + { + TerminatorInstruction::Jmp { destination, .. } => visit(*destination), + TerminatorInstruction::JmpIf { then_destination, else_destination, .. } => { + visit(*then_destination); + visit(*else_destination); + } + TerminatorInstruction::Return { .. } => { + // The last block of the control flow - no successors + } + } +} diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/cfg.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/cfg.rs new file mode 100644 index 00000000000..05b64e30ed8 --- /dev/null +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/cfg.rs @@ -0,0 +1,251 @@ +use std::collections::{HashMap, HashSet}; + +use super::{ + basic_block::{BasicBlock, BasicBlockId}, + basic_block_visitors, + function::Function, +}; + +/// A container for the successors and predecessors of some Block. +#[derive(Clone, Default)] +struct CfgNode { + /// Set of blocks that containing jumps that target this block. + /// The predecessor set has no meaningful order. + pub(crate) predecessors: HashSet, + + /// Set of blocks that are the targets of jumps in this block. + /// The successors set has no meaningful order. + pub(crate) successors: HashSet, +} + +/// The Control Flow Graph maintains a mapping of blocks to their predecessors +/// and successors where predecessors are basic blocks and successors are +/// basic blocks. +pub(crate) struct ControlFlowGraph { + data: HashMap, +} + +impl ControlFlowGraph { + /// Allocate and compute the control flow graph for `func`. + pub(crate) fn with_function(func: &Function) -> Self { + let mut cfg = ControlFlowGraph { data: HashMap::new() }; + cfg.compute(func); + cfg + } + + fn compute(&mut self, func: &Function) { + for (basic_block_id, basic_block) in func.dfg.basic_blocks_iter() { + self.compute_block(basic_block_id, basic_block); + } + } + + fn compute_block(&mut self, basic_block_id: BasicBlockId, basic_block: &BasicBlock) { + basic_block_visitors::visit_block_succs(basic_block, |dest| { + self.add_edge(basic_block_id, dest); + }); + } + + fn invalidate_block_successors(&mut self, basic_block_id: BasicBlockId) { + let node = self + .data + .get_mut(&basic_block_id) + .expect("ICE: Attempted to invalidate cfg node successors for non-existent node."); + let old_successors = node.successors.clone(); + node.successors.clear(); + for successor_id in old_successors { + self.data + .get_mut(&successor_id) + .expect("ICE: Cfg node successor doesn't exist.") + .predecessors + .remove(&basic_block_id); + } + } + + /// Recompute the control flow graph of `block`. + /// + /// This is for use after modifying instructions within a specific block. It recomputes all edges + /// from `basic_block_id` while leaving edges to `basic_block_id` intact. + pub(crate) fn recompute_block(&mut self, func: &Function, basic_block_id: BasicBlockId) { + self.invalidate_block_successors(basic_block_id); + let basic_block = &func.dfg[basic_block_id]; + self.compute_block(basic_block_id, basic_block); + } + + fn add_edge(&mut self, from: BasicBlockId, to: BasicBlockId) { + let predecessor_node = self.data.entry(from).or_default(); + assert!( + predecessor_node.successors.len() < 2, + "ICE: A cfg node cannot have more than two successors" + ); + predecessor_node.successors.insert(to); + let successor_node = self.data.entry(to).or_default(); + assert!( + successor_node.predecessors.len() < 2, + "ICE: A cfg node cannot have more than two predecessors" + ); + successor_node.predecessors.insert(from); + } + + /// Get an iterator over the CFG predecessors to `basic_block_id`. + pub(crate) fn pred_iter( + &self, + basic_block_id: BasicBlockId, + ) -> impl ExactSizeIterator + '_ { + self.data + .get(&basic_block_id) + .expect("ICE: Attempted to iterate predecessors of block not found within cfg.") + .predecessors + .iter() + .copied() + } + + /// Get an iterator over the CFG successors to `basic_block_id`. + pub(crate) fn succ_iter( + &self, + basic_block_id: BasicBlockId, + ) -> impl ExactSizeIterator + '_ { + self.data + .get(&basic_block_id) + .expect("ICE: Attempted to iterate successors of block not found within cfg.") + .successors + .iter() + .copied() + } +} + +#[cfg(test)] +mod tests { + use crate::ssa_refactor::ir::{instruction::TerminatorInstruction, types::Type}; + + use super::{super::function::Function, ControlFlowGraph}; + + #[test] + fn empty() { + let mut func = Function::new("func".into()); + let block_id = func.entry_block(); + func.dfg[block_id].set_terminator(TerminatorInstruction::Return { return_values: vec![] }); + + ControlFlowGraph::with_function(&func); + } + + #[test] + fn jumps() { + // Build function of form + // fn func { + // block0(cond: u1): + // jmpif cond(), then: block2, else: block1 + // block1(): + // jmpif cond(), then: block1, else: block2 + // block2(): + // return + // } + let mut func = Function::new("func".into()); + let block0_id = func.entry_block(); + let cond = func.dfg.add_block_parameter(block0_id, Type::unsigned(1)); + let block1_id = func.dfg.new_block(); + let block2_id = func.dfg.new_block(); + + func.dfg[block0_id].set_terminator(TerminatorInstruction::JmpIf { + condition: cond, + then_destination: block2_id, + else_destination: block1_id, + arguments: vec![], + }); + func.dfg[block1_id].set_terminator(TerminatorInstruction::JmpIf { + condition: cond, + then_destination: block1_id, + else_destination: block2_id, + arguments: vec![], + }); + func.dfg[block2_id].set_terminator(TerminatorInstruction::Return { return_values: vec![] }); + + let mut cfg = ControlFlowGraph::with_function(&func); + + { + let block0_predecessors = cfg.pred_iter(block0_id).collect::>(); + let block1_predecessors = cfg.pred_iter(block1_id).collect::>(); + let block2_predecessors = cfg.pred_iter(block2_id).collect::>(); + + let block0_successors = cfg.succ_iter(block0_id).collect::>(); + let block1_successors = cfg.succ_iter(block1_id).collect::>(); + let block2_successors = cfg.succ_iter(block2_id).collect::>(); + + assert_eq!(block0_predecessors.len(), 0); + assert_eq!(block1_predecessors.len(), 2); + assert_eq!(block2_predecessors.len(), 2); + + assert_eq!(block1_predecessors.contains(&block0_id), true); + assert_eq!(block1_predecessors.contains(&block1_id), true); + assert_eq!(block2_predecessors.contains(&block0_id), true); + assert_eq!(block2_predecessors.contains(&block1_id), true); + + assert_eq!(block0_successors.len(), 2); + assert_eq!(block1_successors.len(), 2); + assert_eq!(block2_successors.len(), 0); + + assert_eq!(block0_successors.contains(&block1_id), true); + assert_eq!(block0_successors.contains(&block2_id), true); + assert_eq!(block1_successors.contains(&block1_id), true); + assert_eq!(block1_successors.contains(&block2_id), true); + } + + // 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.new_block(); + func.dfg[ret_block_id] + .set_terminator(TerminatorInstruction::Return { return_values: vec![] }); + func.dfg[block2_id].set_terminator(TerminatorInstruction::Jmp { + destination: ret_block_id, + arguments: vec![], + }); + func.dfg[block0_id].set_terminator(TerminatorInstruction::JmpIf { + condition: cond, + then_destination: block1_id, + else_destination: ret_block_id, + arguments: vec![], + }); + + // Recompute new and changed blocks + cfg.recompute_block(&mut func, block0_id); + cfg.recompute_block(&mut func, block2_id); + cfg.recompute_block(&mut func, ret_block_id); + + { + let block0_predecessors = cfg.pred_iter(block0_id).collect::>(); + let block1_predecessors = cfg.pred_iter(block1_id).collect::>(); + let block2_predecessors = cfg.pred_iter(block2_id).collect::>(); + + let block0_successors = cfg.succ_iter(block0_id).collect::>(); + let block1_successors = cfg.succ_iter(block1_id).collect::>(); + let block2_successors = cfg.succ_iter(block2_id).collect::>(); + + assert_eq!(block0_predecessors.len(), 0); + assert_eq!(block1_predecessors.len(), 2); + assert_eq!(block2_predecessors.len(), 1); + + assert_eq!(block1_predecessors.contains(&block0_id), true); + assert_eq!(block1_predecessors.contains(&block1_id), true); + assert_eq!(block2_predecessors.contains(&block0_id), false); + assert_eq!(block2_predecessors.contains(&block1_id), true); + + assert_eq!(block0_successors.len(), 2); + assert_eq!(block1_successors.len(), 2); + assert_eq!(block2_successors.len(), 1); + + assert_eq!(block0_successors.contains(&block1_id), true); + assert_eq!(block0_successors.contains(&ret_block_id), true); + assert_eq!(block1_successors.contains(&block1_id), true); + assert_eq!(block1_successors.contains(&block2_id), true); + assert_eq!(block2_successors.contains(&ret_block_id), true); + } + } +} diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs index f92cae79b75..c21fc2c3f35 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs @@ -95,6 +95,16 @@ impl DataFlowGraph { }) } + /// Get an iterator over references to each basic block within the dfg, paired with the basic + /// block's id. + /// + /// The pairs are order by id, which is not guaranteed to be meaningful. + pub(crate) fn basic_blocks_iter( + &self, + ) -> impl ExactSizeIterator { + self.blocks.iter() + } + pub(crate) fn block_parameters(&self, block: BasicBlockId) -> &[ValueId] { self.blocks[block].parameters() } @@ -237,6 +247,13 @@ impl std::ops::Index for DataFlowGraph { } } +impl std::ops::IndexMut for DataFlowGraph { + /// Get a mutable reference to a function's basic block for the given id. + fn index_mut(&mut self, id: BasicBlockId) -> &mut BasicBlock { + &mut self.blocks[id] + } +} + #[cfg(test)] mod tests { use super::DataFlowGraph; diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/map.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/map.rs index bb526076e3b..5937b374726 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir/map.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/map.rs @@ -106,6 +106,14 @@ impl DenseMap { self.storage.push(f(id)); id } + + /// Gets an iterator to a reference to each element in the dense map paired with its id. + /// + /// The id-element pairs are ordered by the numeric values of the ids. + pub(crate) fn iter(&self) -> impl ExactSizeIterator, &T)> { + let ids_iter = (0..self.storage.len()).into_iter().map(|idx| Id::new(idx)); + ids_iter.zip(self.storage.iter()) + } } impl Default for DenseMap {