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

@@ -4,11 +4,15 @@ use sway_types::{Span, Spanned};
 
 use crate::{
     concurrent_slab::ConcurrentSlab,
+    error::{err, ok},
+    namespace::{Path, Root},
     semantic_analysis::{
         TypedImplTrait, TypedStorageDeclaration, TypedStructDeclaration, TypedTraitDeclaration,
         TypedTraitFn,
     },
-    CompileError, TypedFunctionDeclaration,
+    type_system::TypeArgument,
+    type_system::{type_engine::monomorphize, EnforceTypeArguments},
+    CompileError, CompileResult, TypedFunctionDeclaration,
 };
 
 use super::{declaration_id::DeclarationId, declaration_wrapper::DeclarationWrapper};
@@ -24,6 +28,7 @@ pub(crate) struct DeclarationEngine {
     // *declaration_id -> vec of monomorphized copies
     // where the declaration_id is the original declaration
     monomorphized_copies: RwLock<HashMap<usize, Vec<DeclarationId>>>,
+    is_monomorph_cache_enabled: bool,
 }
 
 impl DeclarationEngine {
@@ -39,14 +44,10 @@ impl DeclarationEngine {
 
     fn de_add_monomorphized_copy(&self, original_id: DeclarationId, new_id: DeclarationId) {
         let mut monomorphized_copies = self.monomorphized_copies.write().unwrap();
-        match monomorphized_copies.get_mut(&*original_id) {
-            Some(prev) => {
-                prev.push(new_id);
-            }
-            None => {
-                monomorphized_copies.insert(*original_id, vec![new_id]);
-            }
-        }
+        monomorphized_copies
+            .entry(*original_id)
+            .and_modify(|f| f.push(new_id.clone()))
+            .or_insert_with(|| vec![new_id.clone()]);
     }
 
     fn de_get_monomorphized_copies(&self, original_id: DeclarationId) -> Vec<DeclarationWrapper> {
@@ -76,17 +77,69 @@ impl DeclarationEngine {
         self.slab.get(*index).expect_function(span)
     }
 
-    fn de_add_monomorphized_function_copy(
+    pub(crate) fn de_get_monomorphized_function(
         &self,
         original_id: DeclarationId,
-        new_copy: TypedFunctionDeclaration,
-    ) {
-        let span = new_copy.span();
-        let new_id = DeclarationId::new(
-            self.slab.insert(DeclarationWrapper::Function(new_copy)),
-            span,
+        type_arguments: &Vec<TypeArgument>,
+    ) -> Option<TypedFunctionDeclaration> {
+        for copy in self
+            .de_get_monomorphized_copies(original_id)
+            .iter()
+            .cloned()
+        {
+            let monomorphized_fn_decl = copy.expect_function(&Span::dummy()).unwrap();
+            if &monomorphized_fn_decl.type_parameters == type_arguments {
+                return Some(monomorphized_fn_decl);
+            }
+        }
+        None
+    }
+
+    pub(crate) fn de_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();
+
+        if self.is_monomorph_cache_enabled {
+            let cached_fn_decl =
+                self.de_get_monomorphized_function(decl_id.clone(), 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
+            .de_get_function(decl_id.clone(), call_site_span)
+            .unwrap();
+        check!(
+            monomorphize(
+                &mut typed_function_declaration,
+                type_arguments,
+                enforce_type_arguments,
+                call_site_span,
+                namespace,
+                module_path,
+            ),
+            return err(warnings, errors),
+            warnings,
+            errors
         );
-        self.de_add_monomorphized_copy(original_id, new_id)
+
+        // add the new copy to the declaration engine
+        let new_id = self.slab.insert(DeclarationWrapper::Function(
+            typed_function_declaration.clone(),
+        ));
+        self.de_add_monomorphized_copy(decl_id, DeclarationId::new(new_id, call_site_span.clone()));
+
+        ok(typed_function_declaration, warnings, errors)
     }
 
     fn de_get_monomorphized_function_copies(
@@ -129,7 +182,7 @@ impl DeclarationEngine {
         self.slab.get(*index).expect_trait_fn(span)
     }
 
-    fn insert_trait_impl(&self, trait_impl: TypedImplTrait) -> DeclarationId {
+    fn de_insert_trait_impl(&self, trait_impl: TypedImplTrait) -> DeclarationId {
         let span = trait_impl.span.clone();
         DeclarationId::new(
             self.slab.insert(DeclarationWrapper::TraitImpl(trait_impl)),
@@ -158,26 +211,69 @@ impl DeclarationEngine {
         self.slab.get(*index).expect_struct(span)
     }
 
-    fn de_add_monomorphized_struct_copy(
+    pub(crate) fn de_get_monomorphized_struct(
         &self,
         original_id: DeclarationId,
-        new_copy: TypedStructDeclaration,
-    ) {
-        let span = new_copy.span();
-        let new_id =
-            DeclarationId::new(self.slab.insert(DeclarationWrapper::Struct(new_copy)), span);
-        self.de_add_monomorphized_copy(original_id, new_id)
+        type_arguments: &Vec<TypeArgument>,
+        span: &Span,
+    ) -> Option<TypedStructDeclaration> {
+        for copy in self
+            .de_get_monomorphized_copies(original_id)
+            .iter()
+            .cloned()
+        {
+            let monomorphized_struct_decl = copy.expect_struct(span).unwrap();
+            if &monomorphized_struct_decl.type_parameters == type_arguments {
+                return Some(monomorphized_struct_decl);
+            }
+        }
+        None
     }
 
-    fn de_get_monomorphized_struct_copies(
+    pub(crate) fn de_get_or_create_monomorphized_struct(
         &self,
-        original_id: DeclarationId,
-        span: &Span,
-    ) -> Result<Vec<TypedStructDeclaration>, CompileError> {
-        self.de_get_monomorphized_copies(original_id)
-            .into_iter()
-            .map(|x| x.expect_struct(span))
-            .collect::<Result<_, _>>()
+        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();
+
+        if self.is_monomorph_cache_enabled {
+            let cached_struct_decl =
+                self.de_get_monomorphized_struct(decl_id.clone(), type_arguments, call_site_span);
+            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.de_get_struct(decl_id.clone(), call_site_span).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.de_add_monomorphized_copy(decl_id, DeclarationId::new(new_id, call_site_span.clone()));
+
+        ok(typed_struct_declaration, warnings, errors)
     }
 
     fn de_insert_storage(&self, storage: TypedStorageDeclaration) -> DeclarationId {
@@ -213,20 +309,6 @@ pub(crate) fn de_get_function(
     DECLARATION_ENGINE.de_get_function(index, span)
 }
 
-pub(crate) fn de_add_monomorphized_function_copy(
-    original_id: DeclarationId,
-    new_copy: TypedFunctionDeclaration,
-) {
-    DECLARATION_ENGINE.de_add_monomorphized_function_copy(original_id, new_copy);
-}
-
-pub(crate) fn de_get_monomorphized_function_copies(
-    original_id: DeclarationId,
-    span: &Span,
-) -> Result<Vec<TypedFunctionDeclaration>, CompileError> {
-    DECLARATION_ENGINE.de_get_monomorphized_function_copies(original_id, span)
-}
-
 pub(crate) fn de_insert_trait(r#trait: TypedTraitDeclaration) -> DeclarationId {
     DECLARATION_ENGINE.de_insert_trait(r#trait)
 }
@@ -249,8 +331,8 @@ pub(crate) fn de_get_trait_fn(
     DECLARATION_ENGINE.de_get_trait_fn(index, span)
 }
 
-pub(crate) fn insert_trait_impl(trait_impl: TypedImplTrait) -> DeclarationId {
-    DECLARATION_ENGINE.insert_trait_impl(trait_impl)
+pub(crate) fn de_insert_trait_impl(trait_impl: TypedImplTrait) -> DeclarationId {
+    DECLARATION_ENGINE.de_insert_trait_impl(trait_impl)
 }
 
 pub(crate) fn de_get_trait_impl(
@@ -264,25 +346,29 @@ pub(crate) fn de_insert_struct(r#struct: TypedStructDeclaration) -> DeclarationI
     DECLARATION_ENGINE.de_insert_struct(r#struct)
 }
 
-pub(crate) fn de_get_struct(
+pub fn de_get_struct(
     index: DeclarationId,
     span: &Span,
 ) -> Result<TypedStructDeclaration, CompileError> {
     DECLARATION_ENGINE.de_get_struct(index, span)
 }
 
-pub(crate) fn de_add_monomorphized_struct_copy(
-    original_id: DeclarationId,
-    new_copy: TypedStructDeclaration,
-) {
-    DECLARATION_ENGINE.de_add_monomorphized_struct_copy(original_id, new_copy);
-}
-
-pub(crate) fn de_get_monomorphized_struct_copies(
-    original_id: DeclarationId,
-    span: &Span,
-) -> Result<Vec<TypedStructDeclaration>, CompileError> {
-    DECLARATION_ENGINE.de_get_monomorphized_struct_copies(original_id, span)
+pub(crate) fn de_get_or_create_monomorphized_struct(
+    decl_id: DeclarationId,
+    type_arguments: &mut Vec<TypeArgument>,
+    enforce_type_arguments: EnforceTypeArguments,
+    call_site_span: &Span,
+    namespace: &Root,
+    module_path: &Path,
+) -> CompileResult<TypedStructDeclaration> {
+    DECLARATION_ENGINE.de_get_or_create_monomorphized_struct(
+        decl_id,
+        type_arguments,
+        enforce_type_arguments,
+        call_site_span,
+        namespace,
+        module_path,
+    )
 }
 
 pub(crate) fn de_insert_storage(storage: TypedStorageDeclaration) -> 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

@@ -37,6 +37,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) {