chore(ssa): Add intial control flow graph

crates/noirc_evaluator/src/ssa_refactor.rs

@@ -8,6 +8,8 @@
 #[allow(dead_code)]
 mod basic_block;
 #[allow(dead_code)]
+mod cfg;
+#[allow(dead_code)]
 mod dfg;
 #[allow(dead_code)]
 mod ir;

crates/noirc_evaluator/src/ssa_refactor/basic_block.rs

@@ -25,7 +25,13 @@ pub(crate) struct BasicBlock {
     terminator: TerminatorInstruction,
 }
 
-#[derive(Debug, PartialEq, Eq, Hash, Clone)]
+impl BasicBlock {
+    pub(crate) fn terminator(&self) -> &TerminatorInstruction {
+        &self.terminator
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
 /// An identifier for a Basic Block.
 pub(crate) struct BasicBlockId;
 

crates/noirc_evaluator/src/ssa_refactor/cfg.rs

@@ -0,0 +1,110 @@
+use std::collections::{hash_set, HashMap, HashSet};
+
+use super::{
+    basic_block::BasicBlockId,
+    ir::{function::Function, instruction_predicates},
+};
+
+/// 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<BasicBlockId>,
+
+    /// Set of blocks that are the targets of jumps in this block.
+    /// The successors set has no meaningful order.
+    pub(crate) successors: HashSet<BasicBlockId>,
+}
+
+/// 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<BasicBlockId, CfgNode>,
+}
+
+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 in func.basic_block_ids_iter() {
+            self.compute_block(func, basic_block_id);
+        }
+    }
+
+    fn compute_block(&mut self, func: &Function, basic_block_id: BasicBlockId) {
+        instruction_predicates::visit_block_succs(func, basic_block_id, |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);
+        self.compute_block(func, basic_block_id);
+    }
+
+    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 `block`.
+    pub(crate) fn pred_iter(&self, basic_block_id: BasicBlockId) -> PredIter {
+        self.data
+            .get(&basic_block_id)
+            .expect("ICE: Attempted to iterate predecessors of block not found within cfg.")
+            .predecessors
+            .iter()
+    }
+
+    /// Get an iterator over the CFG successors to `block`.
+    pub(crate) fn succ_iter(&self, basic_block_id: BasicBlockId) -> SuccIter {
+        self.data
+            .get(&basic_block_id)
+            .expect("ICE: Attempted to iterate successors of block not found within cfg.")
+            .successors
+            .iter()
+    }
+}
+
+/// An iterator over block predecessors. The iterator type is `BasicBlockId`.
+pub(crate) type PredIter<'a> = hash_set::Iter<'a, BasicBlockId>;
+
+/// An iterator over block successors. The iterator type is `BasicBlockId`.
+pub(crate) type SuccIter<'a> = hash_set::Iter<'a, BasicBlockId>;

crates/noirc_evaluator/src/ssa_refactor/ir.rs

@@ -1,6 +1,7 @@
 pub(crate) mod extfunc;
-mod function;
+pub(crate) mod function;
 pub(crate) mod instruction;
+pub(crate) mod instruction_predicates;
 pub(crate) mod map;
 pub(crate) mod types;
 pub(crate) mod value;

crates/noirc_evaluator/src/ssa_refactor/ir/function.rs

@@ -3,7 +3,10 @@ use crate::ssa_refactor::basic_block::{BasicBlock, BasicBlockId};
 use super::instruction::Instruction;
 
 use noirc_errors::Location;
-use std::collections::HashMap;
+use std::{
+    collections::{hash_map, HashMap},
+    ops::Index,
+};
 
 /// A function holds a list of instructions.
 /// These instructions are further grouped into
@@ -20,6 +23,34 @@ pub(crate) struct Function {
     entry_block: BasicBlockId,
 }
 
+impl Function {
+    /// Get an iterator over the ids of the basic blocks within the function.
+    ///
+    /// The ids are iterated in no meaningful order.
+    pub(crate) fn basic_block_ids_iter(&self) -> BasicBlockIdsIter {
+        BasicBlockIdsIter(self.basic_blocks.keys())
+    }
+}
+
+impl Index<BasicBlockId> for Function {
+    type Output = BasicBlock;
+    /// Get a function's basic block for the given id.
+    fn index<'a>(&'a self, id: BasicBlockId) -> &'a BasicBlock {
+        &self.basic_blocks[&id]
+    }
+}
+
+/// An iterator over a function's basic block ids. The iterator type is `BasicBlockId`.
+pub(crate) struct BasicBlockIdsIter<'a>(hash_map::Keys<'a, BasicBlockId, BasicBlock>);
+
+impl<'a> Iterator for BasicBlockIdsIter<'a> {
+    type Item = BasicBlockId;
+
+    fn next(&mut self) -> Option<BasicBlockId> {
+        self.0.next().map(|k| *k)
+    }
+}
+
 /// FunctionId is a reference for a function
 #[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
 pub(crate) struct FunctionId(pub(crate) u32);

crates/noirc_evaluator/src/ssa_refactor/ir/instruction_predicates.rs

@@ -0,0 +1,21 @@
+use super::{
+    super::basic_block::BasicBlockId, function::Function, 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<F: FnMut(BasicBlockId)>(
+    func: &Function,
+    basic_block_id: BasicBlockId,
+    mut visit: F,
+) {
+    let terminator = &func[basic_block_id].terminator();
+    match terminator {
+        TerminatorInstruction::Jmp { destination, .. } => visit(*destination),
+        TerminatorInstruction::JmpIf { then_destination, else_destination, .. } => {
+            visit(*then_destination);
+            visit(*else_destination);
+        }
+    }
+}