implement proper lsp-workspace support

fix typo

Co-authored-by: LeoniePhiline <22329650+LeoniePhiline@users.noreply.github.com>

3 files changed, 224 insertions, 70 deletions

diff --git a/helix-lsp/Cargo.toml b/helix-lsp/Cargo.toml
index 33ec5f30..f8526515 100644
--- a/helix-lsp/Cargo.toml
+++ b/helix-lsp/Cargo.toml
@@ -27,3 +27,4 @@ thiserror = "1.0"
 tokio = { version = "1.27", features = ["rt", "rt-multi-thread", "io-util", "io-std", "time", "process", "macros", "fs", "parking_lot", "sync"] }
 tokio-stream = "0.1.12"
 which = "4.4"
+parking_lot = "0.12.1"
diff --git a/helix-lsp/src/client.rs b/helix-lsp/src/client.rs
index 34e4c346..3dab6bc5 100644
--- a/helix-lsp/src/client.rs
+++ b/helix-lsp/src/client.rs
@@ -1,13 +1,17 @@
 use crate::{
-    find_root, jsonrpc,
+    find_lsp_workspace, jsonrpc,
     transport::{Payload, Transport},
     Call, Error, OffsetEncoding, Result,
 };
 
-use helix_core::{ChangeSet, Rope};
+use helix_core::{find_workspace, ChangeSet, Rope};
 use helix_loader::{self, VERSION_AND_GIT_HASH};
-use lsp::PositionEncodingKind;
+use lsp::{
+    notification::DidChangeWorkspaceFolders, DidChangeWorkspaceFoldersParams, OneOf,
+    PositionEncodingKind, WorkspaceFolder, WorkspaceFoldersChangeEvent,
+};
 use lsp_types as lsp;
+use parking_lot::Mutex;
 use serde::Deserialize;
 use serde_json::Value;
 use std::future::Future;
@@ -26,6 +30,17 @@ use tokio::{
     },
 };
 
