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

* Implement first-class functions

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

Co-authored-by: kevaundray <[email protected]>

* Implement intrinsics

* Fix no return & duplicate main

* bad git. remove duplicated functions

* Remove Option<Function> in builder

* Undo debug printing in driver

---------

Co-authored-by: kevaundray <[email protected]>

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

@@ -8,7 +8,7 @@ use crate::ssa_refactor::ir::{
     value::{Value, ValueId},
 };
 
-use super::ir::instruction::Intrinsic;
+use super::{ir::instruction::Intrinsic, ssa_gen::Ssa};
 
 /// The per-function context for each ssa function being generated.
 ///
@@ -20,7 +20,7 @@ use super::ir::instruction::Intrinsic;
 pub(crate) struct FunctionBuilder {
     current_function: Function,
     current_block: BasicBlockId,
-    finished_functions: Vec<(FunctionId, Function)>,
+    finished_functions: Vec<Function>,
 }
 
 impl FunctionBuilder {
@@ -34,14 +34,15 @@ impl FunctionBuilder {
     /// Finish the current function and create a new function
     pub(crate) fn new_function(&mut self, name: String, function_id: FunctionId) {
         let new_function = Function::new(name, function_id);
-        let old_function = std::mem::replace(&mut self.current_function, new_function);
+        self.current_block = new_function.entry_block();
 
-        self.finished_functions.push((self.current_function.id(), old_function));
+        let old_function = std::mem::replace(&mut self.current_function, new_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));
-        self.finished_functions
+    pub(crate) fn finish(mut self) -> Ssa {
+        self.finished_functions.push(self.current_function);
+        Ssa::new(self.finished_functions)
     }
 
     pub(crate) fn add_parameter(&mut self, typ: Type) -> ValueId {

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

@@ -37,18 +37,17 @@ 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 function_id = shared_context
+            .pop_next_function_in_queue()
+            .expect("No function in queue for the FunctionContext to compile")
+            .1;
 
-        let mut this = Self {
-            definitions: HashMap::new(),
-            builder: FunctionBuilder::new(function_name, new_id),
-            shared_context,
-        };
+        let builder = FunctionBuilder::new(function_name, function_id);
+        let mut this = Self { definitions: HashMap::new(), builder, shared_context };
         this.add_parameters_to_scope(parameters);
         this
     }

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

@@ -1,6 +1,9 @@
 mod context;
+mod program;
 mod value;
 
+pub use program::Ssa;
+
 use context::SharedContext;
 use iter_extended::vecmap;
 use noirc_errors::Location;
@@ -13,24 +16,36 @@ use self::{
 
 use super::ir::{instruction::BinaryOp, types::Type, value::ValueId};
 
-pub(crate) fn generate_ssa(program: Program) {
+pub fn generate_ssa(program: Program) -> Ssa {
     let context = SharedContext::new(program);
 
-    let main = context.program.main();
     let main_id = Program::main_id();
-    let main_name = main.name.clone();
+    let main = context.program.main();
+
+    // Queue the main function for compilation
+    context.get_or_queue_function(main_id);
 
-    let mut function_context = FunctionContext::new(main_id, main_name, &main.parameters, &context);
-    function_context.codegen_expression(&main.body);
+    let mut function_context = FunctionContext::new(main.name.clone(), &main.parameters, &context);
+    function_context.codegen_function_body(&main.body);
 
     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),

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

@@ -0,0 +1,23 @@
+use std::fmt::Display;
+
+use crate::ssa_refactor::ir::function::Function;
+
+/// Contains the entire Ssa representation of the program
+pub struct Ssa {
+    functions: Vec<Function>,
+}
+
+impl Ssa {
+    pub fn new(functions: Vec<Function>) -> Self {
+        Self { functions }
+    }
+}
+
+impl Display for Ssa {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        for function in &self.functions {
+            writeln!(f, "{function}")?;
+        }
+        Ok(())
+    }
+}