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 #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, MonomorphizeHelper, 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,30 +77,6 @@ impl DeclarationEngine {
         self.slab.get(*index).expect_function(span)
     }
 
-    fn de_add_monomorphized_function_copy(
-        &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,
-        );
-        self.de_add_monomorphized_copy(original_id, new_id)
-    }
-
-    fn de_get_monomorphized_function_copies(
-        &self,
-        original_id: DeclarationId,
-        span: &Span,
-    ) -> Result<Vec<TypedFunctionDeclaration>, CompileError> {
-        self.de_get_monomorphized_copies(original_id)
-            .into_iter()
-            .map(|x| x.expect_function(span))
-            .collect::<Result<_, _>>()
-    }
-
     fn de_insert_trait(&self, r#trait: TypedTraitDeclaration) -> DeclarationId {
         let span = r#trait.name.span();
         DeclarationId::new(self.slab.insert(DeclarationWrapper::Trait(r#trait)), span)
@@ -129,7 +106,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 +135,78 @@ impl DeclarationEngine {
         self.slab.get(*index).expect_struct(span)
     }
 
-    fn de_add_monomorphized_struct_copy(
+    pub(crate) fn de_get_monomorphized_decl(
         &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,
+    ) -> Result<DeclarationWrapper, CompileError> {
+        for monomorphized_decl in self
+            .de_get_monomorphized_copies(original_id)
+            .iter()
+            .cloned()
+        {
+            if monomorphized_decl.type_parameters() == type_arguments {
+                return Ok(monomorphized_decl);
+            }
+        }
+        Err(CompileError::Internal(
+            "could not find monomorphized decl",
+            span.clone(),
+        ))
     }
 
-    fn de_get_monomorphized_struct_copies(
+    pub(crate) fn de_get_or_create_monomorphized_decl<T>(
         &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<DeclarationWrapper>
+    where
+        T: MonomorphizeHelper,
+    {
+        let mut warnings = Vec::new();
+        let mut errors = Vec::new();
+
+        if self.is_monomorph_cache_enabled {
+            let cached_decl = check!(
+                CompileResult::from(self.de_get_monomorphized_decl(
+                    decl_id,
+                    type_arguments,
+                    call_site_span
+                )),
+                return err(warnings, errors),
+                warnings,
+                errors
+            );
+            return ok(cached_decl, warnings, errors);
+        }
+
+        // monomorphize the declaration into a new copy
+        let mut typed_declaration = self.slab.get(*decl_id);
+
+        check!(
+            monomorphize(
+                &mut typed_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(typed_declaration.to_wrapper());
+        self.de_add_monomorphized_copy(decl_id, DeclarationId::new(new_id, call_site_span.clone()));
+
+        ok(typed_declaration, warnings, errors)
     }
 
     fn de_insert_storage(&self, storage: TypedStorageDeclaration) -> DeclarationId {
@@ -213,20 +242,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 +264,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,27 +279,13 @@ 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_insert_storage(storage: TypedStorageDeclaration) -> DeclarationId {
     DECLARATION_ENGINE.de_insert_storage(storage)
 }
@@ -295,3 +296,24 @@ pub fn de_get_storage(
 ) -> Result<TypedStorageDeclaration, CompileError> {
     DECLARATION_ENGINE.de_get_storage(index, span)
 }
+
+pub(crate) fn de_get_or_create_monomorphized_decl<T>(
+    decl_id: DeclarationId,
+    type_arguments: &mut Vec<TypeArgument>,
+    enforce_type_arguments: EnforceTypeArguments,
+    call_site_span: &Span,
+    namespace: &Root,
+    module_path: &Path,
+) -> CompileResult<DeclarationWrapper>
+where
+    T: MonomorphizeHelper,
+{
+    DECLARATION_ENGINE.de_get_or_create_monomorphized_decl::<T>(
+        decl_id,
+        type_arguments,
+        enforce_type_arguments,
+        call_site_span,
+        namespace,
+        module_path,
+    )
+}

sway-core/src/declaration_engine/declaration_wrapper.rs

@@ -1,14 +1,14 @@
 use std::fmt;
 
-use sway_types::Span;
+use sway_types::{Ident, Span};
 
 use crate::{
     semantic_analysis::{
-        TypedImplTrait, TypedStorageDeclaration, TypedStructDeclaration, TypedTraitDeclaration,
-        TypedTraitFn,
+        TypedEnumDeclaration, TypedImplTrait, TypedStorageDeclaration, TypedStructDeclaration,
+        TypedTraitDeclaration, TypedTraitFn,
     },
     type_system::{CopyTypes, TypeMapping},
-    CompileError, TypedFunctionDeclaration,
+    CompileError, MonomorphizeHelper, TypeParameter, TypedFunctionDeclaration,
 };
 
 /// The [DeclarationWrapper] type is used in the [DeclarationEngine]
@@ -17,6 +17,7 @@ use crate::{
 pub(crate) enum DeclarationWrapper {
     // no-op variant to fulfill the default trait
     Unknown,
+    Enum(TypedEnumDeclaration),
     Function(TypedFunctionDeclaration),
     Trait(TypedTraitDeclaration),
     TraitFn(TypedTraitFn),
@@ -59,6 +60,7 @@ impl CopyTypes for DeclarationWrapper {
     fn copy_types(&mut self, type_mapping: &TypeMapping) {
         match self {
             DeclarationWrapper::Unknown => {}
+            DeclarationWrapper::Enum(decl) => decl.copy_types(type_mapping),
             DeclarationWrapper::Function(decl) => decl.copy_types(type_mapping),
             DeclarationWrapper::Trait(decl) => decl.copy_types(type_mapping),
             DeclarationWrapper::TraitFn(decl) => decl.copy_types(type_mapping),
@@ -69,11 +71,47 @@ impl CopyTypes for DeclarationWrapper {
     }
 }
 
+impl MonomorphizeHelper for DeclarationWrapper {
+    fn type_parameters(&self) -> &[TypeParameter] {
+        match self {
+            DeclarationWrapper::Trait(_)
+            | DeclarationWrapper::TraitFn(_)
+            | DeclarationWrapper::TraitImpl(_)
+            | DeclarationWrapper::Storage(_)
+            | DeclarationWrapper::Unknown => {
+                panic!("declaration type does not support type parameters")
+            }
+            DeclarationWrapper::Enum(decl) => decl.type_parameters(),
+            DeclarationWrapper::Function(decl) => decl.type_parameters(),
+            DeclarationWrapper::Struct(decl) => decl.type_parameters(),
+        }
+    }
+
+    fn name(&self) -> &Ident {
+        match self {
+            DeclarationWrapper::Storage(_) | DeclarationWrapper::Unknown => {
+                panic!("declaration type is not named")
+            }
+            DeclarationWrapper::Enum(decl) => decl.name(),
+            DeclarationWrapper::Function(decl) => decl.name(),
+            DeclarationWrapper::Trait(decl) => &decl.name,
+            DeclarationWrapper::TraitFn(decl) => &decl.name,
+            DeclarationWrapper::TraitImpl(decl) => &decl.trait_name.suffix,
+            DeclarationWrapper::Struct(decl) => decl.name(),
+        }
+    }
+
+    fn to_wrapper(&self) -> DeclarationWrapper {
+        panic!("not expected to be called for DeclarationWrapper")
+    }
+}
+
 impl DeclarationWrapper {
     /// friendly name string used for error reporting.
     fn friendly_name(&self) -> &'static str {
         match self {
             DeclarationWrapper::Unknown => "unknown",
+            DeclarationWrapper::Enum(_) => "enum",
             DeclarationWrapper::Function(_) => "function",
             DeclarationWrapper::Trait(_) => "trait",
             DeclarationWrapper::Struct(_) => "struct",

sway-core/src/semantic_analysis/ast_node/declaration/enum.rs

@@ -1,4 +1,5 @@
 use crate::{
+    declaration_engine::declaration_wrapper::DeclarationWrapper,
     error::*,
     parse_tree::*,
     semantic_analysis::*,
@@ -67,6 +68,10 @@ impl MonomorphizeHelper for TypedEnumDeclaration {
     fn name(&self) -> &Ident {
         &self.name
     }
+
+    fn to_wrapper(&self) -> DeclarationWrapper {
+        DeclarationWrapper::Enum(self.clone())
+    }
 }
 
 impl TypedEnumDeclaration {

sway-core/src/semantic_analysis/ast_node/declaration/function.rs

@@ -1,7 +1,10 @@
 mod function_parameter;
 pub use function_parameter::*;
 
-use crate::{error::*, parse_tree::*, semantic_analysis::*, style::*, type_system::*};
+use crate::{
+    declaration_engine::declaration_wrapper::DeclarationWrapper, error::*, parse_tree::*,
+    semantic_analysis::*, style::*, type_system::*,
+};
 use sha2::{Digest, Sha256};
 use sway_types::{Ident, JsonABIFunction, JsonTypeApplication, JsonTypeDeclaration, Span, Spanned};
 
@@ -79,6 +82,10 @@ impl MonomorphizeHelper for TypedFunctionDeclaration {
     fn name(&self) -> &Ident {
         &self.name
     }
+
+    fn to_wrapper(&self) -> DeclarationWrapper {
+        DeclarationWrapper::Function(self.clone())
+    }
 }
 
 impl TypedFunctionDeclaration {