From d9235efa76229708586d3c9db3dcbac46127ca0a Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Tue, 2 Apr 2024 11:56:29 +0100
Subject: [PATCH] refactor(lsp): move workspace folder logic into the client

- Changed `reuse_client` to check workspace folders in addition to
  root_dir.
---
 runtime/lua/vim/lsp.lua        | 28 ++++++++++++++++++----
 runtime/lua/vim/lsp/buf.lua    | 40 ++++---------------------------
 runtime/lua/vim/lsp/client.lua | 43 +++++++++++++++++++++++++++++++++-
 3 files changed, 70 insertions(+), 41 deletions(-)

diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index 0c9de607fea6db..eb604caacd8e76 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -164,6 +164,28 @@ local function once(fn)
   end
 end
 
+--- @param client vim.lsp.Client
+--- @param config vim.lsp.ClientConfig
+--- @return boolean
+local function reuse_client_default(client, config)
+  if client.name ~= config.name then
+    return false
+  end
+
+  if config.root_dir then
+    for _, dir in ipairs(client.workspace_folders or {}) do
+      -- note: do not need to check client.root_dir since that should be client.workspace_folders[1]
+      if config.root_dir == dir.name then
+        return true
+      end
+    end
+  end
+
+  -- TODO(lewis6991): also check config.workspace_folders
+
+  return false
+end
+
 --- @class vim.lsp.start.Opts
 --- @inlinedoc
 ---
@@ -216,11 +238,7 @@ end
 --- @return integer? client_id
 function lsp.start(config, opts)
   opts = opts or {}
-  local reuse_client = opts.reuse_client
-    or function(client, conf)
-      return client.root_dir == conf.root_dir and client.name == conf.name
-    end
-
+  local reuse_client = opts.reuse_client or reuse_client_default
   local bufnr = resolve_bufnr(opts.bufnr)
 
   for _, client in pairs(all_clients) do
diff --git a/runtime/lua/vim/lsp/buf.lua b/runtime/lua/vim/lsp/buf.lua
index 43f52b8116026c..17cc698b76c573 100644
--- a/runtime/lua/vim/lsp/buf.lua
+++ b/runtime/lua/vim/lsp/buf.lua
@@ -521,28 +521,9 @@ function M.add_workspace_folder(workspace_folder)
     print(workspace_folder, ' is not a valid directory')
     return
   end
-  local new_workspace = {
-    uri = vim.uri_from_fname(workspace_folder),
-    name = workspace_folder,
-  }
-  local params = { event = { added = { new_workspace }, removed = {} } }
-  local bufnr = vim.api.nvim_get_current_buf()
+  local bufnr = api.nvim_get_current_buf()
   for _, client in pairs(vim.lsp.get_clients({ bufnr = bufnr })) do
-    local found = false
-    for _, folder in pairs(client.workspace_folders or {}) do
-      if folder.name == workspace_folder then
-        found = true
-        print(workspace_folder, 'is already part of this workspace')
-        break
-      end
-    end
-    if not found then
-      client.notify(ms.workspace_didChangeWorkspaceFolders, params)
-      if not client.workspace_folders then
-        client.workspace_folders = {}
-      end
-      table.insert(client.workspace_folders, new_workspace)
-    end
+    client:_add_workspace_folder(workspace_folder)
   end
 end
 
@@ -554,23 +535,12 @@ function M.remove_workspace_folder(workspace_folder)
   workspace_folder = workspace_folder
     or npcall(vim.fn.input, 'Workspace Folder: ', vim.fn.expand('%:p:h'))
   api.nvim_command('redraw')
-  if not (workspace_folder and #workspace_folder > 0) then
+  if not workspace_folder or #workspace_folder == 0 then
     return
   end
-  local workspace = {
-    uri = vim.uri_from_fname(workspace_folder),
-    name = workspace_folder,
-  }
-  local params = { event = { added = {}, removed = { workspace } } }
-  local bufnr = vim.api.nvim_get_current_buf()
+  local bufnr = api.nvim_get_current_buf()
   for _, client in pairs(vim.lsp.get_clients({ bufnr = bufnr })) do
-    for idx, folder in pairs(client.workspace_folders) do
-      if folder.name == workspace_folder then
-        client.notify(ms.workspace_didChangeWorkspaceFolders, params)
-        client.workspace_folders[idx] = nil
-        return
-      end
-    end
+    client:_remove_workspace_folder(workspace_folder)
   end
   print(workspace_folder, 'is not currently part of the workspace')
 end
diff --git a/runtime/lua/vim/lsp/client.lua b/runtime/lua/vim/lsp/client.lua
index 303fc55982b3a7..f73f97b8cdbad1 100644
--- a/runtime/lua/vim/lsp/client.lua
+++ b/runtime/lua/vim/lsp/client.lua
@@ -420,7 +420,7 @@ local function get_workspace_folders(workspace_folders, root_dir)
     return {
       {
         uri = vim.uri_from_fname(root_dir),
-        name = string.format('%s', root_dir),
+        name = root_dir,
       },
     }
   end
@@ -1065,4 +1065,45 @@ function Client:_on_exit(code, signal)
   )
 end
 
+--- @package
+--- Add a directory to the workspace folders.
+--- @param dir string?
+function Client:_add_workspace_folder(dir)
+  for _, folder in pairs(self.workspace_folders or {}) do
+    if folder.name == dir then
+      print(dir, 'is already part of this workspace')
+      return
+    end
+  end
+
+  local wf = assert(get_workspace_folders(nil, dir))
+
+  self:_notify(ms.workspace_didChangeWorkspaceFolders, {
+    event = { added = wf, removed = {} },
+  })
+
+  if not self.workspace_folders then
+    self.workspace_folders = {}
+  end
+  vim.list_extend(self.workspace_folders, wf)
+end
+
+--- @package
+--- Remove a directory to the workspace folders.
+--- @param dir string?
+function Client:_remove_workspace_folder(dir)
+  local wf = assert(get_workspace_folders(nil, dir))
+
+  self:_notify(ms.workspace_didChangeWorkspaceFolders, {
+    event = { added = {}, removed = wf },
+  })
+
+  for idx, folder in pairs(self.workspace_folders) do
+    if folder.name == dir then
+      table.remove(self.workspace_folders, idx)
+      break
+    end
+  end
+end
+
 return Client