Fix function_cache garbage collection bug (#6555)

## Description
Fixes a bug that would crash the language server when a function node
was part of the AST.

#5967 added a function cache to the `QueryEngine`. Garbage collection
needed to be called on this otherwise it returns references to types
that have been cleared from the other engines during garbage collection.

In the future, any other nodes that we decide to cache need to have GC
applied to them or we will end up crashing the language server on the
first key-pressed event.

I've also removed the gc_frequency setting and now run GC on each
keystroke. Without doing this, the spans stored in the TokenMap and what
was returned from the compiler were falling out of sync. This ensures
that everything is always up to date and the correct spans are used.
closes #5260

Finally, I've added a `launch.json` that allows for attaching `lldb` to
a live running `forc-lsp` process. This was a pretty cruical step for
tracking down this bug so would be nice to have it easily accessible for
future debugging sessions.

## 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).
- [ ] 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.

---------

Co-authored-by: Sophie Dankel <47993817+sdankel@users.noreply.github.com>

.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!({