chore(ssa refactor): Handle function parameters (#1203)

* Add Context structs and start ssa gen pass

* Fix block arguments

* Fix clippy lint

* Use the correct dfg

* Rename contexts to highlight the inner contexts are shared rather than used directly

* Correctly handle function parameters

* Rename Nested to Tree; add comment

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

@@ -35,7 +35,24 @@ pub(crate) struct BasicBlock {
 pub(crate) type BasicBlockId = Id<BasicBlock>;
 
 impl BasicBlock {
-    pub(super) fn new(parameters: Vec<ValueId>) -> Self {
+    pub(crate) fn new(parameters: Vec<ValueId>) -> Self {
         Self { parameters, instructions: Vec::new(), is_sealed: false, terminator: None }
     }
+
+    pub(crate) fn parameters(&self) -> &[ValueId] {
+        &self.parameters
+    }
+
+    pub(crate) fn add_parameter(&mut self, parameter: ValueId) {
+        self.parameters.push(parameter);
+    }
+
+    /// Insert an instruction at the end of this block
+    pub(crate) fn insert_instruction(&mut self, instruction: InstructionId) {
+        self.instructions.push(instruction);
+    }
+
+    pub(crate) fn set_terminator(&mut self, terminator: TerminatorInstruction) {
+        self.terminator = Some(terminator);
+    }
 }

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

@@ -7,6 +7,8 @@ use super::{
     value::{Value, ValueId},
 };
 
+use iter_extended::vecmap;
+
 #[derive(Debug, Default)]
 /// A convenience wrapper to store `Value`s.
 pub(crate) struct ValueList(Vec<Id<Value>>);
@@ -61,9 +63,31 @@ pub(crate) struct DataFlowGraph {
 }
 
 impl DataFlowGraph {
-    /// Creates a new `empty` basic block
+    /// Creates a new basic block with no parameters.
+    /// After being created, the block is unreachable in the current function
+    /// until another block is made to jump to it.
     pub(crate) fn new_block(&mut self) -> BasicBlockId {
-        todo!()
+        self.blocks.insert(BasicBlock::new(Vec::new()))
+    }
+
+    /// Creates a new basic block with the given parameters.
+    /// After being created, the block is unreachable in the current function
+    /// until another block is made to jump to it.
+    pub(crate) fn new_block_with_parameters(
+        &mut self,
+        parameter_types: impl Iterator<Item = Type>,
+    ) -> BasicBlockId {
+        self.blocks.insert_with_id(|entry_block| {
+            let parameters = vecmap(parameter_types.enumerate(), |(position, typ)| {
+                self.values.insert(Value::Param { block: entry_block, position, typ })
+            });
+
+            BasicBlock::new(parameters)
+        })
+    }
+
+    pub(crate) fn block_parameters(&self, block: BasicBlockId) -> &[ValueId] {
+        self.blocks[block].parameters()
     }
 
     /// Inserts a new instruction into the DFG.
@@ -149,6 +173,14 @@ impl DataFlowGraph {
     pub(crate) fn instruction_results(&self, instruction_id: InstructionId) -> &[ValueId] {
         self.results.get(&instruction_id).expect("expected a list of Values").as_slice()
     }
+
+    pub(crate) fn add_block_parameter(&mut self, block_id: BasicBlockId, typ: Type) -> Id<Value> {
+        let block = &mut self.blocks[block_id];
+        let position = block.parameters().len();
+        let parameter = self.values.insert(Value::Param { block: block_id, position, typ });
+        block.add_parameter(parameter);
+        parameter
+    }
 }
 
 #[cfg(test)]

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

@@ -1,11 +1,9 @@
-use super::basic_block::{BasicBlock, BasicBlockId};
+use super::basic_block::BasicBlockId;
 use super::dfg::DataFlowGraph;
 use super::instruction::Instruction;
-use super::map::{DenseMap, Id, SecondaryMap};
+use super::map::{Id, SecondaryMap};
 use super::types::Type;
-use super::value::Value;
 
-use iter_extended::vecmap;
 use noirc_errors::Location;
 
 /// A function holds a list of instructions.
@@ -16,35 +14,23 @@ use noirc_errors::Location;
 /// into the current function's context.
 #[derive(Debug)]
 pub(crate) struct Function {
-    /// Basic blocks associated to this particular function
-    basic_blocks: DenseMap<BasicBlock>,
-
     /// Maps instructions to source locations
     source_locations: SecondaryMap<Instruction, Location>,
 
     /// The first basic block in the function
     entry_block: BasicBlockId,
 
-    dfg: DataFlowGraph,
+    pub(crate) dfg: DataFlowGraph,
 }
 
 impl Function {
-    pub(crate) fn new(parameter_count: usize) -> Self {
+    /// Creates a new function with an automatically inserted entry block.
+    ///
+    /// Note that any parameters to the function must be manually added later.
+    pub(crate) fn new() -> Self {
         let mut dfg = DataFlowGraph::default();
-        let mut basic_blocks = DenseMap::default();
-
-        // The parameters for each function are stored as the block parameters
-        // of the function's entry block
-        let entry_block = basic_blocks.insert_with_id(|entry_block| {
-            // TODO: Give each parameter its correct type
-            let parameters = vecmap(0..parameter_count, |i| {
-                dfg.make_value(Value::Param { block: entry_block, position: i, typ: Type::Unit })
-            });
-
-            BasicBlock::new(parameters)
-        });
-
-        Self { basic_blocks, source_locations: SecondaryMap::new(), entry_block, dfg }
+        let entry_block = dfg.new_block();
+        Self { source_locations: SecondaryMap::new(), entry_block, dfg }
     }
 
     pub(crate) fn entry_block(&self) -> BasicBlockId {

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

@@ -129,10 +129,20 @@ pub(crate) enum TerminatorInstruction {
         else_destination: BasicBlockId,
         arguments: Vec<ValueId>,
     },
+
     /// Unconditional Jump
     ///
     /// Jumps to specified `destination` with `arguments`
     Jmp { destination: BasicBlockId, arguments: Vec<ValueId> },
+
+    /// Return from the current function with the given return values.
+    ///
+    /// All finished functions should have exactly 1 return instruction.
+    /// Functions with early returns should instead be structured to
+    /// unconditionally jump to a single exit block with the return values
+    /// as the block arguments. Then the exit block can terminate in a return
+    /// instruction returning these values.
+    Return { return_values: Vec<ValueId> },
 }
 
 /// A binary instruction in the IR.

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

@@ -18,6 +18,27 @@ pub(crate) enum NumericType {
 pub(crate) enum Type {
     /// Represents numeric types in the IR, including field elements
     Numeric(NumericType),
+
+    /// A reference to some value, such as an array
+    Reference,
+
+    /// A function that may be called directly
+    Function,
+
     /// The Unit type with a single value
     Unit,
 }
+
+impl Type {
+    pub(crate) fn signed(bit_size: u32) -> Type {
+        Type::Numeric(NumericType::Signed { bit_size })
+    }
+
+    pub(crate) fn unsigned(bit_size: u32) -> Type {
+        Type::Numeric(NumericType::Unsigned { bit_size })
+    }
+
+    pub(crate) fn field() -> Type {
+        Type::Numeric(NumericType::NativeField)
+    }
+}

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

@@ -1,6 +1,8 @@
 use crate::ssa_refactor::ir::{
     basic_block::BasicBlockId,
     function::{Function, FunctionId},
+    types::Type,
+    value::ValueId,
 };
 
 use super::SharedBuilderContext;
@@ -24,8 +26,8 @@ pub(crate) struct FunctionBuilder<'ssa> {
 }
 
 impl<'ssa> FunctionBuilder<'ssa> {
-    pub(crate) fn new(parameters: usize, context: &'ssa SharedBuilderContext) -> Self {
-        let new_function = Function::new(parameters);
+    pub(crate) fn new(context: &'ssa SharedBuilderContext) -> Self {
+        let new_function = Function::new();
         let current_block = new_function.entry_block();
 
         Self {
@@ -38,17 +40,21 @@ impl<'ssa> FunctionBuilder<'ssa> {
     }
 
     /// Finish the current function and create a new function
-    pub(crate) fn new_function(&mut self, parameters: usize) {
-        let new_function = Function::new(parameters);
+    pub(crate) fn new_function(&mut self) {
+        let new_function = Function::new();
         let old_function = std::mem::replace(&mut self.current_function, new_function);
 
         self.finished_functions.push((self.current_function_id, old_function));
-
         self.current_function_id = self.global_context.next_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 add_parameter(&mut self, typ: Type) -> ValueId {
+        let entry = self.current_function.entry_block();
+        self.current_function.dfg.add_block_parameter(entry, typ)
+    }
 }

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

@@ -1,21 +1,24 @@
 use std::collections::HashMap;
 use std::sync::{Mutex, RwLock};
 
-use noirc_frontend::monomorphization::ast::{self, LocalId};
+use iter_extended::vecmap;
+use noirc_frontend::monomorphization::ast::{self, LocalId, Parameters};
 use noirc_frontend::monomorphization::ast::{FuncId, Program};
+use noirc_frontend::Signedness;
 
+use crate::ssa_refactor::ir::types::Type;
 use crate::ssa_refactor::ssa_builder::SharedBuilderContext;
 use crate::ssa_refactor::{
     ir::function::FunctionId as IrFunctionId, ssa_builder::function_builder::FunctionBuilder,
 };
 
-use super::value::Value;
+use super::value::{Tree, Values};
 
 // TODO: Make this a threadsafe queue so we can compile functions in parallel
 type FunctionQueue = Vec<(ast::FuncId, IrFunctionId)>;
 
 pub(super) struct FunctionContext<'a> {
-    definitions: HashMap<LocalId, Value>,
+    definitions: HashMap<LocalId, Values>,
     function_builder: FunctionBuilder<'a>,
     shared_context: &'a SharedContext,
 }
@@ -29,22 +32,90 @@ pub(super) struct SharedContext {
 
 impl<'a> FunctionContext<'a> {
     pub(super) fn new(
-        parameter_count: usize,
+        parameters: &Parameters,
         shared_context: &'a SharedContext,
         shared_builder_context: &'a SharedBuilderContext,
     ) -> Self {
-        Self {
+        let mut this = Self {
             definitions: HashMap::new(),
-            function_builder: FunctionBuilder::new(parameter_count, shared_builder_context),
+            function_builder: FunctionBuilder::new(shared_builder_context),
             shared_context,
+        };
+        this.add_parameters_to_scope(parameters);
+        this
+    }
+
+    pub(super) fn new_function(&mut self, parameters: &Parameters) {
+        self.definitions.clear();
+        self.function_builder.new_function();
+        self.add_parameters_to_scope(parameters);
+    }
+
+    /// Add each parameter to the current scope, and return the list of parameter types.
+    ///
+    /// The returned parameter type list will be flattened, so any struct parameters will
+    /// be returned as one entry for each field (recursively).
+    fn add_parameters_to_scope(&mut self, parameters: &Parameters) {
+        for (id, _, _, typ) in parameters {
+            self.add_parameter_to_scope(*id, typ);
+        }
+    }
+
+    /// Adds a "single" parameter to scope.
+    ///
+    /// Single is in quotes here because in the case of tuple parameters, the tuple is flattened
+    /// into a new parameter for each field recursively.
+    fn add_parameter_to_scope(&mut self, parameter_id: LocalId, parameter_type: &ast::Type) {
+        // Add a separate parameter for each field type in 'parameter_type'
+        let parameter_value = self
+            .map_type(parameter_type, |this, typ| this.function_builder.add_parameter(typ).into());
+
+        self.definitions.insert(parameter_id, parameter_value);
+    }
+
+    /// Maps the given type to a Tree of the result type.
+    ///
+    /// This can be used to (for example) flatten a tuple type, creating
+    /// and returning a new parameter for each field type.
+    pub(super) fn map_type<T>(
+        &mut self,
+        typ: &ast::Type,
+        mut f: impl FnMut(&mut Self, Type) -> T,
+    ) -> Tree<T> {
+        self.map_type_helper(typ, &mut f)
+    }
+
+    // This helper is needed because we need to take f by mutable reference,
+    // otherwise we cannot move it multiple times each loop of vecmap.
+    fn map_type_helper<T>(
+        &mut self,
+        typ: &ast::Type,
+        f: &mut impl FnMut(&mut Self, Type) -> T,
+    ) -> Tree<T> {
+        match typ {
+            ast::Type::Tuple(fields) => {
+                Tree::Branch(vecmap(fields, |field| self.map_type_helper(field, f)))
+            }
+            other => Tree::Leaf(f(self, Self::convert_non_tuple_type(other))),
         }
     }
 
-    pub(super) fn new_function(&mut self, parameters: impl ExactSizeIterator<Item = LocalId>) {
-        self.function_builder.new_function(parameters.len());
+    pub(super) fn convert_non_tuple_type(typ: &ast::Type) -> Type {
+        match typ {
+            ast::Type::Field => Type::field(),
+            ast::Type::Array(_, _) => Type::Reference,
+            ast::Type::Integer(Signedness::Signed, bits) => Type::signed(*bits),
+            ast::Type::Integer(Signedness::Unsigned, bits) => Type::unsigned(*bits),
+            ast::Type::Bool => Type::unsigned(1),
+            ast::Type::String(_) => Type::Reference,
+            ast::Type::Unit => Type::Unit,
+            ast::Type::Tuple(_) => panic!("convert_non_tuple_type called on a tuple: {typ}"),
+            ast::Type::Function(_, _) => Type::Function,
 
-        for (_i, _parameter) in parameters.enumerate() {
-            todo!("Add block param to definitions")
+            // How should we represent Vecs?
+            // Are they a struct of array + length + capacity?
+            // Or are they just references?
+            ast::Type::Vec(_) => Type::Reference,
         }
     }
 }