Implement a map from typed to parsed decl ids in the declaration engine.

sway-core/src/decl_engine/engine.rs

@@ -9,12 +9,19 @@ use sway_types::{Named, ProgramId, Spanned};
 
 use crate::{
     concurrent_slab::ConcurrentSlab,
-    decl_engine::*,
+    decl_engine::{parsed_id::ParsedDeclId, *},
     engine_threading::*,
-    language::ty::{
-        self, TyAbiDecl, TyConfigurableDecl, TyConstantDecl, TyEnumDecl, TyFunctionDecl,
-        TyImplSelfOrTrait, TyStorageDecl, TyStructDecl, TyTraitDecl, TyTraitFn, TyTraitType,
-        TyTypeAliasDecl,
+    language::{
+        parsed::{
+            AbiDeclaration, ConfigurableDeclaration, ConstantDeclaration, EnumDeclaration,
+            FunctionDeclaration, ImplSelfOrTrait, StorageDeclaration, StructDeclaration,
+            TraitDeclaration, TraitFn, TraitTypeDeclaration, TypeAliasDeclaration,
+        },
+        ty::{
+            self, TyAbiDecl, TyConfigurableDecl, TyConstantDecl, TyDeclParsedType, TyEnumDecl,
+            TyFunctionDecl, TyImplSelfOrTrait, TyStorageDecl, TyStructDecl, TyTraitDecl, TyTraitFn,
+            TyTraitType, TyTypeAliasDecl,
+        },
     },
 };
 
@@ -34,6 +41,27 @@ pub struct DeclEngine {
     enum_slab: ConcurrentSlab<TyEnumDecl>,
     type_alias_slab: ConcurrentSlab<TyTypeAliasDecl>,
 
+    function_parsed_decl_id_map:
+        RwLock<HashMap<DeclId<TyFunctionDecl>, ParsedDeclId<FunctionDeclaration>>>,
+    trait_parsed_decl_id_map: RwLock<HashMap<DeclId<TyTraitDecl>, ParsedDeclId<TraitDeclaration>>>,
+    trait_fn_parsed_decl_id_map: RwLock<HashMap<DeclId<TyTraitFn>, ParsedDeclId<TraitFn>>>,
+    trait_type_parsed_decl_id_map:
+        RwLock<HashMap<DeclId<TyTraitType>, ParsedDeclId<TraitTypeDeclaration>>>,
+    impl_self_or_trait_parsed_decl_id_map:
+        RwLock<HashMap<DeclId<TyImplSelfOrTrait>, ParsedDeclId<ImplSelfOrTrait>>>,
+    struct_parsed_decl_id_map:
+        RwLock<HashMap<DeclId<TyStructDecl>, ParsedDeclId<StructDeclaration>>>,
+    storage_parsed_decl_id_map:
+        RwLock<HashMap<DeclId<TyStorageDecl>, ParsedDeclId<StorageDeclaration>>>,
+    abi_parsed_decl_id_map: RwLock<HashMap<DeclId<TyAbiDecl>, ParsedDeclId<AbiDeclaration>>>,
+    constant_parsed_decl_id_map:
+        RwLock<HashMap<DeclId<TyConstantDecl>, ParsedDeclId<ConstantDeclaration>>>,
+    configurable_parsed_decl_id_map:
+        RwLock<HashMap<DeclId<TyConfigurableDecl>, ParsedDeclId<ConfigurableDeclaration>>>,
+    enum_parsed_decl_id_map: RwLock<HashMap<DeclId<TyEnumDecl>, ParsedDeclId<EnumDeclaration>>>,
+    type_alias_parsed_decl_id_map:
+        RwLock<HashMap<DeclId<TyTypeAliasDecl>, ParsedDeclId<TypeAliasDeclaration>>>,
+
     parents: RwLock<HashMap<AssociatedItemDeclId, Vec<AssociatedItemDeclId>>>,
 }
 
@@ -52,6 +80,32 @@ impl Clone for DeclEngine {
             configurable_slab: self.configurable_slab.clone(),
             enum_slab: self.enum_slab.clone(),
             type_alias_slab: self.type_alias_slab.clone(),
+            function_parsed_decl_id_map: RwLock::new(
+                self.function_parsed_decl_id_map.read().clone(),
+            ),
+            trait_parsed_decl_id_map: RwLock::new(self.trait_parsed_decl_id_map.read().clone()),
+            trait_fn_parsed_decl_id_map: RwLock::new(
+                self.trait_fn_parsed_decl_id_map.read().clone(),
+            ),
+            trait_type_parsed_decl_id_map: RwLock::new(
+                self.trait_type_parsed_decl_id_map.read().clone(),
+            ),
+            impl_self_or_trait_parsed_decl_id_map: RwLock::new(
+                self.impl_self_or_trait_parsed_decl_id_map.read().clone(),
+            ),
+            struct_parsed_decl_id_map: RwLock::new(self.struct_parsed_decl_id_map.read().clone()),
+            storage_parsed_decl_id_map: RwLock::new(self.storage_parsed_decl_id_map.read().clone()),
+            abi_parsed_decl_id_map: RwLock::new(self.abi_parsed_decl_id_map.read().clone()),
+            constant_parsed_decl_id_map: RwLock::new(
+                self.constant_parsed_decl_id_map.read().clone(),
+            ),
+            configurable_parsed_decl_id_map: RwLock::new(
+                self.configurable_parsed_decl_id_map.read().clone(),
+            ),
+            enum_parsed_decl_id_map: RwLock::new(self.enum_parsed_decl_id_map.read().clone()),
+            type_alias_parsed_decl_id_map: RwLock::new(
+                self.type_alias_parsed_decl_id_map.read().clone(),
+            ),
             parents: RwLock::new(self.parents.read().clone()),
         }
     }
