Add AST module caching support for LSP (#5251)

## Description
This PR supersedes #4988 and builds on backend work done by @tritao in
#4733.

A new `ModuleId` is created for each compiled module. The `SourceId`
type now also stores the `ModuleId` that was used to create it. This
allows us to perform garbage collection on types in the Declaration and
Type engines that refer the provided module.

```rust
if let Some(module_id) = engines.se().get_module_id(&path) {
    engines.clear_module(&module_id);
}
```

This method isn't super cheap to call. My tests had this averaging at
54ms. We are currently performing garbage collection on every key stroke
in LSP. We should look to have another metric such as only clearing once
every 20 keystrokes or so in the future.

I ran benchmarking on recompiling the `doc_comments` example 20 times to
simulate rapid keystrokes during coding. Here are the results.

| Without Caching | With Caching |
|---------------------|------------------|
| 89.489 s             | 4.6236  s         |

A speed up of 1835.49% is certainly great, yet 4.6 seconds still sounds
like an eternity. Still have lots of room to improve but this should
improve our current system dramatically.

closes #4994

Co-authored-by: Joao Matos <[email protected]>

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 {