fix session cache bug

forc-pkg/src/pkg.rs

@@ -2754,7 +2754,11 @@ pub fn check(
             return Ok(results);
         }
         results.push((programs_res.ok(), handler));
-        eprintln!("⏱️ Compiling package {:?} took {:?}", pkg.name, now.elapsed());
+        eprintln!(
+            "⏱️ Compiling package {:?} took {:?}",
+            pkg.name,
+            now.elapsed()
+        );
     }
 
     if results.is_empty() {

sway-core/src/semantic_analysis/module.rs

@@ -296,7 +296,10 @@ impl ty::TyModule {
                     } else {
                         "🔄 Cache miss"
                     };
-                    eprintln!("{} for module {:?} (up to date: {})", status, relevant_path, is_up_to_date);
+                    eprintln!(
+                        "{} for module {:?} (up to date: {})",
+                        status, relevant_path, is_up_to_date
+                    );
                     if is_up_to_date {
                         return Some(typed.module);
                     }
@@ -329,7 +332,10 @@ impl ty::TyModule {
         } = parsed;
 
         // Check if the root module cache is up to date
-        eprintln!("Root Module: {:?}", parsed.span.source_id().map(|x| engines.se().get_path(x)));
+        eprintln!(
+            "Root Module: {:?}",
+            parsed.span.source_id().map(|x| engines.se().get_path(x))
+        );
         if let Some(module) =
             ty::TyModule::check_cache(parsed.span.source_id(), engines, build_config)
         {
@@ -373,8 +379,6 @@ impl ty::TyModule {
             })
             .collect::<Result<Vec<_>, _>>();
 
-
-
         // TODO: Ordering should be solved across all modules prior to the beginning of type-check.
         let ordered_nodes = node_dependencies::order_ast_nodes_by_dependency(
             handler,

sway-lsp/src/core/session.rs

@@ -26,7 +26,6 @@ use pkg::{
     BuildPlan,
 };
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
-use swayfmt::parse;
 use std::{
     path::PathBuf,
     sync::{atomic::AtomicBool, Arc},
@@ -44,6 +43,7 @@ use sway_core::{
 use sway_error::{error::CompileError, handler::Handler, warning::CompileWarning};
 use sway_types::{ProgramId, SourceEngine, Spanned};
 use sway_utils::{helpers::get_sway_files, PerformanceData};
+use swayfmt::parse;
 
 pub type RunnableMap = DashMap<PathBuf, Vec<Box<dyn Runnable>>>;
 pub type ProjectDirectory = PathBuf;
@@ -144,7 +144,10 @@ impl Session {
         uri: &Url,
     ) -> Result<(), LanguageServerError> {
         let path = uri.to_file_path().unwrap();
-        eprintln!("🗑️ 🗑️ 🗑️ 🗑️ 🗑️   Garbage collecting module {:?}   🗑️ 🗑️ 🗑️ 🗑️ 🗑️", path);
+        eprintln!(
+            "🗑️ 🗑️ 🗑️ 🗑️ 🗑️   Garbage collecting module {:?}   🗑️ 🗑️ 🗑️ 🗑️ 🗑️",
+            path
+        );
         let source_id = { engines.se().get_source_id(&path) };
         engines.clear_module(&source_id);
         Ok(())
@@ -280,16 +283,12 @@ pub fn compile(
     experimental: sway_core::ExperimentalFlags,
 ) -> Result<Vec<(Option<Programs>, Handler)>, LanguageServerError> {
     let _p = tracing::trace_span!("compile").entered();
-    let build_plan_now = std::time::Instant::now();
-    let build_plan = build_plan(uri)?;
-    eprintln!("⏱️ Build Plan took {:?}", build_plan_now.elapsed());
-    let tests_enabled = true;
     pkg::check(
         build_plan,
         BuildTarget::default(),
         true,
         lsp_mode,
-        tests_enabled,
+        true,
         engines,
         retrigger_compilation,
         experimental,
@@ -407,9 +406,12 @@ pub fn parse_project(
     experimental: sway_core::ExperimentalFlags,
 ) -> Result<(), LanguageServerError> {
     let _p = tracing::trace_span!("parse_project").entered();
+    let build_plan_now = std::time::Instant::now();
     let build_plan = session
         .build_plan_cache
         .get_or_update(&session.sync.manifest_path(), || build_plan(uri))?;
+    eprintln!("⏱️ Build Plan took {:?}", build_plan_now.elapsed());
+
     let parse_now = std::time::Instant::now();
     let results = compile(
         &build_plan,
@@ -441,7 +443,10 @@ pub fn parse_project(
         create_runnables(&session.runnables, typed, engines.de(), engines.se());
     }
     eprintln!("⏱️ creating runnables took: {:?}", runnables_now.elapsed());
-    eprintln!("⏱️ TOTAL COMPILATION AND TRAVERSAL TIME: {:?}", parse_now.elapsed());
+    eprintln!(
+        "⏱️ TOTAL COMPILATION AND TRAVERSAL TIME: {:?}",
+        parse_now.elapsed()
+    );
     Ok(())
 }
 

sway-lsp/src/server_state.rs

@@ -41,7 +41,7 @@ pub struct ServerState {
     /// A Least Recently Used (LRU) cache of [Session]s, each representing a project opened in the user's workspace.
     /// This cache limits memory usage by maintaining a fixed number of active sessions, automatically
     /// evicting the least recently used sessions when the capacity is reached.
-    pub(crate) sessions: LruSessionCache,
+    pub sessions: LruSessionCache,
     pub documents: Documents,
     // Compilation thread related fields
     pub(crate) retrigger_compilation: Arc<AtomicBool>,
@@ -182,7 +182,10 @@ impl ServerState {
                                             // It's very important to check if the workspace AST was reused to determine if we need to overwrite the engines.
                                             // 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.
-                                            eprintln!("👨‍💻  metrics.reused_programs {} 👨‍💻", metrics.reused_programs);
+                                            eprintln!(
+                                                "👨‍💻  metrics.reused_programs {} 👨‍💻",
+                                                metrics.reused_programs
+                                            );
                                             if metrics.reused_programs == 0 {
                                                 // The compiler did not reuse the workspace AST.
                                                 // We need to overwrite the old engines with the engines clone.
@@ -353,7 +356,7 @@ impl ServerState {
 
     /// Constructs and returns a tuple of `(Url, Arc<Session>)` from a given workspace URI.
     /// The returned URL represents the temp directory workspace.
-    pub(crate) async fn uri_and_session_from_workspace(
+    pub async fn uri_and_session_from_workspace(
         &self,
         workspace_uri: &Url,
     ) -> Result<(Url, Arc<Session>), LanguageServerError> {
@@ -377,23 +380,23 @@ impl ServerState {
             .ok_or(DirectoryError::ManifestDirNotFound)?
             .to_path_buf();
 
-        let session = self.sessions.get(&manifest_dir).unwrap_or({
-            // If no session can be found, then we need to call init and insert a new session into the map
-            eprintln!("💥💥💥💥💥💥💥💥💥💥💥💥💥💥💥 Initializing new session for manifest_dir = {:?}", manifest_dir);
+        // If the session is already in the cache, return it
+        if let Some(session) = self.sessions.get(&manifest_dir) {
+            return Ok(session);
+        }
 
+        // If no session can be found, then we need to call init and insert a new session into the map
+        let session = Arc::new(Session::new());
+        session.init(uri, &self.documents).await?;
+        self.sessions.insert(manifest_dir.clone(), session.clone());
 
-            self.init_session(uri).await?;
-            self.sessions
-                .get(&manifest_dir)
-                .expect("no session found even though it was just inserted into the map")
-        });
         Ok(session)
     }
 }
 
 /// A Least Recently Used (LRU) cache for storing and managing `Session` objects.
 /// This cache helps limit memory usage by maintaining a fixed number of active sessions.
-pub(crate) struct LruSessionCache {
+pub struct LruSessionCache {
     /// Stores the actual `Session` objects, keyed by their file paths.
     sessions: Arc<DashMap<PathBuf, Arc<Session>>>,
     /// Keeps track of the order in which sessions were accessed, with most recent at the front.
@@ -404,20 +407,20 @@ pub(crate) struct LruSessionCache {
 
 impl LruSessionCache {
     /// Creates a new `LruSessionCache` with the specified capacity.
-    pub(crate) fn new(capacity: usize) -> Self {
+    pub fn new(capacity: usize) -> Self {
         LruSessionCache {
             sessions: Arc::new(DashMap::new()),
             usage_order: Arc::new(Mutex::new(VecDeque::with_capacity(capacity))),
             capacity,
         }
     }
 
-    pub(crate) fn iter(&self) -> impl Iterator<Item = RefMulti<'_, PathBuf, Arc<Session>>> {
+    pub fn iter(&self) -> impl Iterator<Item = RefMulti<'_, PathBuf, Arc<Session>>> {
         self.sessions.iter()
     }
 
     /// Retrieves a session from the cache and updates its position to the front of the usage order.
-    pub(crate) fn get(&self, path: &PathBuf) -> Option<Arc<Session>> {
+    pub fn get(&self, path: &PathBuf) -> Option<Arc<Session>> {
         if let Some(session) = self.sessions.try_get(path).try_unwrap() {
             if self.sessions.len() >= self.capacity {
                 self.move_to_front(path);
@@ -431,16 +434,13 @@ impl LruSessionCache {
     /// Inserts or updates a session in the cache.
     /// If at capacity and inserting a new session, evicts the least recently used one.
     /// For existing sessions, updates their position in the usage order if at capacity.
-    pub(crate) fn insert(&self, path: PathBuf, session: Arc<Session>) {
-        if self.sessions.get(&path).is_some() {
-            tracing::trace!("Updating existing session for path: {:?}", path);
-            // Session already exists, just update its position in the usage order if at capacity
-            if self.sessions.len() >= self.capacity {
-                self.move_to_front(&path);
-            }
+    pub fn insert(&self, path: PathBuf, session: Arc<Session>) {
+        if let Some(mut entry) = self.sessions.get_mut(&path) {
+            // Session already exists, update it
+            *entry = session;
+            self.move_to_front(&path);
         } else {
             // New session
-            tracing::trace!("Inserting new session for path: {:?}", path);
             if self.sessions.len() >= self.capacity {
                 self.evict_least_used();
             }
@@ -472,3 +472,81 @@ impl LruSessionCache {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::path::PathBuf;
+    use std::sync::Arc;
+
+    #[test]
+    fn test_lru_session_cache_insertion_and_retrieval() {
+        let cache = LruSessionCache::new(2);
+        let path1 = PathBuf::from("/path/1");
+        let path2 = PathBuf::from("/path/2");
+        let session1 = Arc::new(Session::new());
+        let session2 = Arc::new(Session::new());
+
+        cache.insert(path1.clone(), session1.clone());
+        cache.insert(path2.clone(), session2.clone());
+
+        assert!(Arc::ptr_eq(&cache.get(&path1).unwrap(), &session1));
+        assert!(Arc::ptr_eq(&cache.get(&path2).unwrap(), &session2));
+    }
+
+    #[test]
+    fn test_lru_session_cache_capacity() {
+        let cache = LruSessionCache::new(2);
+        let path1 = PathBuf::from("/path/1");
+        let path2 = PathBuf::from("/path/2");
+        let path3 = PathBuf::from("/path/3");
+        let session1 = Arc::new(Session::new());
+        let session2 = Arc::new(Session::new());
+        let session3 = Arc::new(Session::new());
+
+        cache.insert(path1.clone(), session1);
+        cache.insert(path2.clone(), session2);
+        cache.insert(path3.clone(), session3);
+
+        assert!(cache.get(&path1).is_none());
+        assert!(cache.get(&path2).is_some());
+        assert!(cache.get(&path3).is_some());
+    }
+
+    #[test]
+    fn test_lru_session_cache_update_order() {
+        let cache = LruSessionCache::new(2);
+        let path1 = PathBuf::from("/path/1");
+        let path2 = PathBuf::from("/path/2");
+        let path3 = PathBuf::from("/path/3");
+        let session1 = Arc::new(Session::new());
+        let session2 = Arc::new(Session::new());
+        let session3 = Arc::new(Session::new());
+
+        cache.insert(path1.clone(), session1.clone());
+        cache.insert(path2.clone(), session2.clone());
+
+        // Access path1 to move it to the front
+        cache.get(&path1);
+
+        // Insert path3, which should evict path2
+        cache.insert(path3.clone(), session3);
+
+        assert!(cache.get(&path1).is_some());
+        assert!(cache.get(&path2).is_none());
+        assert!(cache.get(&path3).is_some());
+    }
+
+    #[test]
+    fn test_lru_session_cache_overwrite() {
+        let cache = LruSessionCache::new(2);
+        let path1 = PathBuf::from("/path/1");
+        let session1 = Arc::new(Session::new());
+        let session1_new = Arc::new(Session::new());
+
+        cache.insert(path1.clone(), session1);
+        cache.insert(path1.clone(), session1_new.clone());
+
+        assert!(Arc::ptr_eq(&cache.get(&path1).unwrap(), &session1_new));
+    }
+}

sway-lsp/tests/lib.rs

@@ -171,7 +171,7 @@ fn did_open_fluid_libraries() {
             .finish();
         let uri = init_and_open(
             &mut service,
-            PathBuf::from("/Users/joshuabatty/Documents/rust/fuel/user_projects/fluid-protocol/libraries")
+            PathBuf::from("/Users/josh/Documents/rust/fuel/user_projects/fluid-protocol/libraries")
                 .join("src/interface.sw"),
         )
         .await;
@@ -239,21 +239,23 @@ fn did_change_stress_test_random_wait() {
             let (mut service, _) = LspService::new(ServerState::new);
             // let example_dir = sway_workspace_dir()
             //     .join(e2e_language_dir())
-                // .join("generics_in_contract");
-//            let uri = init_and_open(&mut service, example_dir.join("src/main.sw")).await;
+            // .join("generics_in_contract");
+            //            let uri = init_and_open(&mut service, example_dir.join("src/main.sw")).await;
 
             let uri = init_and_open(
                 &mut service,
-                PathBuf::from("/Users/joshuabatty/Documents/rust/fuel/user_projects/fluid-protocol/libraries")
-                    .join("src/fpt_staking_interface.sw"),
+                PathBuf::from(
+                    "/Users/josh/Documents/rust/fuel/user_projects/fluid-protocol/libraries",
+                )
+                .join("src/fpt_staking_interface.sw"),
             )
             .await;
 
-        // 1. randomise the file that is changed out of all the files in the project.
-        // 2. change the file
-        // 3. wait for the file to be parsed
-        // 4. try and do a hover or goto def for a random type
-        // 5. repeat.
+            // 1. randomise the file that is changed out of all the files in the project.
+            // 2. change the file
+            // 3. wait for the file to be parsed
+            // 4. try and do a hover or goto def for a random type
+            // 5. repeat.
 
             let times = 6000;
             for version in 0..times {
@@ -2191,3 +2193,36 @@ async fn write_all_example_asts() {
     }
     let _ = server.shutdown_server();
 }
+
+#[test]
+fn test_url_to_session_existing_session() {
+    use std::sync::Arc;
+    run_async!({
+        let (mut service, _) = LspService::new(ServerState::new);
+        let uri = init_and_open(&mut service, doc_comments_dir().join("src/main.sw")).await;
+
+        // First call to uri_and_session_from_workspace
+        let (first_uri, first_session) = service
+            .inner()
+            .uri_and_session_from_workspace(&uri)
+            .await
+            .unwrap();
+
+        // Second call to uri_and_session_from_workspace
+        let (second_uri, second_session) = service
+            .inner()
+            .uri_and_session_from_workspace(&uri)
+            .await
+            .unwrap();
+
+        // Assert that the URIs are the same
+        assert_eq!(first_uri, second_uri, "URIs should be identical");
+
+        // Assert that the sessions are the same (they should point to the same Arc)
+        assert!(
+            Arc::ptr_eq(&first_session, &second_session),
+            "Sessions should be identical"
+        );
+        shutdown_and_exit(&mut service).await;
+    });
+}