@@ -61,26 +115,40 @@ pub trait DeclEngineGet<I, U> {
     fn get(&self, index: &I) -> Arc<U>;
 }
 
+pub trait DeclEngineGetParsedDeclId<T>
+where
+    T: TyDeclParsedType,
+{
+    fn get_parsed_decl_id(&self, decl_id: &DeclId<T>) -> Option<ParsedDeclId<T::ParsedType>>;
+}
+
 pub trait DeclEngineInsert<T>
 where
-    T: Named + Spanned,
+    T: Named + Spanned + TyDeclParsedType,
 {
-    fn insert(&self, decl: T) -> DeclRef<DeclId<T>>;
+    fn insert(
+        &self,
+        decl: T,
+        parsed_decl_id: Option<&ParsedDeclId<T::ParsedType>>,
+    ) -> DeclRef<DeclId<T>>;
 }
 
 pub trait DeclEngineInsertArc<T>
 where
-    T: Named + Spanned,
+    T: Named + Spanned + TyDeclParsedType,
 {
-    fn insert_arc(&self, decl: Arc<T>) -> DeclRef<DeclId<T>>;
+    fn insert_arc(
+        &self,
+        decl: Arc<T>,
+        parsed_decl_id: Option<&ParsedDeclId<T::ParsedType>>,
+    ) -> DeclRef<DeclId<T>>;
 }
 
 pub trait DeclEngineReplace<T> {
     fn replace(&self, index: DeclId<T>, decl: T);
 }
 