+fn workspace_for_uri(uri: lsp::Url) -> WorkspaceFolder {
+    lsp::WorkspaceFolder {
+        name: uri
+            .path_segments()
+            .and_then(|segments| segments.last())
+            .map(|basename| basename.to_string())
+            .unwrap_or_default(),
+        uri,
+    }
+}
+
 #[derive(Debug)]
 pub struct Client {
     id: usize,
@@ -36,11 +51,120 @@ pub struct Client {
     config: Option<Value>,
     root_path: std::path::PathBuf,
     root_uri: Option<lsp::Url>,
-    workspace_folders: Vec<lsp::WorkspaceFolder>,
+    workspace_folders: Mutex<Vec<lsp::WorkspaceFolder>>,
+    initalize_notify: Arc<Notify>,
+    /// workspace folders added while the server is still initalizing
     req_timeout: u64,
 }
 
 impl Client {
+    pub fn try_add_doc(
+        self: &Arc<Self>,
+        root_markers: &[String],
+        manual_roots: &[PathBuf],
+        doc_path: Option<&std::path::PathBuf>,
+        may_support_workspace: bool,
+    ) -> bool {
+        let (workspace, workspace_is_cwd) = find_workspace();
+        let root = find_lsp_workspace(
+            doc_path
+                .and_then(|x| x.parent().and_then(|x| x.to_str()))
+                .unwrap_or("."),
+            root_markers,
+            manual_roots,
+            &workspace,
+            workspace_is_cwd,
+        );
+        let root_uri = root
+            .as_ref()
+            .and_then(|root| lsp::Url::from_file_path(root).ok());
+
+        if self.root_path == root.unwrap_or(workspace)
+            || root_uri.as_ref().map_or(false, |root_uri| {
+                self.workspace_folders
+                    .lock()
+                    .iter()
+                    .any(|workspace| &workspace.uri == root_uri)
+            })
+        {
+            // workspace URI is already registered so we can use this client
+            return true;
+        }
+
+        // this server definitly doesn't support multiple workspace, no need to check capabilities
+        if !may_support_workspace {
+            return false;
+        }
+
+        let Some(capabilities) = self.capabilities.get() else {
+            let client = Arc::clone(self);
+            // initalization hasn't finished yet, deal with this new root later
+            // TODO: In the edgecase that a **new root** is added
+            // for an LSP that **doesn't support workspace_folders** before initaliation is finished
+            // the new roots are ignored.
+            // That particular edgecase would require retroactively spawning new LSP
+            // clients and therefore also require us to retroactively update the corresponding
+            // documents LSP client handle. It's doable but a pretty weird edgecase so let's
+            // wait and see if anyone ever runs into it.
+            tokio::spawn(async move {
+                client.initalize_notify.notified().await;
+                if let Some(workspace_folders_caps) = client
+                    .capabilities()
+                    .workspace
+                    .as_ref()
+                    .and_then(|cap| cap.workspace_folders.as_ref())
+                    .filter(|cap| cap.supported.unwrap_or(false))
+                {
+                    client.add_workspace_folder(
+                        root_uri,
+                        &workspace_folders_caps.change_notifications,
+                    );
+                }
+            });
+            return true;
+        };
+
+        if let Some(workspace_folders_caps) = capabilities
+            .workspace
+            .as_ref()
+            .and_then(|cap| cap.workspace_folders.as_ref())
+            .filter(|cap| cap.supported.unwrap_or(false))
+        {
+            self.add_workspace_folder(root_uri, &workspace_folders_caps.change_notifications);
+            true
+        } else {
+            // the server doesn't support multi workspaces, we need a new client
+            false
+        }
+    }
+
+    fn add_workspace_folder(
+        &self,
+        root_uri: Option<lsp::Url>,
+        change_notifications: &Option<OneOf<bool, String>>,
+    ) {
+        // root_uri is None just means that there isn't really any LSP workspace
+        // associated with this file. For servers that support multiple workspaces
+        // there is just one server so we can always just use that shared instance.
+        // No need to add a new workspace root here as there is no logical root for this file
+        // let the server deal with this
+        let Some(root_uri) = root_uri else {
+            return;
+        };
+
+        // server supports workspace folders, let's add the new root to the list
+        self.workspace_folders
+            .lock()
+            .push(workspace_for_uri(root_uri.clone()));
+        if &Some(OneOf::Left(false)) == change_notifications {
+            // server specifically opted out of DidWorkspaceChange notifications
+            // let's assume the server will request the workspace folders itself
+            // and that we can therefore reuse the client (but are done now)
+            return;
+        }
+        tokio::spawn(self.did_change_workspace(vec![workspace_for_uri(root_uri)], Vec::new()));
+    }
+
     #[allow(clippy::type_complexity)]
     #[allow(clippy::too_many_arguments)]
     pub fn start(
@@ -76,30 +200,25 @@ impl Client {
 
         let (server_rx, server_tx, initialize_notify) =
             Transport::start(reader, writer, stderr, id);
-
-        let root_path = find_root(
+        let (workspace, workspace_is_cwd) = find_workspace();
+        let root = find_lsp_workspace(
             doc_path
                 .and_then(|x| x.parent().and_then(|x| x.to_str()))
                 .unwrap_or("."),
             root_markers,
             manual_roots,
+            &workspace,
+            workspace_is_cwd,
         );
 
-        let root_uri = lsp::Url::from_file_path(root_path.clone()).ok();
+        // `root_uri` and `workspace_folder` can be empty in case there is no workspace
+        // `root_url` can not, use `workspace` as a fallback
+        let root_path = root.clone().unwrap_or_else(|| workspace.clone());
+        let root_uri = root.and_then(|root| lsp::Url::from_file_path(root).ok());
 
-        // TODO: support multiple workspace folders
         let workspace_folders = root_uri
             .clone()
-            .map(|root| {
-                vec![lsp::WorkspaceFolder {
-                    name: root
-                        .path_segments()
-                        .and_then(|segments| segments.last())
-                        .map(|basename| basename.to_string())
-                        .unwrap_or_default(),
-                    uri: root,
-                }]
-            })
+            .map(|root| vec![workspace_for_uri(root)])
             .unwrap_or_default();
 
         let client = Self {
@@ -110,10 +229,10 @@ impl Client {
             capabilities: OnceCell::new(),
             config,
             req_timeout,
-
             root_path,
             root_uri,
-            workspace_folders,
+            workspace_folders: Mutex::new(workspace_folders),
+            initalize_notify: initialize_notify.clone(),
         };
 
         Ok((client, server_rx, initialize_notify))
@@ -169,8 +288,10 @@ impl Client {
         self.config.as_ref()
     }
 
-    pub fn workspace_folders(&self) -> &[lsp::WorkspaceFolder] {
-        &self.workspace_folders
+    pub async fn workspace_folders(
+        &self,
+    ) -> parking_lot::MutexGuard<'_, Vec<lsp::WorkspaceFolder>> {
+        self.workspace_folders.lock()
     }
 
     /// Execute a RPC request on the language server.
@@ -298,7 +419,7 @@ impl Client {
         #[allow(deprecated)]
         let params = lsp::InitializeParams {
             process_id: Some(std::process::id()),
-            workspace_folders: Some(self.workspace_folders.clone()),
+            workspace_folders: Some(self.workspace_folders.lock().clone()),
             // root_path is obsolete, but some clients like pyright still use it so we specify both.
             // clients will prefer _uri if possible
             root_path: self.root_path.to_str().map(|path| path.to_owned()),
@@ -469,6 +590,16 @@ impl Client {
         )
     }
 
+    pub fn did_change_workspace(
+        &self,
+        added: Vec<WorkspaceFolder>,
+        removed: Vec<WorkspaceFolder>,
+    ) -> impl Future<Output = Result<()>> {
+        self.notify::<DidChangeWorkspaceFolders>(DidChangeWorkspaceFoldersParams {
+            event: WorkspaceFoldersChangeEvent { added, removed },
+        })
+    }
+
     // -------------------------------------------------------------------------------------------
     // Text document
     // -------------------------------------------------------------------------------------------
diff --git a/helix-lsp/src/lib.rs b/helix-lsp/src/lib.rs
index e4b00946..d56148a4 100644
--- a/helix-lsp/src/lib.rs
+++ b/helix-lsp/src/lib.rs
@@ -10,15 +10,12 @@ pub use lsp::{Position, Url};
 pub use lsp_types as lsp;
 
 use futures_util::stream::select_all::SelectAll;
-use helix_core::{
-    find_workspace,
-    syntax::{LanguageConfiguration, LanguageServerConfiguration},
-};
+use helix_core::syntax::{LanguageConfiguration, LanguageServerConfiguration};
 use tokio::sync::mpsc::UnboundedReceiver;
 
 use std::{
     collections::{hash_map::Entry, HashMap},
-    path::PathBuf,
+    path::{Path, PathBuf},
     sync::{
         atomic::{AtomicUsize, Ordering},
         Arc,
@@ -609,7 +606,7 @@ impl Notification {
 
 #[derive(Debug)]
 pub struct Registry {
-    inner: HashMap<LanguageId, (usize, Arc<Client>)>,
+    inner: HashMap<LanguageId, Vec<(usize, Arc<Client>)>>,
 
     counter: AtomicUsize,
     pub incoming: SelectAll<UnboundedReceiverStream<(usize, Call)>>,
@@ -633,12 +630,16 @@ impl Registry {
     pub fn get_by_id(&self, id: usize) -> Option<&Client> {
         self.inner
             .values()
+            .flatten()
             .find(|(client_id, _)| client_id == &id)
             .map(|(_, client)| client.as_ref())
     }
 
     pub fn remove_by_id(&mut self, id: usize) {
-        self.inner.retain(|_, (client_id, _)| client_id != &id)
+        self.inner.retain(|_, clients| {
+            clients.retain(|&(client_id, _)| client_id != id);
+            !clients.is_empty()
+        })
     }
 
     pub fn restart(
@@ -664,11 +665,13 @@ impl Registry {
                     start_client(id, language_config, config, doc_path, root_dirs)?;
                 self.incoming.push(UnboundedReceiverStream::new(incoming));
 
-                let (_, old_client) = entry.insert((id, client.clone()));
+                let old_clients = entry.insert(vec![(id, client.clone())]);
 
-                tokio::spawn(async move {
-                    let _ = old_client.force_shutdown().await;
-                });
+                for (_, old_client) in old_clients {
+                    tokio::spawn(async move {
+                        let _ = old_client.force_shutdown().await;
+                    });
+                }
 
                 Ok(Some(client))
             }
@@ -678,10 +681,12 @@ impl Registry {
     pub fn stop(&mut self, language_config: &LanguageConfiguration) {
         let scope = language_config.scope.clone();
 
-        if let Some((_, client)) = self.inner.remove(&scope) {
-            tokio::spawn(async move {
-                let _ = client.force_shutdown().await;
-            });
+        if let Some(clients) = self.inner.remove(&scope) {
+            for (_, client) in clients {
+                tokio::spawn(async move {
+                    let _ = client.force_shutdown().await;
+                });
+            }
         }
     }
 
@@ -696,24 +701,25 @@ impl Registry {
             None => return Ok(None),
         };
 
-        match self.inner.entry(language_config.scope.clone()) {
-            Entry::Occupied(entry) => Ok(Some(entry.get().1.clone())),
-            Entry::Vacant(entry) => {
-                // initialize a new client
-                let id = self.counter.fetch_add(1, Ordering::Relaxed);
-
-                let NewClientResult(client, incoming) =
-                    start_client(id, language_config, config, doc_path, root_dirs)?;
-                self.incoming.push(UnboundedReceiverStream::new(incoming));
-
-                entry.insert((id, client.clone()));
-                Ok(Some(client))
-            }
+        let clients = self.inner.entry(language_config.scope.clone()).or_default();
+        // check if we already have a client for this documents root that we can reuse
+        if let Some((_, client)) = clients.iter_mut().enumerate().find(|(i, (_, client))| {
+            client.try_add_doc(&language_config.roots, root_dirs, doc_path, *i == 0)
+        }) {
+            return Ok(Some(client.1.clone()));
         }
+        // initialize a new client
+        let id = self.counter.fetch_add(1, Ordering::Relaxed);
+
+        let NewClientResult(client, incoming) =
+            start_client(id, language_config, config, doc_path, root_dirs)?;
+        clients.push((id, client.clone()));
+        self.incoming.push(UnboundedReceiverStream::new(incoming));
+        Ok(Some(client))
     }
 
     pub fn iter_clients(&self) -> impl Iterator<Item = &Arc<Client>> {
-        self.inner.values().map(|(_, client)| client)
+        self.inner.values().flatten().map(|(_, client)| client)
     }
 }
 
@@ -850,16 +856,23 @@ fn start_client(
     Ok(NewClientResult(client, incoming))
 }
 
-/// Find an LSP root of a file using the following mechansim:
-/// * start at `file` (either an absolute path or relative to CWD)
-/// * find the top most directory containing a root_marker
-/// * inside the current workspace
-///     * stop the search at the first root_dir that contains `file` or the workspace (obtained from `helix_core::find_workspace`)
-///     * root_dirs only apply inside the workspace. For files outside of the workspace they are ignored
-/// * outside the current workspace: keep searching to the top of the file hiearchy
-pub fn find_root(file: &str, root_markers: &[String], root_dirs: &[PathBuf]) -> PathBuf {
+/// Find an LSP workspace of a file using the following mechanism:
+/// * if the file is outside `workspace` return `None`
+/// * start at `file` and search the file tree upward
+/// * stop the search at the first `root_dirs` entry that contains `file`
+/// * if no `root_dirs` matchs `file` stop at workspace
+/// * Returns the top most directory that contains a `root_marker`
+/// * If no root marker and we stopped at a `root_dirs` entry, return the directory we stopped at
+/// * If we stopped at `workspace` instead and `workspace_is_cwd == false` return `None`
+/// * If we stopped at `workspace` instead and `workspace_is_cwd == true` return `workspace`
+pub fn find_lsp_workspace(
+    file: &str,
+    root_markers: &[String],
+    root_dirs: &[PathBuf],
+    workspace: &Path,
+    workspace_is_cwd: bool,
+) -> Option<PathBuf> {
     let file = std::path::Path::new(file);
-    let workspace = find_workspace();
     let file = if file.is_absolute() {
         file.to_path_buf()
     } else {
@@ -867,7 +880,9 @@ pub fn find_root(file: &str, root_markers: &[String], root_dirs: &[PathBuf]) ->
         current_dir.join(file)
     };
 
-    let inside_workspace = file.strip_prefix(&workspace).is_ok();
+    if !file.starts_with(workspace) {
+        return None;
+    }
 
     let mut top_marker = None;
     for ancestor in file.ancestors() {
@@ -878,18 +893,25 @@ pub fn find_root(file: &str, root_markers: &[String], root_dirs: &[PathBuf]) ->
             top_marker = Some(ancestor);
         }
 
-        if inside_workspace
-            && (ancestor == workspace
-                || root_dirs
-                    .iter()
-                    .any(|root_dir| root_dir == ancestor.strip_prefix(&workspace).unwrap()))
+        if root_dirs
+            .iter()
+            .any(|root_dir| root_dir == ancestor.strip_prefix(workspace).unwrap())
         {
-            return top_marker.unwrap_or(ancestor).to_owned();
+            // if the worskapce is the cwd do not search any higher for workspaces
+            // but specify
+            return Some(top_marker.unwrap_or(workspace).to_owned());
+        }
+        if ancestor == workspace {
+            // if the workspace is the CWD, let the LSP decide what the workspace
+            // is
+            return top_marker
+                .or_else(|| (!workspace_is_cwd).then_some(workspace))
+                .map(Path::to_owned);
         }
     }
 
-    // If no root was found use the workspace as a fallback
-    workspace
+    debug_assert!(false, "workspace must be an ancestor of <file>");
+    None
 }
 
 #[cfg(test)]