feat(lsp): support function for client root_dir (#31630)

If root_dir is a function it is evaluated when the client is created to
determine the root directory.

This enables dynamically determining the root directory based on e.g.
project or directory structure (example: finding a parent Cargo.toml
file that contains "[workspace]" in a Rust project).

3 files changed, 62 insertions, 5 deletions

diff --git a/runtime/doc/lsp.txt b/runtime/doc/lsp.txt
index 8b822daf9e..16f543088b 100644
--- a/runtime/doc/lsp.txt
+++ b/runtime/doc/lsp.txt
@@ -683,6 +683,13 @@ Lua module: vim.lsp                                                 *lsp-core*
                          the LSP server will base its workspaceFolders,
                          rootUri, and rootPath on initialization. Unused if
                          `root_dir` is provided.
+      • {root_dir}?      (`string|fun(cb:fun(string))`) Directory where the
+                         LSP server will base its workspaceFolders, rootUri,
+                         and rootPath on initialization. If a function, it
+                         accepts a single callback argument which must be
+                         called with the value of root_dir to use. The LSP
+                         server will not be started until the callback is
+                         called.
       • {reuse_client}?  (`fun(client: vim.lsp.Client, config: vim.lsp.ClientConfig): boolean`)
                          Predicate used to decide if a client should be
                          re-used. Used on all running clients. The default
diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index 19d0377585..1c8356d64d 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -334,6 +334,11 @@ end
 --- rootUri, and rootPath on initialization. Unused if `root_dir` is provided.
 --- @field root_markers? string[]
 ---
+--- Directory where the LSP server will base its workspaceFolders, rootUri, and rootPath on
+--- initialization. If a function, it accepts a single callback argument which must be called with
+--- the value of root_dir to use. The LSP server will not be started until the callback is called.
+--- @field root_dir? string|fun(cb:fun(string))
+---
 --- Predicate used to decide if a client should be re-used. Used on all
 --- running clients. The default implementation re-uses a client if name and
 --- root_dir matches.
@@ -499,6 +504,15 @@ local function lsp_enable_callback(bufnr)
     return true
   end
 
+  --- @param config vim.lsp.Config
+  local function start(config)
+    return vim.lsp.start(config, {
+      bufnr = bufnr,
+      reuse_client = config.reuse_client,
+      _root_markers = config.root_markers,
+    })
+  end
+
   for name in vim.spairs(lsp._enabled_configs) do
     local config = lsp._resolve_config(name)
 
@@ -507,11 +521,14 @@ local function lsp_enable_callback(bufnr)
       -- do not propagate back to the enabled configs.
       config = vim.deepcopy(config)
 
-      vim.lsp.start(config, {
-        bufnr = bufnr,
-        reuse_client = config.reuse_client,
-        _root_markers = config.root_markers,
-      })
+      if type(config.root_dir) == 'function' then
+        config.root_dir(function(root_dir)
+          config.root_dir = root_dir
+          start(config)
+        end)
+      else
+        start(config)
+      end
     end
   end
 end
diff --git a/test/functional/plugin/lsp_spec.lua b/test/functional/plugin/lsp_spec.lua
index 1f246b0914..f396c837f9 100644
--- a/test/functional/plugin/lsp_spec.lua
+++ b/test/functional/plugin/lsp_spec.lua
@@ -6245,5 +6245,38 @@ describe('LSP', function()
         end)
       )
     end)
+
+    it('supports a function for root_dir', function()
+      exec_lua(create_server_definition)
+
+      local tmp1 = t.tmpname(true)
+
+      eq(
+        'some_dir',
+        exec_lua(function()
+          local server = _G._create_server({
+            handlers = {
+              initialize = function(_, _, callback)
+                callback(nil, { capabilities = {} })
+              end,
+            },
+          })
+
+          vim.lsp.config('foo', {
+            cmd = server.cmd,
+            filetypes = { 'foo' },
+            root_dir = function(cb)
+              cb('some_dir')
+            end,
+          })
+          vim.lsp.enable('foo')
+
+          vim.cmd.edit(assert(tmp1))
+          vim.bo.filetype = 'foo'
+
+          return vim.lsp.get_clients({ bufnr = vim.api.nvim_get_current_buf() })[1].root_dir
+        end)
+      )
+    end)
   end)
 end)