chore(ssa refactor): Fix no returns & duplicate main

crates/noirc_evaluator/src/ssa_refactor.rs

@@ -9,4 +9,4 @@
 
 mod ir;
 mod ssa_builder;
-mod ssa_gen;
+pub mod ssa_gen;

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

@@ -15,7 +15,7 @@ use noirc_errors::Location;
 /// To reference external functions, one must first import the function signature
 /// into the current function's context.
 #[derive(Debug)]
-pub(crate) struct Function {
+pub struct Function {
     /// Maps instructions to source locations
     source_locations: HashMap<InstructionId, Location>,
 

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

@@ -46,7 +46,7 @@ pub(crate) fn display_block(
 ) -> Result {
     let block = &function.dfg[block_id];
 
-    writeln!(f, "  {}({}):", block_id, value_list(function, block.parameters()))?;
+    writeln!(f, "  {}({}):", block_id, value_list_with_types(function, block.parameters()))?;
 
     for instruction in block.instructions() {
         display_instruction(function, *instruction, f)?;
@@ -70,6 +70,16 @@ fn value(function: &Function, id: ValueId) -> String {
     }
 }
 
+/// Display each value along with its type. E.g. `v0: Field, v1: u64, v2: u1`
+fn value_list_with_types(function: &Function, values: &[ValueId]) -> String {
+    vecmap(values, |id| {
+        let value = value(function, *id);
+        let typ = function.dfg.type_of_value(*id);
+        format!("{value}: {typ}")
+    })
+    .join(", ")
+}
+
 fn value_list(function: &Function, values: &[ValueId]) -> String {
     vecmap(values, |id| value(function, *id)).join(", ")
 }
@@ -87,7 +97,9 @@ pub(crate) fn display_terminator(
             writeln!(
                 f,
                 "    jmpif {} then: {}, else: {}",
-                condition, then_destination, else_destination
+                value(function, *condition),
+                then_destination,
+                else_destination
             )
         }
         Some(TerminatorInstruction::Return { return_values }) => {

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

@@ -17,36 +17,48 @@ use super::ir::instruction::Intrinsic;
 ///
 /// Contrary to the name, this struct has the capacity to build as many
 /// functions as needed, although it is limited to one function at a time.
+#[derive(Default)]
 pub(crate) struct FunctionBuilder {
-    current_function: Function,
-    current_block: BasicBlockId,
-    finished_functions: Vec<(FunctionId, Function)>,
+    current_function: Option<Function>,
+    current_block: Option<BasicBlockId>,
+    finished_functions: Vec<Function>,
 }
 
 impl FunctionBuilder {
-    pub(crate) fn new(function_name: String, function_id: FunctionId) -> Self {
-        let new_function = Function::new(function_name, function_id);
-        let current_block = new_function.entry_block();
-
-        Self { current_function: new_function, current_block, finished_functions: Vec::new() }
-    }
-
     /// Finish the current function and create a new function
     pub(crate) fn new_function(&mut self, name: String, function_id: FunctionId) {
+        let old_function = std::mem::take(&mut self.current_function);
         let new_function = Function::new(name, function_id);
-        let old_function = std::mem::replace(&mut self.current_function, new_function);
+        self.current_block = Some(new_function.entry_block());
+        self.current_function = Some(new_function);
 
-        self.finished_functions.push((self.current_function.id(), old_function));
+        if let Some(old_function) = old_function {
+            self.finished_functions.push(old_function);
+        }
     }
 
-    pub(crate) fn finish(mut self) -> Vec<(FunctionId, Function)> {
-        self.finished_functions.push((self.current_function.id(), self.current_function));
+    pub(crate) fn finish(mut self) -> Vec<Function> {
+        if let Some(old_function) = self.current_function {
+            self.finished_functions.push(old_function);
+        }
         self.finished_functions
     }
 
+    pub(crate) fn current_function(&self) -> &Function {
+        self.current_function.as_ref().expect("This builder is not currently building a Function")
+    }
+
+    pub(crate) fn current_function_mut(&mut self) -> &mut Function {
+        self.current_function.as_mut().expect("This builder is not currently building a Function")
+    }
+
+    pub(crate) fn current_block(&self) -> BasicBlockId {
+        self.current_block.expect("This builder has no current_block set")
+    }
+
     pub(crate) fn add_parameter(&mut self, typ: Type) -> ValueId {
-        let entry = self.current_function.entry_block();
-        self.current_function.dfg.add_block_parameter(entry, typ)
+        let entry = self.current_function().entry_block();
+        self.current_function_mut().dfg.add_block_parameter(entry, typ)
     }
 
     /// Insert a numeric constant into the current function
@@ -55,7 +67,7 @@ impl FunctionBuilder {
         value: impl Into<FieldElement>,
         typ: Type,
     ) -> ValueId {
-        self.current_function.dfg.make_constant(value.into(), typ)
+        self.current_function_mut().dfg.make_constant(value.into(), typ)
     }
 
     /// Insert a numeric constant into the current function of type Field
@@ -64,15 +76,15 @@ impl FunctionBuilder {
     }
 
     pub(crate) fn type_of_value(&self, value: ValueId) -> Type {
-        self.current_function.dfg.type_of_value(value)
+        self.current_function().dfg.type_of_value(value)
     }
 
     pub(crate) fn insert_block(&mut self) -> BasicBlockId {
-        self.current_function.dfg.make_block()
+        self.current_function_mut().dfg.make_block()
     }
 
     pub(crate) fn add_block_parameter(&mut self, block: BasicBlockId, typ: Type) -> ValueId {
-        self.current_function.dfg.add_block_parameter(block, typ)
+        self.current_function_mut().dfg.add_block_parameter(block, typ)
     }
 
     /// Inserts a new instruction at the end of the current block and returns its results
@@ -81,16 +93,17 @@ impl FunctionBuilder {
         instruction: Instruction,
         ctrl_typevars: Option<Vec<Type>>,
     ) -> &[ValueId] {
-        let id = self.current_function.dfg.make_instruction(instruction, ctrl_typevars);
-        self.current_function.dfg.insert_instruction_in_block(self.current_block, id);
-        self.current_function.dfg.instruction_results(id)
+        let id = self.current_function_mut().dfg.make_instruction(instruction, ctrl_typevars);
+        let current_block = self.current_block();
+        self.current_function_mut().dfg.insert_instruction_in_block(current_block, id);
+        self.current_function_mut().dfg.instruction_results(id)
     }
 
     /// Switch to inserting instructions in the given block.
     /// Expects the given block to be within the same function. If you want to insert
     /// instructions into a new function, call new_function instead.
     pub(crate) fn switch_to_block(&mut self, block: BasicBlockId) {
-        self.current_block = block;
+        self.current_block = Some(block);
     }
 
     /// Insert an allocate instruction at the end of the current block, allocating the
@@ -113,7 +126,7 @@ impl FunctionBuilder {
         offset: ValueId,
         type_to_load: Type,
     ) -> ValueId {
-        if let Some(offset) = self.current_function.dfg.get_numeric_constant(offset) {
+        if let Some(offset) = self.current_function().dfg.get_numeric_constant(offset) {
             if !offset.is_zero() {
                 let offset = self.field_constant(offset);
                 address = self.insert_binary(address, BinaryOp::Add, offset);
@@ -171,7 +184,8 @@ impl FunctionBuilder {
 
     /// Terminates the current block with the given terminator instruction
     fn terminate_block_with(&mut self, terminator: TerminatorInstruction) {
-        self.current_function.dfg.set_block_terminator(self.current_block, terminator);
+        let current_block = self.current_block();
+        self.current_function_mut().dfg.set_block_terminator(current_block, terminator);
     }
 
     /// Terminate the current block with a jmp instruction to jmp to the given
@@ -212,9 +226,9 @@ impl FunctionBuilder {
     /// (which should always be present to load a mutable value) with a store of the
     /// assigned value.
     pub(crate) fn mutate_load_into_store(&mut self, load_result: ValueId, value_to_store: ValueId) {
-        let (instruction, address) = match &self.current_function.dfg[load_result] {
+        let (instruction, address) = match &self.current_function().dfg[load_result] {
             Value::Instruction { instruction, .. } => {
-                match &self.current_function.dfg[*instruction] {
+                match &self.current_function().dfg[*instruction] {
                     Instruction::Load { address } => (*instruction, *address),
                     other => {
                         panic!("mutate_load_into_store: Expected Load instruction, found {other:?}")
@@ -225,21 +239,21 @@ impl FunctionBuilder {
         };
 
         let store = Instruction::Store { address, value: value_to_store };
-        self.current_function.dfg.replace_instruction(instruction, store);
+        self.current_function_mut().dfg.replace_instruction(instruction, store);
         // Clear the results of the previous load for safety
-        self.current_function.dfg.make_instruction_results(instruction, None);
+        self.current_function_mut().dfg.make_instruction_results(instruction, None);
     }
 
     /// Returns a ValueId pointing to the given function or imports the function
     /// into the current function if it was not already, and returns that ID.
     pub(crate) fn import_function(&mut self, function: FunctionId) -> ValueId {
-        self.current_function.dfg.import_function(function)
+        self.current_function_mut().dfg.import_function(function)
     }
 
     /// Retrieve a value reference to the given intrinsic operation.
     /// Returns None if there is no intrinsic matching the given name.
     pub(crate) fn import_intrinsic(&mut self, name: &str) -> Option<ValueId> {
         Intrinsic::lookup(name)
-            .map(|intrinsic| self.current_function.dfg.import_intrinsic(intrinsic))
+            .map(|intrinsic| self.current_function_mut().dfg.import_intrinsic(intrinsic))
     }
 }

crates/noirc_evaluator/src/ssa_refactor/ssa_gen/context.rs

@@ -36,21 +36,8 @@ pub(super) struct SharedContext {
 }
 
 impl<'a> FunctionContext<'a> {
-    pub(super) fn new(
-        function_id: FuncId,
-        function_name: String,
-        parameters: &Parameters,
-        shared_context: &'a SharedContext,
-    ) -> Self {
-        let new_id = shared_context.get_or_queue_function(function_id);
-
-        let mut this = Self {
-            definitions: HashMap::new(),
-            builder: FunctionBuilder::new(function_name, new_id),
-            shared_context,
-        };
-        this.add_parameters_to_scope(parameters);
-        this
+    pub(super) fn new(shared_context: &'a SharedContext) -> Self {
+        Self { definitions: HashMap::new(), builder: FunctionBuilder::default(), shared_context }
     }
 
     pub(super) fn new_function(&mut self, id: IrFunctionId, name: String, parameters: &Parameters) {

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

@@ -11,26 +11,35 @@ use self::{
     value::{Tree, Values},
 };
 
-use super::ir::{instruction::BinaryOp, types::Type, value::ValueId};
+use super::ir::{function::Function, instruction::BinaryOp, types::Type, value::ValueId};
 
-pub(crate) fn generate_ssa(program: Program) {
+pub fn generate_ssa(program: Program) -> Vec<Function> {
     let context = SharedContext::new(program);
 
-    let main = context.program.main();
     let main_id = Program::main_id();
-    let main_name = main.name.clone();
 
-    let mut function_context = FunctionContext::new(main_id, main_name, &main.parameters, &context);
-    function_context.codegen_expression(&main.body);
+    // Queue the main function for compilation
+    context.get_or_queue_function(main_id);
 
+    let mut function_context = FunctionContext::new(&context);
     while let Some((src_function_id, dest_id)) = context.pop_next_function_in_queue() {
         let function = &context.program[src_function_id];
         function_context.new_function(dest_id, function.name.clone(), &function.parameters);
-        function_context.codegen_expression(&function.body);
+        function_context.codegen_function_body(&function.body);
     }
+
+    function_context.builder.finish()
 }
 
 impl<'a> FunctionContext<'a> {
+    /// Codegen a function's body and set its return value to that of its last parameter.
+    /// For functions returning nothing, this will be an empty list.
+    fn codegen_function_body(&mut self, body: &Expression) {
+        let return_value = self.codegen_expression(body);
+        let results = return_value.into_value_list(self);
+        self.builder.terminate_with_return(results);
+    }
+
     fn codegen_expression(&mut self, expr: &Expression) -> Values {
         match expr {
             Expression::Ident(ident) => self.codegen_ident(ident),