LSP: Support multiple workspaces

sway-lsp/benches/lsp_benchmarks/compile.rs

@@ -14,7 +14,7 @@ fn benchmarks(c: &mut Criterion) {
         .unwrap()
         .block_on(async { black_box(super::compile_test_project().await) });
 
-    let sync = state.sync_workspace.get().unwrap();
+    let sync = state.get_sync_workspace_for_uri(&uri).unwrap();
     let build_plan = session
         .build_plan_cache
         .get_or_update(&sync.workspace_manifest_path(), || {

sway-lsp/benches/lsp_benchmarks/mod.rs

@@ -17,7 +17,7 @@ pub async fn compile_test_project() -> (Url, Arc<Session>, ServerState, Engines)
     let state = ServerState::default();
     let engines_clone = state.engines.read().clone();
     let session = Arc::new(Session::new());
-    let sync = state.get_or_init_global_sync_workspace(&uri).await.unwrap();
+    let sync = state.get_or_init_sync_workspace(&uri).await.unwrap();
     let temp_uri = sync.workspace_to_temp_url(&uri).unwrap();
 
     state.documents.handle_open_file(&temp_uri).await;

sway-lsp/benches/lsp_benchmarks/requests.rs

@@ -10,7 +10,7 @@ fn benchmarks(c: &mut Criterion) {
     let (uri, session, state, engines) = Runtime::new()
         .unwrap()
         .block_on(async { black_box(super::compile_test_project().await) });
-    let sync = state.sync_workspace.get().unwrap();
+    let sync = state.get_sync_workspace_for_uri(&uri).unwrap();
     let config = sway_lsp::config::Config::default();
     let position = Position::new(1717, 24);
     let range = Range::new(Position::new(1628, 0), Position::new(1728, 0));
@@ -55,12 +55,12 @@ fn benchmarks(c: &mut Criterion) {
     });
 
     c.bench_function("find_all_references", |b| {
-        b.iter(|| session.token_references(&uri, position, &state.token_map, &engines, sync))
+        b.iter(|| session.token_references(&uri, position, &state.token_map, &engines, &sync))
     });
 
     c.bench_function("goto_definition", |b| {
         b.iter(|| {
-            session.token_definition_response(&uri, position, &engines, &state.token_map, sync)
+            session.token_definition_response(&uri, position, &engines, &state.token_map, &sync)
         })
     });
 
@@ -78,7 +78,7 @@ fn benchmarks(c: &mut Criterion) {
 
     c.bench_function("prepare_rename", |b| {
         b.iter(|| {
-            capabilities::rename::prepare_rename(&engines, &state.token_map, &uri, position, sync)
+            capabilities::rename::prepare_rename(&engines, &state.token_map, &uri, position, &sync)
         })
     });
 
@@ -90,7 +90,7 @@ fn benchmarks(c: &mut Criterion) {
                 "new_token_name".to_string(),
                 &uri,
                 position,
-                sync,
+                &sync,
             )
         })
     });

sway-lsp/benches/lsp_benchmarks/token_map.rs

