Fix function_cache garbage collection bug

.vscode/launch.json

@@ -0,0 +1,23 @@
+// Configuration for debugging the Sway Language Server (forc-lsp)
+// Usage instructions:
+// 1. Ensure you've built forc-lsp with debug symbols:
+//    cargo build -p forc-lsp
+// 2. Install the debug version:
+//    cargo install --path ./forc-plugins/forc-lsp --debug
+// 3. Open your Sway project in a separate VSCode window (this starts forc-lsp)
+// 4. In the forc-lsp project window, set breakpoints in the code
+// 5. Go to Run and Debug view, select "Attach to forc-lsp", and start debugging
+// 6. When prompted, select the forc-lsp process
+// 7. Debug forc-lsp as it responds to actions in your Sway project
+{
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "type": "lldb",
+            "request": "attach",
+            "name": "Attach to forc-lsp",
+            "pid": "${command:pickProcess}",
+            "program": "${env:HOME}/.cargo/bin/forc-lsp"
+        }
+    ]
+}

sway-core/src/engine_threading.rs

@@ -46,6 +46,7 @@ impl Engines {
         self.type_engine.clear_program(program_id);
         self.decl_engine.clear_program(program_id);
         self.parsed_decl_engine.clear_program(program_id);
+        self.query_engine.clear_program(program_id);
     }
 
     /// Removes all data associated with `source_id` from the declaration and type engines.
@@ -54,6 +55,7 @@ impl Engines {
         self.type_engine.clear_module(source_id);
         self.decl_engine.clear_module(source_id);
         self.parsed_decl_engine.clear_module(source_id);
+        self.query_engine.clear_module(source_id);
     }
 
     /// Helps out some `thing: T` by adding `self` as context.

sway-core/src/query_engine/mod.rs

@@ -7,7 +7,7 @@ use std::{
     time::SystemTime,
 };
 use sway_error::{error::CompileError, warning::CompileWarning};
-use sway_types::IdentUnique;
+use sway_types::{IdentUnique, ProgramId, SourceId, Spanned};
 
 use crate::{
     decl_engine::{DeclId, DeclRef},
@@ -141,17 +141,18 @@ pub struct FunctionCacheEntry {
 #[derive(Debug, Default)]
 pub struct QueryEngine {
     // We want the below types wrapped in Arcs to optimize cloning from LSP.
-    programs_cache: Arc<RwLock<ProgramsCacheMap>>,
-    function_cache: Arc<RwLock<FunctionsCacheMap>>,
+    programs_cache: CowCache<ProgramsCacheMap>,
     pub module_cache: CowCache<ModuleCacheMap>,
+    // NOTE: Any further AstNodes that are cached need to have garbage collection applied, see clear_module()
+    function_cache: CowCache<FunctionsCacheMap>,
 }
 
 impl Clone for QueryEngine {
     fn clone(&self) -> Self {
         Self {
-            programs_cache: self.programs_cache.clone(),
-            function_cache: self.function_cache.clone(),
+            programs_cache: CowCache::new(self.programs_cache.read().clone()),
             module_cache: CowCache::new(self.module_cache.read().clone()),
+            function_cache: CowCache::new(self.function_cache.read().clone()),
         }
     }
 }
@@ -205,6 +206,30 @@ impl QueryEngine {
             FunctionCacheEntry { fn_decl },
         );
     }
+
+    /// Removes all data associated with the `source_id` from the function cache.
+    pub fn clear_module(&mut self, source_id: &SourceId) {
+        self.function_cache
+            .write()
+            .retain(|(ident, _), _| ident.span().source_id().map_or(true, |id| id != source_id));
+    }
+
+    /// Removes all data associated with the `program_id` from the function cache.
+    pub fn clear_program(&mut self, program_id: &ProgramId) {
+        self.function_cache.write().retain(|(ident, _), _| {
+            ident
+                .span()
+                .source_id()
+                .map_or(true, |id| id.program_id() != *program_id)
+        });
+    }
+
+    ///  Commits all changes to their respective caches.
+    pub fn commit(&self) {
+        self.programs_cache.commit();
+        self.module_cache.commit();
+        self.function_cache.commit();
+    }
 }
 
 /// Thread-safe, copy-on-write cache optimized for LSP operations.

sway-lsp/src/config.rs

