Add AST module caching support for LSP

sway-core/src/concurrent_slab.rs

@@ -1,9 +1,7 @@
+use crate::{decl_engine::*, engine_threading::*, type_system::*};
 use std::{fmt, sync::RwLock};
-
 use sway_types::{Named, Spanned};
 
-use crate::{decl_engine::*, engine_threading::*, type_system::*};
-
 #[derive(Debug)]
 pub(crate) struct ConcurrentSlab<T> {
     inner: RwLock<Vec<T>>,
@@ -70,16 +68,21 @@ where
         let inner = self.inner.read().unwrap();
         inner[index].clone()
     }
+
+    pub fn retain(&self, predicate: impl Fn(&T) -> bool) {
+        let mut inner = self.inner.write().unwrap();
+        inner.retain(predicate);
+    }
 }
 
-impl ConcurrentSlab<TypeInfo> {
+impl ConcurrentSlab<TypeSourceInfo> {
     pub fn replace(
         &self,
         index: TypeId,
-        prev_value: &TypeInfo,
-        new_value: TypeInfo,
+        prev_value: &TypeSourceInfo,
+        new_value: TypeSourceInfo,
         engines: &Engines,
-    ) -> Option<TypeInfo> {
+    ) -> Option<TypeSourceInfo> {
         let index = index.index();
         // The comparison below ends up calling functions in the slab, which
         // can lead to deadlocks if we used a single read/write lock.
@@ -89,7 +92,10 @@ impl ConcurrentSlab<TypeInfo> {
         {
             let inner = self.inner.read().unwrap();
             let actual_prev_value = &inner[index];
-            if !actual_prev_value.eq(prev_value, engines) {
+            if !actual_prev_value
+                .type_info
+                .eq(&prev_value.type_info, engines)
+            {
                 return Some(actual_prev_value.clone());
             }
         }

sway-core/src/decl_engine/engine.rs

@@ -4,7 +4,7 @@ use std::{
     sync::RwLock,
 };
 
-use sway_types::{Named, Spanned};
+use sway_types::{ModuleId, Named, Spanned};
 
 use crate::{
     concurrent_slab::{ConcurrentSlab, ListDisplay},
@@ -147,6 +147,35 @@ decl_engine_index!(constant_slab, ty::TyConstantDecl);
 decl_engine_index!(enum_slab, ty::TyEnumDecl);
 decl_engine_index!(type_alias_slab, ty::TyTypeAliasDecl);
 
+macro_rules! decl_engine_clear_module {
+    ($($slab:ident, $decl:ty);* $(;)?) => {
+        impl DeclEngine {
+            pub fn clear_module(&mut self, module_id: &ModuleId) {
+                $(
+                    self.$slab.retain(|ty| match ty.span().source_id() {
+                        Some(source_id) => &source_id.module_id() != module_id,
+                        None => false,
+                    });
+                )*
+            }
+        }
+    };
+}
+
+decl_engine_clear_module!(
+    function_slab, ty::TyFunctionDecl;
+    trait_slab, ty::TyTraitDecl;
+    trait_fn_slab, ty::TyTraitFn;
+    trait_type_slab, ty::TyTraitType;
+    impl_trait_slab, ty::TyImplTrait;
+    struct_slab, ty::TyStructDecl;
+    storage_slab, ty::TyStorageDecl;
+    abi_slab, ty::TyAbiDecl;
+    constant_slab, ty::TyConstantDecl;
+    enum_slab, ty::TyEnumDecl;
+    type_alias_slab, ty::TyTypeAliasDecl;
+);
+
 impl DeclEngine {
     /// Given a [DeclRef] `index`, finds all the parents of `index` and all the
     /// recursive parents of those parents, and so on. Does not perform

sway-core/src/engine_threading.rs

@@ -1,13 +1,11 @@
+use crate::{decl_engine::DeclEngine, query_engine::QueryEngine, type_system::TypeEngine};
 use std::{
     cmp::Ordering,
     fmt,
     hash::{BuildHasher, Hash, Hasher},
 };
-
 use sway_types::SourceEngine;
 
-use crate::{decl_engine::DeclEngine, query_engine::QueryEngine, type_system::TypeEngine};
-
 #[derive(Debug, Default)]
 pub struct Engines {
     type_engine: TypeEngine,
@@ -47,6 +45,13 @@ impl Engines {
         &self.source_engine
     }
 
+    /// Removes all data associated with `module_id` from the declaration and type engines.
+    /// It is intended to be used during garbage collection to remove any data that is no longer needed.
+    pub fn clear_module(&mut self, module_id: &sway_types::ModuleId) {
+        self.type_engine.clear_module(module_id);
+        self.decl_engine.clear_module(module_id);
+    }
+
     /// Helps out some `thing: T` by adding `self` as context.
     pub fn help_out<T>(&self, thing: T) -> WithEngines<'_, T> {
         WithEngines {

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

@@ -72,7 +72,7 @@ impl CreateTypeId for TyAbiDecl {
             abi_name: AbiName::Known(self.name.clone().into()),
             address: None,
         };
-        type_engine.insert(engines, ty)
+        type_engine.insert(engines, ty, self.name.span().source_id())
     }
 }
 

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

@@ -784,6 +784,7 @@ impl TyDecl {
             }) => type_engine.insert(
                 engines,
                 TypeInfo::Struct(DeclRef::new(name.clone(), *decl_id, decl_span.clone())),
+                name.span().source_id(),
             ),
             TyDecl::EnumDecl(EnumDecl {
                 name,
@@ -793,6 +794,7 @@ impl TyDecl {
             }) => type_engine.insert(
                 engines,
                 TypeInfo::Enum(DeclRef::new(name.clone(), *decl_id, decl_span.clone())),
+                name.span().source_id(),
             ),
             TyDecl::StorageDecl(StorageDecl { decl_id, .. }) => {
                 let storage_decl = decl_engine.get_storage(decl_id);
@@ -801,6 +803,7 @@ impl TyDecl {
                     TypeInfo::Storage {
                         fields: storage_decl.fields_as_typed_struct_fields(),
                     },
+                    storage_decl.span().source_id(),
                 )
             }
             TyDecl::TypeAliasDecl(TypeAliasDecl { decl_id, .. }) => {

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

@@ -67,6 +67,7 @@ impl CreateTypeId for TyTypeAliasDecl {
                 name: self.name.clone(),
                 ty: self.ty.clone(),
             },
+            self.name.span().source_id(),
         )
     }
 }

sway-core/src/language/ty/expression/expression.rs

@@ -414,7 +414,7 @@ impl TyExpression {
         let type_engine = engines.te();
         TyExpression {
             expression: TyExpressionVariant::Tuple { fields: vec![] },
-            return_type: type_engine.insert(engines, TypeInfo::ErrorRecovery(err)),
+            return_type: type_engine.insert(engines, TypeInfo::ErrorRecovery(err), None),
             span,
         }
     }

sway-core/src/semantic_analysis/ast_node/code_block.rs

@@ -100,14 +100,15 @@ impl ty::TyCodeBlock {
                         return ctx.engines().te().insert(
                             engines,
                             TypeInfo::Enum(DeclRef::new(name.clone(), decl_id, decl_span.clone())),
+                            name.span().source_id(),
                         );
                     }
 
-                    ctx.engines.te().insert(engines, TypeInfo::Unknown)
+                    ctx.engines.te().insert(engines, TypeInfo::Unknown, None)
                 } else {
                     ctx.engines
                         .te()
-                        .insert(engines, TypeInfo::Tuple(Vec::new()))
+                        .insert(engines, TypeInfo::Tuple(Vec::new()), span.source_id())
                 }
             });
         (block_type, span)

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

@@ -41,7 +41,7 @@ impl ty::TyConstantDecl {
                 EnforceTypeArguments::No,
                 None,
             )
-            .unwrap_or_else(|err| type_engine.insert(engines, TypeInfo::ErrorRecovery(err)));
+            .unwrap_or_else(|err| type_engine.insert(engines, TypeInfo::ErrorRecovery(err), None));
 
         // this subst is required to replace associated types, namely TypeInfo::TraitType.
         type_ascription.type_id.subst(&ctx.type_subst(), engines);
@@ -121,7 +121,7 @@ impl ty::TyConstantDecl {
             call_path,
             span,
             attributes: Default::default(),
-            return_type: type_engine.insert(engines, TypeInfo::Unknown),
+            return_type: type_engine.insert(engines, TypeInfo::Unknown, None),
             type_ascription,
             is_configurable: false,
             value: None,

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

@@ -42,7 +42,7 @@ impl TyDecl {
                         None,
                     )
                     .unwrap_or_else(|err| {
-                        type_engine.insert(engines, TypeInfo::ErrorRecovery(err))
+                        type_engine.insert(engines, TypeInfo::ErrorRecovery(err), None)
                     });
                 let mut ctx = ctx
                     .with_type_annotation(type_ascription.type_id)
@@ -106,7 +106,7 @@ impl TyDecl {
             parsed::Declaration::FunctionDeclaration(fn_decl) => {
                 let span = fn_decl.span.clone();
                 let mut ctx =
-                    ctx.with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown));
+                    ctx.with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown, None));
                 let fn_decl = match ty::TyFunctionDecl::type_check(
                     handler,
                     ctx.by_ref(),
@@ -350,7 +350,7 @@ impl TyDecl {
                 let new_ty = ctx
                     .resolve_type(handler, ty.type_id, &span, EnforceTypeArguments::Yes, None)
                     .unwrap_or_else(|err| {
-                        type_engine.insert(engines, TypeInfo::ErrorRecovery(err))
+                        type_engine.insert(engines, TypeInfo::ErrorRecovery(err), None)
                     });
 
                 // create the type alias decl using the resolved type above

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

@@ -1,10 +1,9 @@
-use sway_error::handler::{ErrorEmitted, Handler};
-
 use crate::{
     language::{parsed::*, ty, CallPath},
     semantic_analysis::{type_check_context::EnforceTypeArguments, *},
     type_system::*,
 };
+use sway_error::handler::{ErrorEmitted, Handler};
 
 impl ty::TyEnumDecl {
     pub fn type_check(
@@ -74,7 +73,7 @@ impl ty::TyEnumVariant {
                 EnforceTypeArguments::Yes,
                 None,
             )
-            .unwrap_or_else(|err| type_engine.insert(engines, TypeInfo::ErrorRecovery(err)));
+            .unwrap_or_else(|err| type_engine.insert(engines, TypeInfo::ErrorRecovery(err), None));
         Ok(ty::TyEnumVariant {
             name: variant.name.clone(),
             type_argument,

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

@@ -115,7 +115,7 @@ impl ty::TyFunctionDecl {
                 EnforceTypeArguments::Yes,
                 None,
             )
-            .unwrap_or_else(|err| type_engine.insert(engines, TypeInfo::ErrorRecovery(err)));
+            .unwrap_or_else(|err| type_engine.insert(engines, TypeInfo::ErrorRecovery(err), None));
 
         let (visibility, is_contract_call) = if is_method {
             if is_in_impl_self {
@@ -345,6 +345,7 @@ fn test_function_selector_behavior() {
                     .insert(
                         &engines,
                         TypeInfo::StringArray(Length::new(5, Span::dummy())),
+                        None,
                     )
                     .into(),
             },
@@ -354,12 +355,15 @@ fn test_function_selector_behavior() {
                 is_mutable: false,
                 mutability_span: Span::dummy(),
                 type_argument: TypeArgument {
-                    type_id: engines
-                        .te()
-                        .insert(&engines, TypeInfo::UnsignedInteger(IntegerBits::ThirtyTwo)),
+                    type_id: engines.te().insert(
+                        &engines,
+                        TypeInfo::UnsignedInteger(IntegerBits::ThirtyTwo),
+                        None,
+                    ),
                     initial_type_id: engines.te().insert(
                         &engines,
                         TypeInfo::StringArray(Length::new(5, Span::dummy())),
+                        None,
                     ),
                     span: Span::dummy(),
                     call_path_tree: None,

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

@@ -35,7 +35,7 @@ impl ty::TyFunctionParameter {
                 EnforceTypeArguments::Yes,
                 None,
             )
-            .unwrap_or_else(|err| type_engine.insert(engines, TypeInfo::ErrorRecovery(err)));
+            .unwrap_or_else(|err| type_engine.insert(engines, TypeInfo::ErrorRecovery(err), None));
 
         type_argument.type_id.check_type_parameter_bounds(
             handler,
@@ -89,7 +89,7 @@ impl ty::TyFunctionParameter {
                 EnforceTypeArguments::Yes,
                 None,
             )
-            .unwrap_or_else(|err| type_engine.insert(engines, TypeInfo::ErrorRecovery(err)));
+            .unwrap_or_else(|err| type_engine.insert(engines, TypeInfo::ErrorRecovery(err), None));
 
         let typed_parameter = ty::TyFunctionParameter {
             name,

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

@@ -115,7 +115,7 @@ impl TyImplTrait {
         // Update the context
         let mut ctx = ctx
             .with_help_text("")
-            .with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown))
+            .with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown, None))
             .with_self_type(Some(implementing_for.type_id));
 
         let impl_trait = match ctx
@@ -362,7 +362,7 @@ impl TyImplTrait {
 
         let mut ctx = ctx
             .with_help_text("")
-            .with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown));
+            .with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown, None));
 
         // type check the items inside of the impl block
         let mut new_items = vec![];
@@ -778,14 +778,23 @@ fn type_check_trait_implementation(
                             name: Ident::new_with_override("Self".into(), Span::dummy()),
                             trait_constraints: VecSet(vec![]),
                         },
+                        None,
                     ),
                 };
                 trait_type_mapping.extend(TypeSubstMap::from_type_parameters_and_type_arguments(
-                    vec![type_engine.insert(engines, old_type_decl_info1)],
+                    vec![type_engine.insert(
+                        engines,
+                        old_type_decl_info1,
+                        type_decl.name.span().source_id(),
+                    )],
                     vec![type_decl.ty.clone().unwrap().type_id],
                 ));
                 trait_type_mapping.extend(TypeSubstMap::from_type_parameters_and_type_arguments(
-                    vec![type_engine.insert(engines, old_type_decl_info2)],
+                    vec![type_engine.insert(
+                        engines,
+                        old_type_decl_info2,
+                        type_decl.name.span().source_id(),
+                    )],
                     vec![type_decl.ty.clone().unwrap().type_id],
                 ));
             }
@@ -965,7 +974,7 @@ fn type_check_impl_method(
     let mut ctx = ctx
         .by_ref()
         .with_help_text("")
-        .with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown));
+        .with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown, None));
 
     let interface_name = || -> InterfaceName {
         if is_contract {
@@ -1188,7 +1197,7 @@ fn type_check_const_decl(
     let mut ctx = ctx
         .by_ref()
         .with_help_text("")
-        .with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown));
+        .with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown, None));
 
     let interface_name = || -> InterfaceName {
         if is_contract {
@@ -1270,7 +1279,7 @@ fn type_check_type_decl(
     let mut ctx = ctx
         .by_ref()
         .with_help_text("")
-        .with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown));
+        .with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown, None));
 
     let interface_name = || -> InterfaceName {
         if is_contract {