@@ -6,7 +6,7 @@ fn benchmarks(c: &mut Criterion) {
     let (uri, _, state, engines) = Runtime::new()
         .unwrap()
         .block_on(async { black_box(super::compile_test_project().await) });
-    let sync = state.sync_workspace.get().unwrap();
+    let sync = state.get_sync_workspace_for_uri(&uri).unwrap();
     let position = Position::new(1716, 24);
     let path = uri.to_file_path().unwrap();
     let program_id = sway_lsp::core::session::program_id_from_path(&path, &engines).unwrap();

sway-lsp/src/capabilities/hover/mod.rs

@@ -52,6 +52,7 @@ pub fn hover_data(
     }
 
     let client_config = state.config.read().client.clone();
+    let sync = state.get_sync_workspace_for_uri(url).unwrap();
     let contents = match &token.declared_token_ident(engines) {
         Some(decl_ident) => {
             let t = state.token_map.try_get(decl_ident).try_unwrap()?;
@@ -62,7 +63,7 @@ pub fn hover_data(
                 decl_token,
                 &decl_ident.name,
                 client_config,
-                state.sync_workspace(),
+                &sync,
             )
         }
         // The `TypeInfo` of the token does not contain an `Ident`. In this case,
@@ -73,7 +74,7 @@ pub fn hover_data(
             token,
             &ident.name,
             client_config,
-            state.sync_workspace(),
+            &sync,
         ),
     };
 

sway-lsp/src/handlers/notification.rs

@@ -2,7 +2,7 @@
 //! Protocol. This module specifically handles notification messages sent by the Client.
 
 use crate::{
-    core::{document::Documents, session::Session},
+    core::{document::Documents, session::Session, sync::SyncWorkspace},
     error::LanguageServerError,
     server_state::{CompilationContext, ServerState, TaskMessage},
 };
@@ -21,14 +21,14 @@ pub async fn handle_did_open_text_document(
     params: DidOpenTextDocumentParams,
 ) -> Result<(), LanguageServerError> {
     let file_uri = &params.text_document.uri;
-    // Initialize the SyncWorkspace if it doesn't exist.
-    let _ = state.get_or_init_global_sync_workspace(file_uri).await?;
+    // Initialize the SyncWorkspace for this file if it doesn't exist.
+    let sync_workspace = state.get_or_init_sync_workspace(file_uri).await?;
 
     // Get or create a session for the original file URI.
     let (uri, session) = state.uri_and_session_from_workspace(&params.text_document.uri)?;
     state.documents.handle_open_file(&uri).await;
 
-    send_new_compilation_request(state, session.clone(), &uri, None, false);
+    send_new_compilation_request(state, session.clone(), &uri, None, false, sync_workspace);
     state.is_compiling.store(true, Ordering::SeqCst);
     state.wait_for_parsing().await;
     state
@@ -44,6 +44,7 @@ fn send_new_compilation_request(
     uri: &Url,
     version: Option<i32>,
     optimized_build: bool,
+    sync_workspace: Arc<SyncWorkspace>,
 ) {
     let file_versions = file_versions(&state.documents, uri, version.map(|v| v as u64));
 
@@ -72,7 +73,7 @@ fn send_new_compilation_request(
             optimized_build,
             gc_options: state.config.read().garbage_collection.clone(),
             file_versions,
-            sync: Some(state.sync_workspace.get().unwrap().clone()),
+            sync: Some(sync_workspace),
         }));
 }
 
@@ -88,6 +89,7 @@ pub async fn handle_did_change_text_document(
     }
 
     let (uri, session) = state.uri_and_session_from_workspace(&params.text_document.uri)?;
+    let sync_workspace = state.get_sync_workspace_for_uri(&params.text_document.uri)?;
     state
         .documents
         .write_changes_to_file(&uri, &params.content_changes)
@@ -100,6 +102,7 @@ pub async fn handle_did_change_text_document(
         Some(params.text_document.version),
         // TODO: Set this back to true once https://github.com/FuelLabs/sway/issues/6576 is fixed.
         false,
+        sync_workspace,
     );
     Ok(())
 }