@@ -71,17 +71,11 @@ impl Default for DiagnosticConfig {
 #[serde(rename_all = "camelCase")]
 pub struct GarbageCollectionConfig {
     pub gc_enabled: bool,
-    pub gc_frequency: i32,
 }
 
 impl Default for GarbageCollectionConfig {
     fn default() -> Self {
-        Self {
-            gc_enabled: true,
-            // Garbage collection is fairly expsensive so we default to only clear on every 3rd keystroke.
-            // Waiting too long to clear can cause a stack overflow to occur.
-            gc_frequency: 3,
-        }
+        Self { gc_enabled: true }
     }
 }
 

sway-lsp/src/server_state.rs

@@ -137,21 +137,17 @@ impl ServerState {
                         let session = ctx.session.as_ref().unwrap().clone();
                         let mut engines_clone = session.engines.read().clone();
 
-                        if let Some(version) = ctx.version {
-                            // Perform garbage collection at configured intervals if enabled to manage memory usage.
-                            if ctx.gc_options.gc_enabled
-                                && version % ctx.gc_options.gc_frequency == 0
+                        // Perform garbage collection if enabled to manage memory usage.
+                        if ctx.gc_options.gc_enabled {
+                            // Call this on the engines clone so we don't clear types that are still in use
+                            // and might be needed in the case cancel compilation was triggered.
+                            if let Err(err) =
+                                session.garbage_collect_module(&mut engines_clone, &uri)
                             {
-                                // Call this on the engines clone so we don't clear types that are still in use
-                                // and might be needed in the case cancel compilation was triggered.
-                                if let Err(err) =
-                                    session.garbage_collect_module(&mut engines_clone, &uri)
-                                {
-                                    tracing::error!(
-                                        "Unable to perform garbage collection: {}",
-                                        err.to_string()
-                                    );
-                                }
+                                tracing::error!(
+                                    "Unable to perform garbage collection: {}",
+                                    err.to_string()
+                                );
                             }
                         }
 
@@ -180,10 +176,10 @@ impl ServerState {
                                             // Because the engines_clone has garbage collection applied. If the workspace AST was reused, we need to keep the old engines
                                             // as the engines_clone might have cleared some types that are still in use.
                                             if metrics.reused_programs == 0 {
-                                                // Commit local changes in the module cache to the shared state.
+                                                // Commit local changes in the programs, module, and function caches to the shared state.
                                                 // This ensures that any modifications made during compilation are preserved
                                                 // before we swap the engines.
-                                                engines_clone.qe().module_cache.commit();
+                                                engines_clone.qe().commit();
                                                 // The compiler did not reuse the workspace AST.
                                                 // We need to overwrite the old engines with the engines clone.
                                                 mem::swap(

sway-lsp/tests/fixtures/garbage_collection/minimal_script/.gitignore

@@ -0,0 +1,2 @@
+out
+target

sway-lsp/tests/fixtures/garbage_collection/minimal_script/Forc.toml

@@ -0,0 +1,8 @@
+[project]
+authors = ["Fuel Labs <contact@fuel.sh>"]
+entry = "main.sw"
+license = "Apache-2.0"
+name = "minimal_script"
+implicit-std = false
+
+[dependencies]

sway-lsp/tests/fixtures/garbage_collection/minimal_script/src/main.sw

@@ -0,0 +1,14 @@
+script;
+
+struct MyStruct {
+    field1: u16,
+}
+
+fn func(s: MyStruct) -> u16 {
+    s.field1
+}
+
+fn main() {
+    let x = MyStruct { field1: 10 };
+    let y = func(x);
+}

sway-lsp/tests/lib.rs

@@ -261,13 +261,6 @@ fn garbage_collection_runner(path: PathBuf) {
     run_async!({
         setup_panic_hook();
         let (mut service, _) = LspService::new(ServerState::new);
-        // set the garbage collection frequency to 1
-        service
-            .inner()
-            .config
-            .write()
-            .garbage_collection
-            .gc_frequency = 1;
         let uri = init_and_open(&mut service, path).await;
         let times = 60;
 
@@ -302,6 +295,14 @@ fn garbage_collection_paths() {
     garbage_collection_runner(p);
 }
 
+#[test]
+fn garbage_collection_minimal_script() {
+    let p = sway_workspace_dir()
+        .join("sway-lsp/tests/fixtures/garbage_collection/minimal_script")
+        .join("src/main.sw");
+    garbage_collection_runner(p);
+}
+
 #[test]
 fn lsp_syncs_with_workspace_edits() {
     run_async!({