chore(ssa refactor): Add loop unrolling pass (#1364)

* Start inlining pass

* Get most of pass working

* Finish function inlining pass

* Add basic test

* Address PR comments

* Start block inlining

* Add basic instruction simplification

* Cargo fmt

* Add comments

* Add context object

* Add push_instruction

* Fix bug in inlining pass

* Reorder loop unrolling pass

* Get it working for most loops. Still missing loops with if inside

* Rework entire pass from scratch

* Finish loop unrolling

* Add doc comment

* Fix bad merge

* Add test

* Remove outdated parts of PR

* Correctly handle loops with non-const indices

* Address PR comments

* Fix inlining bug and add a test for loops which fail to unroll

* Update simplify_cfg to use new inline_block method

* Remove now-unneeded test helper function

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
@@ -52,6 +52,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);
@@ -83,6 +89,32 @@ 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")
+    }
+
+    /// Take ownership of this block's terminator, replacing it with an empty return terminator
+    /// so that no clone is needed.
+    ///
+    /// It is expected that this function is used as an optimization on blocks that are no longer
+    /// reachable or will have their terminator overwritten afterwards. Using this on a reachable
+    /// block without setting the terminator afterward will result in the empty return terminator
+    /// being kept, which is likely unwanted.
+    pub(crate) fn take_terminator(&mut self) -> TerminatorInstruction {
+        let terminator = self.terminator.as_mut().expect("Expected block to have a terminator");
+        std::mem::replace(terminator, TerminatorInstruction::Return { return_values: Vec::new() })
+    }
+
     /// 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.
@@ -107,23 +139,4 @@ impl BasicBlock {
             });
         self.instructions.remove(index);
     }
-
-    /// Take ownership of this block's terminator, replacing it with an empty return terminator
-    /// so that no clone is needed.
-    ///
-    /// It is expected that this function is used as an optimization on blocks that are no longer
-    /// reachable or will have their terminator overwritten afterwards. Using this on a reachable
-    /// block without setting the terminator afterward will result in the empty return terminator
-    /// being kept, which is likely unwanted.
-    pub(crate) fn take_terminator(&mut self) -> TerminatorInstruction {
-        let terminator = self.terminator.as_mut().expect("Expected block to have a terminator");
-        std::mem::replace(terminator, TerminatorInstruction::Return { return_values: Vec::new() })
-    }
-
-    /// 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")
-    }
 }

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
@@ -65,7 +68,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
@@ -101,17 +124,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))
             }
         }
@@ -246,16 +271,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> {
@@ -282,6 +297,20 @@ impl DataFlowGraph {
     ) {
         self.blocks[block].set_terminator(terminator);
     }
+
+    /// Moves the entirety of the given block's contents into the destination block.
+    /// The source block afterward will be left in a valid but emptied state. The
+    /// destination block will also have its terminator overwritten with that of the
+    /// source block.
+    pub(crate) fn inline_block(&mut self, source: BasicBlockId, destination: BasicBlockId) {
+        let source = &mut self.blocks[source];
+        let mut instructions = std::mem::take(source.instructions_mut());
+        let terminator = source.take_terminator();
+
+        let destination = &mut self.blocks[destination];
+        destination.instructions_mut().append(&mut instructions);
+        destination.set_terminator(terminator);
+    }
 }
 
 impl std::ops::Index<InstructionId> for DataFlowGraph {