@@ -129,7 +132,8 @@ pub(crate) async fn handle_did_save_text_document(
         .pid_locked_files
         .remove_dirty_flag(&params.text_document.uri)?;
     let (uri, session) = state.uri_and_session_from_workspace(&params.text_document.uri)?;
-    send_new_compilation_request(state, session.clone(), &uri, None, false);
+    let sync_workspace = state.get_sync_workspace_for_uri(&params.text_document.uri)?;
+    send_new_compilation_request(state, session.clone(), &uri, None, false, sync_workspace);
     state.wait_for_parsing().await;
     state
         .publish_diagnostics(uri, params.text_document.uri, session)
@@ -147,7 +151,8 @@ pub(crate) fn handle_did_change_watched_files(
         match event.typ {
             FileChangeType::CHANGED => {
                 if event.uri.to_string().contains("Forc.toml") {
-                    state.sync_workspace.get().unwrap().sync_manifest()?;
+                    let sync_workspace = state.get_sync_workspace_for_uri(&event.uri)?;
+                    sync_workspace.sync_manifest()?;
                     // TODO: Recompile the project | see https://github.com/FuelLabs/sway/issues/7103
                 }
             }

sway-lsp/src/handlers/request.rs

@@ -80,14 +80,14 @@ pub fn handle_goto_definition(
         .uri_and_session_from_workspace(&params.text_document_position_params.text_document.uri)
     {
         Ok((uri, session)) => {
-            let sync = state.sync_workspace();
+            let sync = state.get_sync_workspace_for_uri(&uri).unwrap();
             let position = params.text_document_position_params.position;
             Ok(session.token_definition_response(
                 &uri,
                 position,
                 &state.engines.read(),
                 &state.token_map,
-                sync,
+                &sync,
             ))
         }
         Err(err) => {
@@ -154,13 +154,13 @@ pub fn handle_prepare_rename(
 ) -> Result<Option<PrepareRenameResponse>> {
     match state.uri_from_workspace(&params.text_document.uri) {
         Ok(uri) => {
-            let sync = state.sync_workspace();
+            let sync = state.get_sync_workspace_for_uri(&uri).unwrap();
             match capabilities::rename::prepare_rename(
                 &state.engines.read(),
                 &state.token_map,
                 &uri,
                 params.position,
-                sync,
+                &sync,
             ) {
                 Ok(res) => Ok(Some(res)),
                 Err(err) => {
@@ -181,14 +181,14 @@ pub fn handle_rename(state: &ServerState, params: RenameParams) -> Result<Option
         Ok(uri) => {
             let new_name = params.new_name;
             let position = params.text_document_position.position;
-            let sync = state.sync_workspace();
+            let sync = state.get_sync_workspace_for_uri(&uri).unwrap();
             match capabilities::rename::rename(
                 &state.engines.read(),
                 &state.token_map,
                 new_name,
                 &uri,
                 position,
-                sync,
+                &sync,
             ) {
                 Ok(res) => Ok(Some(res)),
                 Err(err) => {
@@ -237,13 +237,13 @@ pub async fn handle_references(
     match state.uri_and_session_from_workspace(&params.text_document_position.text_document.uri) {
         Ok((uri, session)) => {
             let position = params.text_document_position.position;
-            let sync = state.sync_workspace();
+            let sync = state.get_sync_workspace_for_uri(&uri).unwrap();
             Ok(session.token_references(
                 &uri,
                 position,
                 &state.token_map,
                 &state.engines.read(),
-                sync,
+                &sync,
             ))
         }
         Err(err) => {

sway-lsp/src/server.rs

@@ -31,14 +31,14 @@ impl LanguageServer for ServerState {
         if let Err(err) = self.register_forc_toml_watcher().await {
             tracing::error!("Failed to register Forc.toml file watcher: {}", err);
         }
-        // Populate documents from temp dir
-        if let Some(sw) = self.sync_workspace.get() {
-            if let Ok(temp_dir) = sw.temp_dir() {
-                if let Err(err) = self.documents.store_sway_files_from_temp(temp_dir).await {
-                    tracing::warn!("Failed to populate documents from temp dir: {}", err);
-                }
-            }
-        }
+        // // Populate documents from temp dir
+        // if let Some(sw) = self.sync_workspace.get() {
+        //     if let Ok(temp_dir) = sw.temp_dir() {
+        //         if let Err(err) = self.documents.store_sway_files_from_temp(temp_dir).await {
+        //             tracing::warn!("Failed to populate documents from temp dir: {}", err);
+        //         }
+        //     }
+        // }
         tracing::info!("Sway Language Server Initialized");
     }