-pub trait DeclEngineIndex<T>:
-    DeclEngineGet<DeclId<T>, T> + DeclEngineInsert<T> + DeclEngineReplace<T>
+pub trait DeclEngineIndex<T>: DeclEngineGet<DeclId<T>, T> + DeclEngineReplace<T>
 where
     T: Named + Spanned,
 {
@@ -115,41 +183,110 @@ decl_engine_get!(enum_slab, ty::TyEnumDecl);
 decl_engine_get!(type_alias_slab, ty::TyTypeAliasDecl);
 
 macro_rules! decl_engine_insert {
-    ($slab:ident, $decl:ty) => {
+    ($slab:ident, $parsed_slab:ident, $decl:ty) => {
         impl DeclEngineInsert<$decl> for DeclEngine {
-            fn insert(&self, decl: $decl) -> DeclRef<DeclId<$decl>> {
+            fn insert(
+                &self,
+                decl: $decl,
+                parsed_decl_id: Option<&ParsedDeclId<<$decl as TyDeclParsedType>::ParsedType>>,
+            ) -> DeclRef<DeclId<$decl>> {
                 let span = decl.span();
-                DeclRef::new(
-                    decl.name().clone(),
-                    DeclId::new(self.$slab.insert(decl)),
-                    span,
-                )
+                let decl_name = decl.name().clone();
+                let decl_id = DeclId::new(self.$slab.insert(decl));
+                if let Some(parsed_decl_id) = parsed_decl_id {
+                    self.$parsed_slab
+                        .write()
+                        .insert(decl_id, parsed_decl_id.clone());
+                }
+                DeclRef::new(decl_name, decl_id, span)
             }
         }
         impl DeclEngineInsertArc<$decl> for DeclEngine {
-            fn insert_arc(&self, decl: Arc<$decl>) -> DeclRef<DeclId<$decl>> {
+            fn insert_arc(
+                &self,
+                decl: Arc<$decl>,
+                parsed_decl_id: Option<&ParsedDeclId<<$decl as TyDeclParsedType>::ParsedType>>,
+            ) -> DeclRef<DeclId<$decl>> {
                 let span = decl.span();
-                DeclRef::new(
-                    decl.name().clone(),
-                    DeclId::new(self.$slab.insert_arc(decl)),
-                    span,
-                )
+                let decl_name = decl.name().clone();
+                let decl_id = DeclId::new(self.$slab.insert_arc(decl));
+                if let Some(parsed_decl_id) = parsed_decl_id {
+                    self.$parsed_slab
+                        .write()
+                        .insert(decl_id, parsed_decl_id.clone());
+                }
+                DeclRef::new(decl_name, decl_id, span)
+            }
+        }
+    };
+}
+decl_engine_insert!(
+    function_slab,
+    function_parsed_decl_id_map,
+    ty::TyFunctionDecl
+);
+decl_engine_insert!(trait_slab, trait_parsed_decl_id_map, ty::TyTraitDecl);
+decl_engine_insert!(trait_fn_slab, trait_fn_parsed_decl_id_map, ty::TyTraitFn);
+decl_engine_insert!(
+    trait_type_slab,
+    trait_type_parsed_decl_id_map,
+    ty::TyTraitType
+);
+decl_engine_insert!(
+    impl_self_or_trait_slab,
+    impl_self_or_trait_parsed_decl_id_map,
+    ty::TyImplSelfOrTrait
+);
+decl_engine_insert!(struct_slab, struct_parsed_decl_id_map, ty::TyStructDecl);
+decl_engine_insert!(storage_slab, storage_parsed_decl_id_map, ty::TyStorageDecl);
+decl_engine_insert!(abi_slab, abi_parsed_decl_id_map, ty::TyAbiDecl);
+decl_engine_insert!(
+    constant_slab,
+    constant_parsed_decl_id_map,
+    ty::TyConstantDecl
+);
+decl_engine_insert!(
+    configurable_slab,
+    configurable_parsed_decl_id_map,
+    ty::TyConfigurableDecl
+);
+decl_engine_insert!(enum_slab, enum_parsed_decl_id_map, ty::TyEnumDecl);
+decl_engine_insert!(
+    type_alias_slab,
+    type_alias_parsed_decl_id_map,
+    ty::TyTypeAliasDecl
+);
+
+macro_rules! decl_engine_parsed_decl_id {
+    ($slab:ident, $decl:ty) => {
+        impl DeclEngineGetParsedDeclId<$decl> for DeclEngine {
+            fn get_parsed_decl_id(
+                &self,
+                decl_id: &DeclId<$decl>,
+            ) -> Option<ParsedDeclId<<$decl as TyDeclParsedType>::ParsedType>> {
+                let parsed_decl_id_map = self.$slab.read();
+                if let Some(parsed_decl_id) = parsed_decl_id_map.get(&decl_id) {
+                    return Some(parsed_decl_id.clone());
+                } else {
+                    None
+                }
             }
         }
     };
 }
-decl_engine_insert!(function_slab, ty::TyFunctionDecl);
-decl_engine_insert!(trait_slab, ty::TyTraitDecl);
-decl_engine_insert!(trait_fn_slab, ty::TyTraitFn);
-decl_engine_insert!(trait_type_slab, ty::TyTraitType);
-decl_engine_insert!(impl_self_or_trait_slab, ty::TyImplSelfOrTrait);
-decl_engine_insert!(struct_slab, ty::TyStructDecl);
-decl_engine_insert!(storage_slab, ty::TyStorageDecl);
-decl_engine_insert!(abi_slab, ty::TyAbiDecl);
-decl_engine_insert!(constant_slab, ty::TyConstantDecl);
-decl_engine_insert!(configurable_slab, ty::TyConfigurableDecl);
-decl_engine_insert!(enum_slab, ty::TyEnumDecl);
-decl_engine_insert!(type_alias_slab, ty::TyTypeAliasDecl);
+
+decl_engine_parsed_decl_id!(function_parsed_decl_id_map, ty::TyFunctionDecl);
+decl_engine_parsed_decl_id!(trait_parsed_decl_id_map, ty::TyTraitDecl);
+decl_engine_parsed_decl_id!(trait_fn_parsed_decl_id_map, ty::TyTraitFn);
+decl_engine_parsed_decl_id!(trait_type_parsed_decl_id_map, ty::TyTraitType);
+decl_engine_parsed_decl_id!(impl_self_or_trait_parsed_decl_id_map, ty::TyImplSelfOrTrait);
+decl_engine_parsed_decl_id!(struct_parsed_decl_id_map, ty::TyStructDecl);
+decl_engine_parsed_decl_id!(storage_parsed_decl_id_map, ty::TyStorageDecl);
+decl_engine_parsed_decl_id!(abi_parsed_decl_id_map, ty::TyAbiDecl);
+decl_engine_parsed_decl_id!(constant_parsed_decl_id_map, ty::TyConstantDecl);
+decl_engine_parsed_decl_id!(configurable_parsed_decl_id_map, ty::TyConfigurableDecl);
+decl_engine_parsed_decl_id!(enum_parsed_decl_id_map, ty::TyEnumDecl);
+decl_engine_parsed_decl_id!(type_alias_parsed_decl_id_map, ty::TyTypeAliasDecl);
 
 macro_rules! decl_engine_replace {
     ($slab:ident, $decl:ty) => {

sway-core/src/decl_engine/id.rs

@@ -5,7 +5,7 @@ use std::{fmt, hash::Hash};
 
 use sway_types::{Named, Spanned};
 
-use crate::language::ty::TyTraitType;
+use crate::language::ty::{TyDeclParsedType, TyTraitType};
 use crate::{
     decl_engine::*,
     engine_threading::*,
@@ -231,8 +231,8 @@ impl SubstTypes for DeclId<TyTraitType> {
 
 impl<T> DeclId<T>
 where
-    DeclEngine: DeclEngineIndex<T>,
-    T: Named + Spanned + SubstTypes + Clone,
+    DeclEngine: DeclEngineIndex<T> + DeclEngineInsert<T> + DeclEngineGetParsedDeclId<T>,
+    T: Named + Spanned + SubstTypes + Clone + TyDeclParsedType,
 {
     pub(crate) fn subst_types_and_insert_new(
         &self,
@@ -242,7 +242,7 @@ where
         let decl_engine = engines.de();
         let mut decl = (*decl_engine.get(self)).clone();
         if decl.subst(type_mapping, engines).has_changes() {
-            Some(decl_engine.insert(decl))
+            Some(decl_engine.insert(decl, decl_engine.get_parsed_decl_id(self).as_ref()))
         } else {
             None
         }

sway-core/src/decl_engine/ref.rs

@@ -29,8 +29,8 @@ use crate::{
     decl_engine::*,
     engine_threading::*,
     language::ty::{
-        self, TyAbiDecl, TyConstantDecl, TyEnumDecl, TyFunctionDecl, TyImplSelfOrTrait,
-        TyStorageDecl, TyStructDecl, TyTraitDecl, TyTraitFn, TyTraitType,
+        self, TyAbiDecl, TyConstantDecl, TyDeclParsedType, TyEnumDecl, TyFunctionDecl,
+        TyImplSelfOrTrait, TyStorageDecl, TyStructDecl, TyTraitDecl, TyTraitFn, TyTraitType,
     },
     semantic_analysis::TypeCheckContext,
     type_system::*,
@@ -96,8 +96,8 @@ impl<T> DeclRef<DeclId<T>> {
 
 impl<T> DeclRef<DeclId<T>>
 where
-    DeclEngine: DeclEngineIndex<T>,
-    T: Named + Spanned + SubstTypes + Clone,
+    DeclEngine: DeclEngineIndex<T> + DeclEngineInsert<T> + DeclEngineGetParsedDeclId<T>,
+    T: Named + Spanned + SubstTypes + Clone + TyDeclParsedType,
 {
     pub(crate) fn subst_types_and_insert_new(
         &self,
@@ -107,7 +107,7 @@ where
         let decl_engine = engines.de();
         let mut decl = (*decl_engine.get(&self.id)).clone();
         if decl.subst(type_mapping, engines).has_changes() {
-            Some(decl_engine.insert(decl))
+            Some(decl_engine.insert(decl, decl_engine.get_parsed_decl_id(&self.id).as_ref()))
         } else {
             None
         }
@@ -132,8 +132,8 @@ where
 impl<T> DeclRef<DeclId<T>>
 where
     AssociatedItemDeclId: From<DeclId<T>>,
-    DeclEngine: DeclEngineIndex<T>,
-    T: Named + Spanned + SubstTypes + Clone,
+    DeclEngine: DeclEngineIndex<T> + DeclEngineInsert<T> + DeclEngineGetParsedDeclId<T>,
+    T: Named + Spanned + SubstTypes + Clone + TyDeclParsedType,
 {
     pub(crate) fn subst_types_and_insert_new_with_parent(
         &self,
@@ -145,7 +145,7 @@ where
         if decl.subst(type_mapping, engines).has_changes() {
             Some(
                 decl_engine
-                    .insert(decl)
+                    .insert(decl, decl_engine.get_parsed_decl_id(&self.id).as_ref())
                     .with_parent(decl_engine, self.id.into()),
             )
         } else {

sway-core/src/language/ty/declaration/abi.rs

@@ -1,9 +1,14 @@
-use crate::{engine_threading::*, language::parsed, transform, type_system::*};
+use crate::{
+    engine_threading::*,
+    language::parsed::{self, AbiDeclaration},
+    transform,
+    type_system::*,
+};
 use std::hash::{Hash, Hasher};
 
 use sway_types::{Ident, Named, Span, Spanned};
 
-use super::{TyTraitInterfaceItem, TyTraitItem};
+use super::{TyDeclParsedType, TyTraitInterfaceItem, TyTraitItem};
 
 /// A [TyAbiDecl] contains the type-checked version of the parse tree's
 /// `AbiDeclaration`.
@@ -19,6 +24,10 @@ pub struct TyAbiDecl {
     pub attributes: transform::AttributesMap,
 }
 
+impl TyDeclParsedType for TyAbiDecl {
+    type ParsedType = AbiDeclaration;
+}
+
 impl EqWithEngines for TyAbiDecl {}
 impl PartialEqWithEngines for TyAbiDecl {
     fn eq(&self, other: &Self, ctx: &PartialEqWithEnginesContext) -> bool {

sway-core/src/language/ty/declaration/configurable.rs

@@ -10,7 +10,7 @@ use crate::{
     decl_engine::{DeclId, DeclMapping, DeclRef, ReplaceDecls},
     engine_threading::*,
     has_changes,
-    language::{ty::*, CallPath, Visibility},
+    language::{parsed::ConfigurableDeclaration, ty::*, CallPath, Visibility},
     semantic_analysis::TypeCheckContext,
     transform,
     type_system::*,
@@ -29,6 +29,10 @@ pub struct TyConfigurableDecl {
     pub decode_fn: Option<DeclRef<DeclId<TyFunctionDecl>>>,
 }
 
+impl TyDeclParsedType for TyConfigurableDecl {
+    type ParsedType = ConfigurableDeclaration;
+}
+
 impl DebugWithEngines for TyConfigurableDecl {
     fn fmt(&self, f: &mut fmt::Formatter<'_>, _engines: &Engines) -> fmt::Result {
         write!(f, "{}", self.call_path)