Implement Typed Module Caching (#6280)

## Description
This PR implements a typed module caching system for the compiler when
interacted with from the language server. The main goal is to improve
performance by caching typed modules and reusing them when possible,
reducing unnecessary recompilation.

## Key Changes
- Introduced `CowCache` for efficient, thread-safe caching with
copy-on-write semantics
- Implemented typed module caching in the `QueryEngine`
- Updated `ServerState` and `Session` to utilize the new caching system


## Performance Improvements
These changes show promising performance improvements in processing
didChange events, particularly for larger projects. Here are some
benchmarks using the FUSD libraries Sway project (19 Sway files):

| Build | Before | After | Improvement |
|-------|--------|-------|-------------|
| Debug | 448.25ms | 123.73ms | 72.4% faster |
| Release | 103.83ms | 53.59ms | 48.4% faster |

These improvements should lead to a more responsive development
experience, especially when making frequent small changes.

closes #6228
improves: FuelLabs/sway-vscode-plugin#172

## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [x] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [x] If my change requires substantial documentation changes, I have
[requested support from the DevRel
team](https://github.com/FuelLabs/devrel-requests/issues/new/choose)
- [x] I have added tests that prove my fix is effective or that my
feature works.
- [x] I have added (or requested a maintainer to add) the necessary
`Breaking*` or `New Feature` labels where relevant.
- [x] I have done my best to ensure that my PR adheres to [the Fuel Labs
Code Review
Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md).
- [x] I have requested a review from the relevant team or maintainers.

sway-core/src/decl_engine/engine.rs

@@ -5,7 +5,7 @@ use std::{
     sync::Arc,
 };
 
-use sway_types::{Named, ProgramId, Spanned};
+use sway_types::{Named, ProgramId, SourceId, Spanned};
 
 use crate::{
     concurrent_slab::ConcurrentSlab,
@@ -459,6 +459,53 @@ decl_engine_clear_program!(
     type_alias_slab, ty::TyTypeAliasDecl;
 );
 
+macro_rules! decl_engine_clear_module {
+    ($($slab:ident, $decl:ty);* $(;)?) => {
+        impl DeclEngine {
+            pub fn clear_module(&mut self, source_id: &SourceId) {
+                self.parents.write().retain(|key, _| {
+                    match key {
+                        AssociatedItemDeclId::TraitFn(decl_id) => {
+                            self.get_trait_fn(decl_id).span().source_id().map_or(true, |src_id| src_id != source_id)
+                        },
+                        AssociatedItemDeclId::Function(decl_id) => {
+                            self.get_function(decl_id).span().source_id().map_or(true, |src_id| src_id != source_id)
+                        },
+                        AssociatedItemDeclId::Type(decl_id) => {
+                            self.get_type(decl_id).span().source_id().map_or(true, |src_id| src_id != source_id)
+                        },
+                        AssociatedItemDeclId::Constant(decl_id) => {
+                            self.get_constant(decl_id).span().source_id().map_or(true, |src_id| src_id != source_id)
+                        },
+                    }
+                });
+
+                $(
+                    self.$slab.retain(|_k, ty| match ty.span().source_id() {
+                        Some(src_id) => src_id != source_id,
+                        None => true,
+                    });
+                )*
+            }
+        }
+    };
+}
+
+decl_engine_clear_module!(
+    function_slab, ty::TyFunctionDecl;
+    trait_slab, ty::TyTraitDecl;
+    trait_fn_slab, ty::TyTraitFn;
+    trait_type_slab, ty::TyTraitType;
+    impl_self_or_trait_slab, ty::TyImplTrait;
+    struct_slab, ty::TyStructDecl;
+    storage_slab, ty::TyStorageDecl;
+    abi_slab, ty::TyAbiDecl;
+    constant_slab, ty::TyConstantDecl;
+    configurable_slab, ty::TyConfigurableDecl;
+    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/decl_engine/parsed_engine.rs

@@ -9,7 +9,7 @@ use crate::{
 };
 
 use std::sync::Arc;
-use sway_types::{ProgramId, Spanned};
+use sway_types::{ProgramId, SourceId, Spanned};
 
 use super::parsed_id::ParsedDeclId;
 
@@ -191,6 +191,46 @@ decl_engine_clear_program!(
         .span()),
 );
 
+macro_rules! decl_engine_clear_module {
+    ($(($slab:ident, $getter:expr)),* $(,)?) => {
+        impl ParsedDeclEngine {
+            pub fn clear_module(&mut self, program_id: &SourceId) {
+                $(
+                    self.$slab.retain(|_k, item| {
+                        #[allow(clippy::redundant_closure_call)]
+                        let span = $getter(item);
+                        match span.source_id() {
+                            Some(src_id) => src_id != program_id,
+                            None => true,
+                        }
+                    });
+                )*
+            }
+        }
+    };
+}
+
+decl_engine_clear_module!(
+    (variable_slab, |item: &VariableDeclaration| item.name.span()),
+    (function_slab, |item: &FunctionDeclaration| item.name.span()),
+    (trait_slab, |item: &TraitDeclaration| item.name.span()),
+    (trait_fn_slab, |item: &TraitFn| item.name.span()),
+    (trait_type_slab, |item: &TraitTypeDeclaration| item
+        .name
+        .span()),
+    (impl_self_or_trait_slab, |item: &ImplSelfOrTrait| item
+        .block_span
+        .clone()),
+    (struct_slab, |item: &StructDeclaration| item.name.span()),
+    (storage_slab, |item: &StorageDeclaration| item.span.clone()),
+    (abi_slab, |item: &AbiDeclaration| item.name.span()),
+    (constant_slab, |item: &ConstantDeclaration| item.name.span()),
+    (enum_slab, |item: &EnumDeclaration| item.name.span()),
+    (type_alias_slab, |item: &TypeAliasDeclaration| item
+        .name
+        .span()),
+);
+
 impl ParsedDeclEngine {
     /// Friendly helper method for calling the `get` method from the
     /// implementation of [ParsedDeclEngineGet] for [ParsedDeclEngine]

sway-core/src/engine_threading.rs

@@ -48,6 +48,14 @@ impl Engines {
         self.parsed_decl_engine.clear_program(program_id);
     }
 
+    /// Removes all data associated with `source_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, source_id: &sway_types::SourceId) {
+        self.type_engine.clear_module(source_id);
+        self.decl_engine.clear_module(source_id);
+        self.parsed_decl_engine.clear_module(source_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/module.rs

@@ -75,7 +75,7 @@ impl TyModule {
 
 #[derive(Clone, Debug)]
 pub struct TySubmodule {
-    pub module: TyModule,
+    pub module: Arc<TyModule>,
     pub mod_name_span: Span,
 }
 

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

@@ -497,7 +497,7 @@ impl CollectTypesMetadata for TyProgram {
                 for module in std::iter::once(&self.root).chain(
                     self.root
                         .submodules_recursive()
-                        .map(|(_, submod)| &submod.module),
+                        .map(|(_, submod)| &*submod.module),
                 ) {
                     for node in module.all_nodes.iter() {
                         let is_generic_function = node.is_generic_function(decl_engine);
@@ -526,7 +526,7 @@ impl CollectTypesMetadata for TyProgram {
         for module in std::iter::once(&self.root).chain(
             self.root
                 .submodules_recursive()
-                .map(|(_, submod)| &submod.module),
+                .map(|(_, submod)| &*submod.module),
         ) {
             for node in module.all_nodes.iter() {
                 if node.is_test_function(decl_engine) {

sway-core/src/lib.rs

@@ -33,7 +33,7 @@ use control_flow_analysis::ControlFlowGraph;
 pub use debug_generation::write_dwarf;
 use indexmap::IndexMap;
 use metadata::MetadataManager;
-use query_engine::{ModuleCacheKey, ModulePath, ProgramsCacheEntry};
+use query_engine::{ModuleCacheKey, ModuleCommonInfo, ParsedModuleInfo, ProgramsCacheEntry};
 use std::collections::hash_map::DefaultHasher;
 use std::hash::{Hash, Hasher};
 use std::path::{Path, PathBuf};
@@ -236,7 +236,7 @@ fn parse_in_memory(
 
 pub struct Submodule {
     name: Ident,
-    path: ModulePath,
+    path: Arc<PathBuf>,
     lexed: lexed::LexedSubmodule,
     parsed: parsed::ParseSubmodule,
 }
@@ -259,7 +259,6 @@ fn parse_submodules(
 ) -> Submodules {
     // Assume the happy path, so there'll be as many submodules as dependencies, but no more.
     let mut submods = Vec::with_capacity(module.submodules().count());
-
     module.submodules().for_each(|submod| {
         // Read the source code from the dependency.
         // If we cannot, record as an error, but continue with other files.
@@ -275,7 +274,6 @@ fn parse_submodules(
                 return;
             }
         };
-
         if let Ok(ParsedModuleTree {
             tree_type: kind,
             lexed_module,
@@ -318,7 +316,6 @@ fn parse_submodules(
             submods.push(submodule);
         }
     });
-
     submods
 }
 
@@ -411,15 +408,19 @@ fn parse_module_tree(
     let version = lsp_mode
         .and_then(|lsp| lsp.file_versions.get(path.as_ref()).copied())
         .unwrap_or(None);
-    let cache_entry = ModuleCacheEntry {
-        path,
-        modified_time,
-        hash,
-        dependencies,
+
+    let common_info = ModuleCommonInfo {
+        path: path.clone(),
         include_tests,
+        dependencies,
+        hash,
+    };
+    let parsed_info = ParsedModuleInfo {
+        modified_time,
         version,
     };
-    query_engine.insert_parse_module_cache_entry(cache_entry);
+    let cache_entry = ModuleCacheEntry::new(common_info, parsed_info);
+    query_engine.update_or_insert_parsed_module_cache_entry(cache_entry);
 
     Ok(ParsedModuleTree {
         tree_type: kind,
@@ -428,59 +429,89 @@ fn parse_module_tree(
     })
 }
 
-fn is_parse_module_cache_up_to_date(
+/// Checks if the typed module cache for a given path is up to date.
+///
+/// This function determines whether the cached typed representation of a module
+/// is still valid based on file versions and dependencies.
+///
+/// Note: This functionality is currently only supported when the compiler is
+/// initiated from the language server.
+pub(crate) fn is_ty_module_cache_up_to_date(
     engines: &Engines,
     path: &Arc<PathBuf>,
     include_tests: bool,
     build_config: Option<&BuildConfig>,
 ) -> bool {
-    let query_engine = engines.qe();
+    let cache = engines.qe().module_cache.read();
     let key = ModuleCacheKey::new(path.clone(), include_tests);
-    let entry = query_engine.get_parse_module_cache_entry(&key);
-    match entry {
-        Some(entry) => {
-            // Let's check if we can re-use the dependency information
-            // we got from the cache.
+    cache.get(&key).map_or(false, |entry| {
+        entry.typed.as_ref().map_or(false, |typed| {
+            // Check if the cache is up to date based on file versions
             let cache_up_to_date = build_config
-                .as_ref()
                 .and_then(|x| x.lsp_mode.as_ref())
-                .and_then(|lsp| {
-                    // First try to get the file version from lsp if it exists
-                    lsp.file_versions.get(path.as_ref())
-                })
-                .map_or_else(
-                    || {
-                        // Otherwise we can safely read the file from disk here, as the LSP has not modified it, or we are not in LSP mode.
-                        // Check if the file has been modified or if its hash is the same as the last compilation
-                        let modified_time = std::fs::metadata(path.as_path())
-                            .ok()
-                            .and_then(|m| m.modified().ok());
-                        entry.modified_time == modified_time || {
-                            let src = std::fs::read_to_string(path.as_path()).unwrap();
-                            let mut hasher = DefaultHasher::new();
-                            src.hash(&mut hasher);
-                            let hash = hasher.finish();
-                            hash == entry.hash
-                        }
-                    },
-                    |version| {
-                        // The cache is invalid if the lsp version is greater than the last compilation
-                        !version.map_or(false, |v| v > entry.version.unwrap_or(0))
-                    },
-                );
-
-            // Look at the dependencies recursively to make sure they have not been
-            // modified either.
-            if cache_up_to_date {
-                entry.dependencies.iter().all(|path| {
-                    is_parse_module_cache_up_to_date(engines, path, include_tests, build_config)
+                .and_then(|lsp| lsp.file_versions.get(path.as_ref()))
+                .map_or(true, |version| {
+                    version.map_or(true, |v| typed.version.map_or(false, |tv| v <= tv))
+                });
+
+            // If the cache is up to date, recursively check all dependencies
+            cache_up_to_date
+                && entry.common.dependencies.iter().all(|dep_path| {
+                    is_ty_module_cache_up_to_date(engines, dep_path, include_tests, build_config)
                 })
-            } else {
-                false
-            }
-        }
-        None => false,
-    }
+        })
+    })
+}
+
+/// Checks if the parsed module cache for a given path is up to date.
+///
+/// This function determines whether the cached parsed representation of a module
+/// is still valid based on file versions, modification times, or content hashes.
+pub(crate) fn is_parse_module_cache_up_to_date(
+    engines: &Engines,
+    path: &Arc<PathBuf>,
+    include_tests: bool,
+    build_config: Option<&BuildConfig>,
+) -> bool {
+    let cache = engines.qe().module_cache.read();
+    let key = ModuleCacheKey::new(path.clone(), include_tests);
+    cache.get(&key).map_or(false, |entry| {
+        // Determine if the cached dependency information is still valid
+        let cache_up_to_date = build_config
+            .and_then(|x| x.lsp_mode.as_ref())
+            .and_then(|lsp| lsp.file_versions.get(path.as_ref()))
+            .map_or_else(
+                || {
+                    // If LSP mode is not active or file version is unavailable, fall back to filesystem checks.
+                    let modified_time = std::fs::metadata(path.as_path())
+                        .ok()
+                        .and_then(|m| m.modified().ok());
+                    // Check if modification time matches, or if not, compare file content hash
+                    entry.parsed.modified_time == modified_time || {
+                        let src = std::fs::read_to_string(path.as_path()).unwrap();
+                        let mut hasher = DefaultHasher::new();
+                        src.hash(&mut hasher);
+                        hasher.finish() == entry.common.hash
+                    }
+                },
+                |version| {
+                    // Determine if the parse cache is up-to-date in LSP mode:
+                    // - If there's no LSP file version (version is None), consider the cache up-to-date.
+                    // - If there is an LSP file version:
+                    //   - If there's no cached version (entry.parsed.version is None), the cache is outdated.
+                    //   - If there's a cached version, compare them: cache is up-to-date if the LSP file version
+                    //     is not greater than the cached version.
+                    version.map_or(true, |v| entry.parsed.version.map_or(false, |ev| v <= ev))
+                },
+            );
+
+        // Checks if the typed module cache for a given path is up to date// If the cache is up to date, recursively check all dependencies to make sure they have not been
+        // modified either.
+        cache_up_to_date
+            && entry.common.dependencies.iter().all(|dep_path| {
+                is_parse_module_cache_up_to_date(engines, dep_path, include_tests, build_config)
+            })
+    })
 }
 
 fn module_path(
@@ -697,12 +728,12 @@ pub fn compile_to_ast(
     retrigger_compilation: Option<Arc<AtomicBool>>,
 ) -> Result<Programs, ErrorEmitted> {
     check_should_abort(handler, retrigger_compilation.clone())?;
+
     let query_engine = engines.qe();
     let mut metrics = PerformanceData::default();
     if let Some(config) = build_config {
         let path = config.canonical_root_module();
         let include_tests = config.include_tests;
-
         // Check if we can re-use the data in the cache.
         if is_parse_module_cache_up_to_date(engines, &path, include_tests, build_config) {
             let mut entry = query_engine.get_programs_cache_entry(&path).unwrap();