chore(ssa refactor): Add loop unrolling pass

crates/noirc_evaluator/src/ssa_refactor.rs

@@ -7,13 +7,13 @@
 //! This module heavily borrows from Cranelift
 #![allow(dead_code)]
 
-use crate::errors::RuntimeError;
+use crate::errors::{RuntimeError, RuntimeErrorKind};
 use acvm::{acir::circuit::Circuit, compiler::transformers::IsOpcodeSupported, Language};
 use noirc_abi::Abi;
 
 use noirc_frontend::monomorphization::ast::Program;
 
-use self::acir_gen::Acir;
+use self::ssa_gen::Ssa;
 
 mod acir_gen;
 mod ir;
@@ -24,9 +24,15 @@ pub mod ssa_gen;
 /// Optimize the given program by converting it into SSA
 /// form and performing optimizations there. When finished,
 /// convert the final SSA into ACIR and return it.
-pub fn optimize_into_acir(program: Program) -> Acir {
-    ssa_gen::generate_ssa(program).inline_functions().into_acir()
+pub fn optimize_into_acir(program: Program) {
+    ssa_gen::generate_ssa(program)
+        .print("Initial SSA:")
+        .inline_functions()
+        .print("After Inlining:")
+        .unroll_loops()
+        .print("After Unrolling:");
 }
+
 /// Compiles the Program into ACIR and applies optimizations to the arithmetic gates
 /// This is analogous to `ssa:create_circuit` and this method is called when one wants
 /// to use the new ssa module to process Noir code.
@@ -37,5 +43,14 @@ pub fn experimental_create_circuit(
     _enable_logging: bool,
     _show_output: bool,
 ) -> Result<(Circuit, Abi), RuntimeError> {
-    todo!("this is a stub function for the new SSA refactor module")
+    optimize_into_acir(_program);
+    let error_kind = RuntimeErrorKind::Spanless("Acir-gen is unimplemented".into());
+    Err(RuntimeError::new(error_kind, None))
+}
+
+impl Ssa {
+    fn print(self, msg: &str) -> Ssa {
+        println!("{msg}\n{self}");
+        self
+    }
 }

crates/noirc_evaluator/src/ssa_refactor/acir_gen/mod.rs

@@ -6,7 +6,7 @@ use super::ssa_gen::Ssa;
 struct Context {}
 
 /// The output of the Acir-gen pass
-pub struct Acir {}
+pub(crate) struct Acir {}
 
 impl Ssa {
     pub(crate) fn into_acir(self) -> Acir {

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

@@ -29,10 +29,10 @@ pub(crate) struct BasicBlock {
 pub(crate) type BasicBlockId = Id<BasicBlock>;
 
 impl BasicBlock {
-    /// Create a new BasicBlock with the given instructions.
+    /// Create a new BasicBlock with the given parameters.
     /// Parameters can also be added later via BasicBlock::add_parameter
-    pub(crate) fn new(instructions: Vec<InstructionId>) -> Self {
-        Self { parameters: Vec::new(), instructions, terminator: None }
+    pub(crate) fn new() -> Self {
+        Self { parameters: Vec::new(), instructions: Vec::new(), terminator: None }
     }
 
     /// Returns the parameters of this block
@@ -47,6 +47,12 @@ impl BasicBlock {
         self.parameters.push(parameter);
     }
 
+    /// Replace this block's current parameters with that of the given Vec.
+    /// This does not perform any checks that any previous parameters were unused.
+    pub(crate) fn set_parameters(&mut self, parameters: Vec<ValueId>) {
+        self.parameters = parameters;
+    }
+
     /// Insert an instruction at the end of this block
     pub(crate) fn insert_instruction(&mut self, instruction: InstructionId) {
         self.instructions.push(instruction);
@@ -78,6 +84,20 @@ impl BasicBlock {
         self.terminator.as_ref()
     }
 
+    /// Returns the terminator of this block, panics if there is None.
+    ///
+    /// Once this block has finished construction, this is expected to always be Some.
+    pub(crate) fn unwrap_terminator(&self) -> &TerminatorInstruction {
+        self.terminator().expect("Expected block to have terminator instruction")
+    }
+
+    /// Returns a mutable reference to the terminator of this block.
+    ///
+    /// Once this block has finished construction, this is expected to always be Some.
+    pub(crate) fn unwrap_terminator_mut(&mut self) -> &mut TerminatorInstruction {
+        self.terminator.as_mut().expect("Expected block to have terminator instruction")
+    }
+
     /// 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.

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

@@ -84,10 +84,6 @@ impl ControlFlowGraph {
         );
         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);
     }
 

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

@@ -1,5 +1,7 @@
 use std::collections::HashMap;
 
+use crate::ssa_refactor::ir::instruction::SimplifyResult;
+
 use super::{
     basic_block::{BasicBlock, BasicBlockId},
     constant::{NumericConstant, NumericConstantId},
@@ -13,6 +15,7 @@ use super::{
 };
 
 use acvm::FieldElement;
+use iter_extended::vecmap;
 
 /// The DataFlowGraph contains most of the actual data in a function including
 /// its blocks, instructions, and values. This struct is largely responsible for
@@ -57,6 +60,8 @@ pub(crate) struct DataFlowGraph {
     signatures: DenseMap<Signature>,
 
     /// All blocks in a function
+    ///
+    /// This map is sparse to allow removing unreachable blocks during optimizations
     blocks: DenseMap<BasicBlock>,
 }
 
@@ -65,7 +70,27 @@ impl DataFlowGraph {
     /// After being created, the block is unreachable in the current function
     /// until another block is made to jump to it.
     pub(crate) fn make_block(&mut self) -> BasicBlockId {
-        self.blocks.insert(BasicBlock::new(Vec::new()))
+        self.blocks.insert(BasicBlock::new())
+    }
+
+    /// Create a new block with the same parameter count and parameter
+    /// types from the given block.
+    /// This is a somewhat niche operation used in loop unrolling but is included
+    /// here as doing it outside the DataFlowGraph would require cloning the parameters.
+    pub(crate) fn make_block_with_parameters_from_block(
+        &mut self,
+        block: BasicBlockId,
+    ) -> BasicBlockId {
+        let new_block = self.make_block();
+        let parameters = self.blocks[block].parameters();
+
+        let parameters = vecmap(parameters.iter().enumerate(), |(position, param)| {
+            let typ = self.values[*param].get_type();
+            self.values.insert(Value::Param { block: new_block, position, typ })
+        });
+
+        self.blocks[new_block].set_parameters(parameters);
+        new_block
     }
 
     /// Get an iterator over references to each basic block within the dfg, paired with the basic
@@ -78,6 +103,11 @@ impl DataFlowGraph {
         self.blocks.iter()
     }
 
+    // Remove all blocks in this DFG that do not satisfy the given predicate
+    // pub(crate) fn retain_blocks(&mut self, mut predicate: impl FnMut(BasicBlockId) -> bool) {
+    //     self.blocks.retain(|id, _| predicate(*id))
+    // }
+
     /// Returns the parameters of the given block
     pub(crate) fn block_parameters(&self, block: BasicBlockId) -> &[ValueId] {
         self.blocks[block].parameters()
@@ -101,17 +131,19 @@ impl DataFlowGraph {
     }
 
     /// Inserts a new instruction at the end of the given block and returns its results
-    pub(crate) fn insert_instruction(
+    pub(crate) fn insert_instruction_and_results(
         &mut self,
         instruction: Instruction,
         block: BasicBlockId,
         ctrl_typevars: Option<Vec<Type>>,
     ) -> InsertInstructionResult {
+        use InsertInstructionResult::*;
         match instruction.simplify(self) {
-            Some(simplification) => InsertInstructionResult::SimplifiedTo(simplification),
-            None => {
+            SimplifyResult::SimplifiedTo(simplification) => SimplifiedTo(simplification),
+            SimplifyResult::Remove => InstructionRemoved,
+            SimplifyResult::None => {
                 let id = self.make_instruction(instruction, ctrl_typevars);
-                self.insert_instruction_in_block(block, id);
+                self.blocks[block].insert_instruction(id);
                 InsertInstructionResult::Results(self.instruction_results(id))
             }
         }
@@ -123,6 +155,15 @@ impl DataFlowGraph {
         self.values.insert(value)
     }
 
+    /// Replaces the value specified by the given ValueId with a new Value.
+    ///
+    /// This is the preferred method to call for optimizations simplifying
+    /// values since other instructions referring to the same ValueId need
+    /// not be modified to refer to a new ValueId.
+    pub(crate) fn set_value(&mut self, value_id: ValueId, new_value: Value) {
+        self.values[value_id] = new_value;
+    }
+
     /// Creates a new constant value, or returns the Id to an existing one if
     /// one already exists.
     pub(crate) fn make_constant(&mut self, value: FieldElement, typ: Type) -> ValueId {
@@ -233,16 +274,6 @@ impl DataFlowGraph {
         parameter
     }
 
-    /// Insert an instruction at the end of a given block.
-    /// If the block already has a terminator, the instruction is inserted before the terminator.
-    pub(crate) fn insert_instruction_in_block(
-        &mut self,
-        block: BasicBlockId,
-        instruction: InstructionId,
-    ) {
-        self.blocks[block].insert_instruction(instruction);
-    }
-
     /// Returns the field element represented by this value if it is a numeric constant.
     /// Returns None if the given value is not a numeric constant.
     pub(crate) fn get_numeric_constant(&self, value: Id<Value>) -> Option<FieldElement> {
@@ -269,6 +300,14 @@ impl DataFlowGraph {
     ) {
         self.blocks[block].set_terminator(terminator);
     }
+
+    /// Sets the terminator instruction for the given basic block
+    pub(crate) fn get_block_terminator_mut(
+        &mut self,
+        block: BasicBlockId,
+    ) -> &mut TerminatorInstruction {
+        self.blocks[block].unwrap_terminator_mut()
+    }
 }
 
 impl std::ops::Index<InstructionId> for DataFlowGraph {

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

@@ -61,6 +61,24 @@ impl Function {
     pub(crate) fn parameters(&self) -> &[ValueId] {
         self.dfg.block_parameters(self.entry_block)
     }
+
+    // Remove any unreachable blocks from this function.
+    // To do this, this method must traverse the cfg of this function in program order.
+    // pub(crate) fn remove_unreachable_blocks(&mut self) {
+    //     let mut reached = HashSet::new();
+    //     let mut stack: Vec<BasicBlockId> = vec![self.entry_block];
+
+    //     while let Some(block) = stack.pop() {
+    //         reached.insert(block);
+    //         for block in self.dfg[block].successors() {
+    //             if !reached.contains(&block) {
+    //                 stack.push(block);
+    //             }
+    //         }
+    //     }
+
+    //     self.dfg.retain_blocks(|block| reached.contains(&block));
+    // }
 }
 
 /// FunctionId is a reference for a function