Add a monomorphization cache to the declaration engine.

This introduces a cache to the declaration engine that keeps track of
functions/structs that have been instantiated from a specific set
of type parameters.

This was originally prototyped in
tritao/declaration-engine-and-collection-context-demo@c3aa5e7,
where these APIs where wired
to the function application and struct expressions instantiation code.

However given the declaration engine is not hooked up yet here, this
PR is just adding the APIs, and more proper testing will be added as we
wire things up.

Closes FuelLabs#2636.

sway-core/src/declaration_engine/declaration_engine.rs

@@ -1,53 +1,50 @@
+use sway_types::Span;
+
+use super::{declaration_id::DeclarationId, declaration_wrapper::DeclarationWrapper};
 use crate::{
     concurrent_slab::ConcurrentSlab,
+    error::{err, ok},
+    namespace::{Path, Root},
     semantic_analysis::{
         TypedImplTrait, TypedStructDeclaration, TypedTraitDeclaration, TypedTraitFn,
     },
-    TypedFunctionDeclaration,
+    type_system::TypeArgument,
+    type_system::{type_engine::monomorphize, EnforceTypeArguments},
+    CompileResult, TypedFunctionDeclaration,
 };
-
-use super::{declaration_id::DeclarationId, declaration_wrapper::DeclarationWrapper};
+use std::collections::HashMap;
 
 /// Used inside of type inference to store declarations.
 pub struct DeclarationEngine {
     slab: ConcurrentSlab<DeclarationId, DeclarationWrapper>,
     // *declaration_id -> vec of monomorphized copies
     // where the declaration_id is the original declaration
-    monomorphized_copies: im::HashMap<usize, Vec<DeclarationId>>,
+    monomorphized_copies: HashMap<DeclarationId, Vec<DeclarationId>>,
 }
 
 impl DeclarationEngine {
     pub(crate) fn new() -> DeclarationEngine {
         DeclarationEngine {
             slab: ConcurrentSlab::default(),
-            monomorphized_copies: im::HashMap::new(),
+            monomorphized_copies: HashMap::new(),
         }
     }
 
     pub(crate) fn look_up_decl_id(&self, index: DeclarationId) -> DeclarationWrapper {
         self.slab.get(index)
     }
-
-    pub(crate) fn add_monomorphized_copy(
-        &mut self,
-        original_id: DeclarationId,
-        new_id: DeclarationId,
-    ) {
-        match self.monomorphized_copies.get_mut(&*original_id) {
-            Some(prev) => {
-                prev.push(new_id);
-            }
-            None => {
-                self.monomorphized_copies.insert(*original_id, vec![new_id]);
-            }
-        }
+    fn add_monomorphized_copy(&mut self, original_id: DeclarationId, new_id: DeclarationId) {
+        self.monomorphized_copies
+            .entry(original_id)
+            .and_modify(|f| f.push(new_id))
+            .or_insert_with(|| vec![new_id]);
     }
 
     pub(crate) fn get_monomorphized_copies(
         &self,
         original_id: DeclarationId,
     ) -> Vec<DeclarationWrapper> {
-        match self.monomorphized_copies.get(&*original_id).cloned() {
+        match self.monomorphized_copies.get(&original_id).cloned() {
             Some(copies) => copies.into_iter().map(|copy| self.slab.get(copy)).collect(),
             None => vec![],
         }
@@ -64,13 +61,64 @@ impl DeclarationEngine {
         self.slab.get(index).expect_function()
     }
 
-    pub(crate) fn add_monomorphized_function_copy(
-        &mut self,
+    pub(crate) fn get_monomorphized_function(
+        &self,
         original_id: DeclarationId,
-        new_copy: TypedFunctionDeclaration,
-    ) {
-        let new_id = self.slab.insert(DeclarationWrapper::Function(new_copy));
-        self.add_monomorphized_copy(original_id, new_id)
+        type_arguments: &Vec<TypeArgument>,
+    ) -> Option<TypedFunctionDeclaration> {
+        for copy in self.get_monomorphized_copies(original_id).iter() {
+            let monomorphized_fn_decl = match copy {
+                DeclarationWrapper::Function(fn_decl) => fn_decl,
+                _ => panic!("expected an existing function decl"),
+            };
+
+            if &monomorphized_fn_decl.type_parameters == type_arguments {
+                return Some(monomorphized_fn_decl.clone());
+            }
+        }
+        None
+    }
+
+    pub(crate) fn get_or_create_monomorphized_function(
+        &mut self,
+        decl_id: DeclarationId,
+        type_arguments: &mut Vec<TypeArgument>,
+        enforce_type_arguments: EnforceTypeArguments,
+        call_site_span: &Span,
+        namespace: &Root,
+        module_path: &Path,
+    ) -> CompileResult<TypedFunctionDeclaration> {
+        let mut warnings = Vec::new();
+        let mut errors = Vec::new();
+
+        let cached_fn_decl = self.get_monomorphized_function(decl_id, type_arguments);
+        if let Some(cached_fn_decl) = cached_fn_decl {
+            return ok(cached_fn_decl, warnings, errors);
+        }
+
+        // monomorphize the function declaration into a new copy
+        let mut typed_function_declaration = self.get_function(decl_id).unwrap();
+        check!(
+            monomorphize(
+                &mut typed_function_declaration,
+                type_arguments,
+                enforce_type_arguments,
+                call_site_span,
+                namespace,
+                module_path,
+            ),
+            return err(warnings, errors),
+            warnings,
+            errors
+        );
+
+        // add the new copy to the declaration engine
+        let new_id = self.slab.insert(DeclarationWrapper::Function(
+            typed_function_declaration.clone(),
+        ));
+        self.add_monomorphized_copy(decl_id, new_id);
+
+        ok(typed_function_declaration, warnings, errors)
     }
 
     pub(crate) fn get_monomorphized_function_copies(
@@ -118,13 +166,64 @@ impl DeclarationEngine {
         self.slab.get(index).expect_struct()
     }
 
-    pub(crate) fn add_monomorphized_struct_copy(
-        &mut self,
+    pub(crate) fn get_monomorphized_struct(
+        &self,
         original_id: DeclarationId,
-        new_copy: TypedStructDeclaration,
-    ) {
-        let new_id = self.slab.insert(DeclarationWrapper::Struct(new_copy));
-        self.add_monomorphized_copy(original_id, new_id)
+        type_arguments: &Vec<TypeArgument>,
+    ) -> Option<TypedStructDeclaration> {
+        for copy in self.get_monomorphized_copies(original_id).iter() {
+            let monomorphized_struct_decl = match copy {
+                DeclarationWrapper::Struct(struct_decl) => struct_decl,
+                _ => panic!("expected an existing struct decl"),
+            };
+
+            if &monomorphized_struct_decl.type_parameters == type_arguments {
+                return Some(monomorphized_struct_decl.clone());
+            }
+        }
+        None
+    }
+
+    pub(crate) fn get_or_create_monomorphized_struct(
+        &mut self,
+        decl_id: DeclarationId,
+        type_arguments: &mut Vec<TypeArgument>,
+        enforce_type_arguments: EnforceTypeArguments,
+        call_site_span: &Span,
+        namespace: &Root,
+        module_path: &Path,
+    ) -> CompileResult<TypedStructDeclaration> {
+        let mut warnings = Vec::new();
+        let mut errors = Vec::new();
+
+        let cached_struct_decl = self.get_monomorphized_struct(decl_id, type_arguments);
+        if let Some(cached_struct_decl) = cached_struct_decl {
+            return ok(cached_struct_decl, warnings, errors);
+        }
+
+        // monomorphize the struct declaration into a new copy
+        let mut typed_struct_declaration = self.get_struct(decl_id).unwrap();
+        check!(
+            monomorphize(
+                &mut typed_struct_declaration,
+                type_arguments,
+                enforce_type_arguments,
+                call_site_span,
+                namespace,
+                module_path,
+            ),
+            return err(warnings, errors),
+            warnings,
+            errors
+        );
+
+        // add the new copy to the declaration engine
+        let new_id = self
+            .slab
+            .insert(DeclarationWrapper::Struct(typed_struct_declaration.clone()));
+        self.add_monomorphized_copy(decl_id, new_id);
+
+        ok(typed_struct_declaration, warnings, errors)
     }
 
     pub(crate) fn get_monomorphized_struct_copies(

sway-core/src/declaration_engine/declaration_id.rs

@@ -1,7 +1,7 @@
 use std::fmt;
 
 /// An ID used to refer to an item in the [DeclarationEngine]
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
 pub struct DeclarationId(usize);
 
 impl fmt::Display for DeclarationId {

sway-core/src/type_system/mod.rs

@@ -6,7 +6,7 @@ mod resolved_type;
 mod trait_constraint;
 mod type_argument;
 mod type_binding;
-mod type_engine;
+pub(crate) mod type_engine;
 mod type_id;
 mod type_info;
 mod type_mapping;

sway-core/src/type_system/type_argument.rs

@@ -30,6 +30,12 @@ impl PartialEq for TypeArgument {
     }
 }
 
+impl PartialEq<TypeParameter> for TypeArgument {
+    fn eq(&self, other: &TypeParameter) -> bool {
+        self.type_id == other.type_id
+    }
+}
+
 impl Default for TypeArgument {
     fn default() -> Self {
         let initial_type_id = insert_type(TypeInfo::Unknown);

sway-core/src/type_system/type_parameter.rs

@@ -42,6 +42,12 @@ impl PartialEq for TypeParameter {
     }
 }
 
+impl PartialEq<TypeArgument> for TypeParameter {
+    fn eq(&self, other: &TypeArgument) -> bool {
+        self.type_id == other.type_id
+    }
+}
+
 impl CopyTypes for TypeParameter {
     fn copy_types(&mut self, type_mapping: &TypeMapping) {
         self.type_id = match look_up_type_id(self.type_id).matches_type_parameter(type_mapping) {