chore(ssa refactor): Implement first-class functions (#1238)

* Implement first-class functions

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

Co-authored-by: kevaundray <kevtheappdev@gmail.com>

---------

Co-authored-by: kevaundray <kevtheappdev@gmail.com>

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

@@ -1,9 +1,11 @@
+use std::collections::HashMap;
+
 use super::{
     basic_block::{BasicBlock, BasicBlockId},
     constant::{NumericConstant, NumericConstantId},
-    function::Signature,
+    function::{FunctionId, Signature},
     instruction::{Instruction, InstructionId, InstructionResultType, TerminatorInstruction},
-    map::{DenseMap, Id, SecondaryMap, TwoWayMap},
+    map::{DenseMap, Id, TwoWayMap},
     types::Type,
     value::{Value, ValueId},
 };
@@ -53,7 +55,7 @@ pub(crate) struct DataFlowGraph {
     /// Currently, we need to define them in a better way
     /// Call instructions require the func signature, but
     /// other instructions may need some more reading on my part
-    results: SecondaryMap<Instruction, ValueList>,
+    results: HashMap<InstructionId, ValueList>,
 
     /// Storage for all of the values defined in this
     /// function.
@@ -64,6 +66,11 @@ pub(crate) struct DataFlowGraph {
     /// twice will return the same ConstantId.
     constants: TwoWayMap<NumericConstant>,
 
+    /// Contains each function that has been imported into the current function.
+    /// Each function's Value::Function is uniqued here so any given FunctionId
+    /// will always have the same ValueId within this function.
+    functions: HashMap<FunctionId, ValueId>,
+
     /// Function signatures of external methods
     signatures: DenseMap<Signature>,
 
@@ -150,6 +157,14 @@ impl DataFlowGraph {
         self.values.insert(Value::NumericConstant { constant, typ })
     }
 
+    /// Gets or creates a ValueId for the given FunctionId.
+    pub(crate) fn import_function(&mut self, function: FunctionId) -> ValueId {
+        if let Some(existing) = self.functions.get(&function) {
+            return *existing;
+        }
+        self.values.insert(Value::Function { id: function })
+    }
+
     /// Attaches results to the instruction, clearing any previous results.
     ///
     /// This does not normally need to be called manually as it is called within

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

@@ -1,7 +1,9 @@
+use std::collections::HashMap;
+
 use super::basic_block::BasicBlockId;
 use super::dfg::DataFlowGraph;
-use super::instruction::Instruction;
-use super::map::{Id, SecondaryMap};
+use super::instruction::InstructionId;
+use super::map::Id;
 use super::types::Type;
 
 use noirc_errors::Location;
@@ -15,7 +17,7 @@ use noirc_errors::Location;
 #[derive(Debug)]
 pub(crate) struct Function {
     /// Maps instructions to source locations
-    source_locations: SecondaryMap<Instruction, Location>,
+    source_locations: HashMap<InstructionId, Location>,
 
     /// The first basic block in the function
     entry_block: BasicBlockId,
@@ -35,7 +37,7 @@ impl Function {
     pub(crate) fn new(name: String, id: FunctionId) -> Self {
         let mut dfg = DataFlowGraph::default();
         let entry_block = dfg.make_block();
-        Self { name, source_locations: SecondaryMap::new(), id, entry_block, dfg }
+        Self { name, source_locations: HashMap::new(), id, entry_block, dfg }
     }
 
     pub(crate) fn name(&self) -> &str {

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

@@ -1,6 +1,4 @@
-use super::{
-    basic_block::BasicBlockId, function::FunctionId, map::Id, types::Type, value::ValueId,
-};
+use super::{basic_block::BasicBlockId, map::Id, types::Type, value::ValueId};
 
 /// Reference to an instruction
 pub(crate) type InstructionId = Id<Instruction>;
@@ -41,7 +39,7 @@ pub(crate) enum Instruction {
     Constrain(ValueId),
 
     /// Performs a function call with a list of its arguments.
-    Call { func: FunctionId, arguments: Vec<ValueId> },
+    Call { func: ValueId, arguments: Vec<ValueId> },
 
     /// Performs a call to an intrinsic function and stores the
     /// results in `return_arguments`.

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

@@ -260,20 +260,6 @@ impl<T> std::ops::Index<Id<T>> for TwoWayMap<T> {
     }
 }
 
-/// A SecondaryMap is for storing secondary data for a given key. Since this
-/// map is for secondary data, it will not return fresh Ids for data, instead
-/// it expects users to provide these ids in order to associate existing ids with
-/// additional data.
-///
-/// Unlike SecondaryMap in cranelift, this version is sparse and thus
-/// does not require inserting default elements for each key in between
-/// the desired key and the previous length of the map.
-///
-/// There is no expectation that there is always secondary data for all relevant
-/// Ids of a given type, so unlike the other Map types, it is possible for
-/// a call to .get(id) to return None.
-pub(crate) type SecondaryMap<K, V> = HashMap<Id<K>, V>;
-
 /// A simple counter to create fresh Ids without any storage.
 /// Useful for assigning ids before the storage is created or assigning ids
 /// for types that have no single owner.

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

@@ -55,12 +55,17 @@ pub(crate) fn display_block(
     display_terminator(function, block.terminator(), f)
 }
 
-/// Specialize displaying value ids so that if they refer to constants we
-/// print the constant directly
+/// Specialize displaying value ids so that if they refer to a numeric
+/// constant or a function we print those directly.
 fn value(function: &Function, id: ValueId) -> String {
-    match function.dfg.get_numeric_constant_with_type(id) {
-        Some((value, typ)) => format!("{} {}", value, typ),
-        None => id.to_string(),
+    use super::value::Value;
+    match &function.dfg[id] {
+        Value::NumericConstant { constant, typ } => {
+            let value = function.dfg[*constant].value();
+            format!("{} {}", typ, value)
+        }
+        Value::Function { id } => id.to_string(),
+        _ => id.to_string(),
     }
 }
 
@@ -120,7 +125,7 @@ pub(crate) fn display_instruction(
             writeln!(f, "constrain {}", show(*value))
         }
         Instruction::Call { func, arguments } => {
-            writeln!(f, "call {func}({})", value_list(function, arguments))
+            writeln!(f, "call {}({})", show(*func), value_list(function, arguments))
         }
         Instruction::Intrinsic { func, arguments } => {
             writeln!(f, "intrinsic {func}({})", value_list(function, arguments))

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

@@ -1,6 +1,9 @@
 use crate::ssa_refactor::ir::basic_block::BasicBlockId;
 
-use super::{constant::NumericConstantId, instruction::InstructionId, map::Id, types::Type};
+use super::{
+    constant::NumericConstantId, function::FunctionId, instruction::InstructionId, map::Id,
+    types::Type,
+};
 
 pub(crate) type ValueId = Id<Value>;
 
@@ -27,6 +30,13 @@ pub(crate) enum Value {
 
     /// This Value originates from a numeric constant
     NumericConstant { constant: NumericConstantId, typ: Type },
+
+    /// This Value refers to a function in the IR.
+    /// Functions always have the type Type::Function.
+    /// If the argument or return types are needed, users should retrieve
+    /// their types via the Call instruction's arguments or the Call instruction's
+    /// result types respectively.
+    Function { id: FunctionId },
 }
 
 impl Value {
@@ -35,6 +45,7 @@ impl Value {
             Value::Instruction { typ, .. } => *typ,
             Value::Param { typ, .. } => *typ,
             Value::NumericConstant { typ, .. } => *typ,
+            Value::Function { .. } => Type::Function,
         }
     }
 }

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

@@ -160,7 +160,7 @@ impl FunctionBuilder {
     /// the results of the call.
     pub(crate) fn insert_call(
         &mut self,
-        func: FunctionId,
+        func: ValueId,
         arguments: Vec<ValueId>,
         result_types: Vec<Type>,
     ) -> &[ValueId] {

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

@@ -164,7 +164,7 @@ impl<'a> FunctionContext<'a> {
     /// back into a Values tree of the proper shape.
     pub(super) fn insert_call(
         &mut self,
-        function: IrFunctionId,
+        function: ValueId,
         arguments: Vec<ValueId>,
         result_type: &ast::Type,
     ) -> Values {

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

@@ -11,7 +11,7 @@ use self::{
     value::{Tree, Values},
 };
 
-use super::ir::{function::FunctionId, instruction::BinaryOp, types::Type, value::ValueId};
+use super::ir::{instruction::BinaryOp, types::Type, value::ValueId};
 
 pub(crate) fn generate_ssa(program: Program) {
     let context = SharedContext::new(program);
@@ -255,18 +255,8 @@ impl<'a> FunctionContext<'a> {
         Self::get_field(tuple, field_index)
     }
 
-    fn codegen_function(&mut self, function: &Expression) -> FunctionId {
-        use crate::ssa_refactor::ssa_gen::value::Value;
-        match self.codegen_expression(function) {
-            Tree::Leaf(Value::Function(id)) => id,
-            other => {
-                panic!("codegen_function: expected function value, found {other:?}")
-            }
-        }
-    }
-
     fn codegen_call(&mut self, call: &ast::Call) -> Values {
-        let function = self.codegen_function(&call.func);
+        let function = self.codegen_non_tuple_expression(&call.func);
 
         let arguments = call
             .arguments