refactor(lsp): fold in dynamic_registration code into the client

Problem:

Capability register logic is spread across 3 files.

Solution:

- Consolidate (and simplify) logic into the client.
- Teach client.supports_method about resolve methods

7 files changed, 188 insertions, 189 deletions

diff --git a/runtime/doc/lsp.txt b/runtime/doc/lsp.txt
index 7d50cb52eb..b7a4f0ea01 100644
--- a/runtime/doc/lsp.txt
+++ b/runtime/doc/lsp.txt
@@ -968,7 +968,7 @@ Lua module: vim.lsp.client                                        *lsp-client*
                                   request before sending kill -15. If set to
                                   false, nvim exits immediately after sending
                                   the "shutdown" request to the server.
-      • {get_language_id}       (`fun(bufnr: integer, filetype: string): string`)
+      • {get_language_id}       (`fun(bufnr: integer, filetype?: string): string`)
       • {capabilities}          (`lsp.ClientCapabilities`) The capabilities
                                 provided by the client (editor or tool)
       • {dynamic_capabilities}  (`lsp.DynamicCapabilities`)
@@ -1089,8 +1089,9 @@ Lua module: vim.lsp.client                                        *lsp-client*
                               `initialize` in the LSP spec.
       • {name}?               (`string`, default: client-id) Name in log
                               messages.
-      • {get_language_id}?    (`fun(bufnr: integer, filetype: string): string`)
-                              Language ID as string. Defaults to the filetype.
+      • {get_language_id}?    (`fun(bufnr: integer, filetype?: string): string`)
+                              Language ID as string. Defaults to the buffer
+                              filetype.
       • {offset_encoding}?    (`'utf-8'|'utf-16'|'utf-32'`) The encoding that
                               the LSP server expects. Client does not verify
                               this is correct.
diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index 42a0ccc3d4..0de3b4ee4d 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -3,7 +3,6 @@ local validate = vim.validate
 
 local lsp = vim._defer_require('vim.lsp', {
   _changetracking = ..., --- @module 'vim.lsp._changetracking'
-  _dynamic = ..., --- @module 'vim.lsp._dynamic'
   _snippet_grammar = ..., --- @module 'vim.lsp._snippet_grammar'
   _tagfunc = ..., --- @module 'vim.lsp._tagfunc'
   _watchfiles = ..., --- @module 'vim.lsp._watchfiles'
@@ -31,6 +30,13 @@ local changetracking = lsp._changetracking
 ---@nodoc
 lsp.rpc_response_error = lsp.rpc.rpc_response_error
 
+lsp._resolve_to_request = {
+  [ms.codeAction_resolve] = ms.textDocument_codeAction,
+  [ms.codeLens_resolve] = ms.textDocument_codeLens,
+  [ms.documentLink_resolve] = ms.textDocument_documentLink,
+  [ms.inlayHint_resolve] = ms.textDocument_inlayHint,
+}
+
 -- maps request name to the required server_capability in the client.
 lsp._request_name_to_capability = {
   [ms.callHierarchy_incomingCalls] = { 'callHierarchyProvider' },
diff --git a/runtime/lua/vim/lsp/_dynamic.lua b/runtime/lua/vim/lsp/_dynamic.lua
deleted file mode 100644
index 27113c0e74..0000000000
--- a/runtime/lua/vim/lsp/_dynamic.lua
+++ /dev/null
@@ -1,110 +0,0 @@
-local glob = vim.glob
-
---- @class lsp.DynamicCapabilities
---- @field capabilities table<string, lsp.Registration[]>
---- @field client_id number
-local M = {}
-
---- @param client_id number
---- @return lsp.DynamicCapabilities
-function M.new(client_id)
-  return setmetatable({
-    capabilities = {},
-    client_id = client_id,
-  }, { __index = M })
-end
-
-function M:supports_registration(method)
-  local client = vim.lsp.get_client_by_id(self.client_id)
-  if not client then
-    return false
-  end
-  local capability = vim.tbl_get(client.capabilities, unpack(vim.split(method, '/')))
-  return type(capability) == 'table' and capability.dynamicRegistration
-end
-
---- @param registrations lsp.Registration[]
-function M:register(registrations)
-  -- remove duplicates
-  self:unregister(registrations)
-  for _, reg in ipairs(registrations) do
-    local method = reg.method
-    if not self.capabilities[method] then
-      self.capabilities[method] = {}
-    end
-    table.insert(self.capabilities[method], reg)
-  end
-end
-
---- @param unregisterations lsp.Unregistration[]
-function M:unregister(unregisterations)
-  for _, unreg in ipairs(unregisterations) do
-    local method = unreg.method
-    if not self.capabilities[method] then
-      return
-    end
-    local id = unreg.id
-    for i, reg in ipairs(self.capabilities[method]) do
-      if reg.id == id then
-        table.remove(self.capabilities[method], i)
-        break
-      end
-    end
-  end
-end
-
---- @param method string
---- @param opts? {bufnr: integer?}
---- @return lsp.Registration? (table|nil) the registration if found
-function M:get(method, opts)
-  opts = opts or {}
-  opts.bufnr = opts.bufnr or vim.api.nvim_get_current_buf()
-  for _, reg in ipairs(self.capabilities[method] or {}) do
-    if not reg.registerOptions then
-      return reg
-    end
-    local documentSelector = reg.registerOptions.documentSelector
-    if not documentSelector then
-      return reg
-    end
-    if self:match(opts.bufnr, documentSelector) then
-      return reg
-    end
-  end
-end
-
---- @param method string
---- @param opts? {bufnr: integer?}
-function M:supports(method, opts)
-  return self:get(method, opts) ~= nil
-end
-
---- @param bufnr number
---- @param documentSelector lsp.DocumentSelector
---- @private
-function M:match(bufnr, documentSelector)
-  local client = vim.lsp.get_client_by_id(self.client_id)
-  if not client then
-    return false
-  end
-  local language = client.get_language_id(bufnr, vim.bo[bufnr].filetype)
-  local uri = vim.uri_from_bufnr(bufnr)
-  local fname = vim.uri_to_fname(uri)
-  for _, filter in ipairs(documentSelector) do
-    local matches = true
-    if filter.language and language ~= filter.language then
-      matches = false
-    end
-    if matches and filter.scheme and not vim.startswith(uri, filter.scheme .. ':') then
-      matches = false
-    end
-    if matches and filter.pattern and not glob.to_lpeg(filter.pattern):match(fname) then
-      matches = false
-    end
-    if matches then
-      return true
-    end
-  end
-end
-
-return M
diff --git a/runtime/lua/vim/lsp/_watchfiles.lua b/runtime/lua/vim/lsp/_watchfiles.lua
index 98e9818bcd..c4cdb5aea8 100644
--- a/runtime/lua/vim/lsp/_watchfiles.lua
+++ b/runtime/lua/vim/lsp/_watchfiles.lua
@@ -44,9 +44,8 @@ M._poll_exclude_pattern = glob.to_lpeg('**/.git/{objects,subtree-cache}/**')
 --- Registers the workspace/didChangeWatchedFiles capability dynamically.
 ---
 ---@param reg lsp.Registration LSP Registration object.
----@param ctx lsp.HandlerContext Context from the |lsp-handler|.
-function M.register(reg, ctx)
-  local client_id = ctx.client_id
+---@param client_id integer Client ID.
+function M.register(reg, client_id)
   local client = assert(vim.lsp.get_client_by_id(client_id), 'Client must be running')
   -- Ill-behaved servers may not honor the client capability and try to register
   -- anyway, so ignore requests when the user has opted out of the feature.
@@ -155,9 +154,8 @@ end
 --- Unregisters the workspace/didChangeWatchedFiles capability dynamically.
 ---
 ---@param unreg lsp.Unregistration LSP Unregistration object.
----@param ctx lsp.HandlerContext Context from the |lsp-handler|.
-function M.unregister(unreg, ctx)
-  local client_id = ctx.client_id
+---@param client_id integer Client ID.
+function M.unregister(unreg, client_id)
   local client_cancels = cancels[client_id]
   local reg_cancels = client_cancels[unreg.id]
   while #reg_cancels > 0 do
diff --git a/runtime/lua/vim/lsp/buf.lua b/runtime/lua/vim/lsp/buf.lua
index a75e322e90..6383855a30 100644
--- a/runtime/lua/vim/lsp/buf.lua
+++ b/runtime/lua/vim/lsp/buf.lua
@@ -1131,12 +1131,7 @@ local function on_code_action_results(results, opts)
     local action = choice.action
     local bufnr = assert(choice.ctx.bufnr, 'Must have buffer number')
 
-    local reg = client.dynamic_capabilities:get(ms.textDocument_codeAction, { bufnr = bufnr })
-
-    local supports_resolve = vim.tbl_get(reg or {}, 'registerOptions', 'resolveProvider')
-      or client.supports_method(ms.codeAction_resolve)
-
-    if not action.edit and client and supports_resolve then
+    if not action.edit and client.supports_method(ms.codeAction_resolve) then
       client.request(ms.codeAction_resolve, action, function(err, resolved_action)
         if err then
           if action.command then
diff --git a/runtime/lua/vim/lsp/client.lua b/runtime/lua/vim/lsp/client.lua
index 2718f40c96..ba12447c40 100644
--- a/runtime/lua/vim/lsp/client.lua
+++ b/runtime/lua/vim/lsp/client.lua
@@ -91,8 +91,8 @@ local validate = vim.validate
 --- (default: client-id)
 --- @field name? string
 ---
---- Language ID as string. Defaults to the filetype.
---- @field get_language_id? fun(bufnr: integer, filetype: string): string
+--- Language ID as string. Defaults to the buffer filetype.
+--- @field get_language_id? fun(bufnr: integer, filetype?: string): string
 ---
 --- The encoding that the LSP server expects. Client does not verify this is correct.
 --- @field offset_encoding? 'utf-8'|'utf-16'|'utf-32'
@@ -212,10 +212,11 @@ local validate = vim.validate
 --- A table with flags for the client. The current (experimental) flags are:
 --- @field flags vim.lsp.Client.Flags
 ---
---- @field get_language_id fun(bufnr: integer, filetype: string): string
+--- @field get_language_id fun(bufnr: integer, filetype?: string): string
 ---
 --- The capabilities provided by the client (editor or tool)
 --- @field capabilities lsp.ClientCapabilities
+--- @field private registrations table<string,lsp.Registration[]>
 --- @field dynamic_capabilities lsp.DynamicCapabilities
 ---
 --- Sends a request to the server.
@@ -339,10 +340,10 @@ end
 --- By default, get_language_id just returns the exact filetype it is passed.
 --- It is possible to pass in something that will calculate a different filetype,
 --- to be sent by the client.
---- @param _bufnr integer
---- @param filetype string
-local function default_get_language_id(_bufnr, filetype)
-  return filetype
+--- @param bufnr integer
+--- @param filetype? string
+local function default_get_language_id(bufnr, filetype)
+  return filetype or vim.bo[bufnr].filetype
 end
 
 --- Validates a client configuration as given to |vim.lsp.start_client()|.
@@ -403,18 +404,16 @@ local function get_name(id, config)
   return tostring(id)
 end
 
---- @param workspace_folders lsp.WorkspaceFolder[]?
---- @param root_dir string?
+--- @param workspace_folders string|lsp.WorkspaceFolder[]?
 --- @return lsp.WorkspaceFolder[]?
-local function get_workspace_folders(workspace_folders, root_dir)
-  if workspace_folders then
+local function get_workspace_folders(workspace_folders)
+  if type(workspace_folders) == 'table' then
     return workspace_folders
-  end
-  if root_dir then
+  elseif type(workspace_folders) == 'string' then
     return {
       {
-        uri = vim.uri_from_fname(root_dir),
-        name = root_dir,
+        uri = vim.uri_from_fname(workspace_folders),
+        name = workspace_folders,
       },
     }
   end
@@ -451,13 +450,13 @@ function Client.create(config)
     requests = {},
     attached_buffers = {},
     server_capabilities = {},
-    dynamic_capabilities = lsp._dynamic.new(id),
+    registrations = {},
     commands = config.commands or {},
     settings = config.settings or {},
     flags = config.flags or {},
     get_language_id = config.get_language_id or default_get_language_id,
     capabilities = config.capabilities or lsp.protocol.make_client_capabilities(),
-    workspace_folders = get_workspace_folders(config.workspace_folders, config.root_dir),
+    workspace_folders = get_workspace_folders(config.workspace_folders or config.root_dir),
     root_dir = config.root_dir,
     _before_init_cb = config.before_init,
     _on_init_cbs = ensure_list(config.on_init),
@@ -478,6 +477,28 @@ function Client.create(config)
     messages = { name = name, messages = {}, progress = {}, status = {} },
   }
 
+  --- @class lsp.DynamicCapabilities
+  --- @nodoc
+  self.dynamic_capabilities = {
+    capabilities = self.registrations,
+    client_id = id,
+    register = function(_, registrations)
+      return self:_register_dynamic(registrations)
+    end,
+    unregister = function(_, unregistrations)
+      return self:_unregister_dynamic(unregistrations)
+    end,
+    get = function(_, method, opts)
+      return self:_get_registration(method, opts and opts.bufnr)
+    end,
+    supports_registration = function(_, method)
+      return self:_supports_registration(method)
+    end,
+    supports = function(_, method, opts)
+      return self:_get_registration(method, opts and opts.bufnr) ~= nil
+    end,
+  }
+
   self.request = method_wrapper(self, Client._request)
   self.request_sync = method_wrapper(self, Client._request_sync)
   self.notify = method_wrapper(self, Client._notify)
@@ -846,6 +867,100 @@ function Client:_stop(force)
   end)
 end
 
+--- Get options for a method that is registered dynamically.
+--- @param method string
+function Client:_supports_registration(method)
+  local capability = vim.tbl_get(self.capabilities, unpack(vim.split(method, '/')))
+  return type(capability) == 'table' and capability.dynamicRegistration
+end
+
+--- @private
+--- @param registrations lsp.Registration[]
+function Client:_register_dynamic(registrations)
+  -- remove duplicates
+  self:_unregister_dynamic(registrations)
+  for _, reg in ipairs(registrations) do
+    local method = reg.method
+    if not self.registrations[method] then
+      self.registrations[method] = {}
+    end
+    table.insert(self.registrations[method], reg)
+  end
+end
+
+--- @param registrations lsp.Registration[]
+function Client:_register(registrations)
+  self:_register_dynamic(registrations)
+
+  local unsupported = {} --- @type string[]
+
+  for _, reg in ipairs(registrations) do
+    local method = reg.method
+    if method == ms.workspace_didChangeWatchedFiles then
+      vim.lsp._watchfiles.register(reg, self.id)
+    elseif not self:_supports_registration(method) then
+      unsupported[#unsupported + 1] = method
+    end
+  end
+
+  if #unsupported > 0 then
+    local warning_tpl = 'The language server %s triggers a registerCapability '
+      .. 'handler for %s despite dynamicRegistration set to false. '
+      .. 'Report upstream, this warning is harmless'
+    log.warn(string.format(warning_tpl, self.name, table.concat(unsupported, ', ')))
+  end
+end
+
+--- @private
+--- @param unregistrations lsp.Unregistration[]
+function Client:_unregister_dynamic(unregistrations)
+  for _, unreg in ipairs(unregistrations) do
+    local sreg = self.registrations[unreg.method]
+    -- Unegister dynamic capability
+    for i, reg in ipairs(sreg or {}) do
+      if reg.id == unreg.id then
+        table.remove(sreg, i)
+        break
+      end
+    end
+  end
+end
+
+--- @param unregistrations lsp.Unregistration[]
+function Client:_unregister(unregistrations)
+  self:_unregister_dynamic(unregistrations)
+  for _, unreg in ipairs(unregistrations) do
+    if unreg.method == ms.workspace_didChangeWatchedFiles then
+      vim.lsp._watchfiles.unregister(unreg, self.id)
+    end
+  end
+end
+
+--- @param method string
+--- @param bufnr? integer
+--- @return lsp.Registration?
+function Client:_get_registration(method, bufnr)
+  bufnr = bufnr or vim.api.nvim_get_current_buf()
+  for _, reg in ipairs(self.registrations[method] or {}) do
+    if not reg.registerOptions or not reg.registerOptions.documentSelector then
+      return reg
+    end
+    local documentSelector = reg.registerOptions.documentSelector
+    local language = self.get_language_id(bufnr)
+    local uri = vim.uri_from_bufnr(bufnr)
+    local fname = vim.uri_to_fname(uri)
+    for _, filter in ipairs(documentSelector) do
+      if
+        not (filter.language and language ~= filter.language)
+        and not (filter.scheme and not vim.startswith(uri, filter.scheme .. ':'))
+        and not (filter.pattern and not vim.glob.to_lpeg(filter.pattern):match(fname))
+      then
+        return reg
+      end
+    end
+  end
+end
+
 --- @private
 --- Checks whether a client is stopped.
 ---
@@ -908,12 +1023,11 @@ function Client:_text_document_did_open_handler(bufnr)
     return
   end
 
-  local filetype = vim.bo[bufnr].filetype
   self.notify(ms.textDocument_didOpen, {
     textDocument = {
       version = lsp.util.buf_versions[bufnr],
       uri = vim.uri_from_bufnr(bufnr),
-      languageId = self.get_language_id(bufnr, filetype),
+      languageId = self.get_language_id(bufnr),
       text = lsp._buf_get_full_text(bufnr),
     },
   })
@@ -978,12 +1092,37 @@ function Client:_supports_method(method, opts)
   if vim.tbl_get(self.server_capabilities, unpack(required_capability)) then
     return true
   end
-  if self.dynamic_capabilities:supports_registration(method) then
-    return self.dynamic_capabilities:supports(method, opts)
+
+  local rmethod = lsp._resolve_to_request[method]
+  if rmethod then
+    if self:_supports_registration(rmethod) then
+      local reg = self:_get_registration(rmethod, opts and opts.bufnr)
+      return vim.tbl_get(reg or {}, 'registerOptions', 'resolveProvider') or false
+    end
+  else
+    if self:_supports_registration(method) then
+      return self:_get_registration(method, opts and opts.bufnr) ~= nil
+    end
   end
   return false
 end
 
+--- Get options for a method that is registered dynamically.
+--- @param method string
+--- @param bufnr? integer
+--- @return lsp.LSPAny?
+function Client:_get_registration_options(method, bufnr)
+  if not self:_supports_registration(method) then
+    return
+  end
+
+  local reg = self:_get_registration(method, bufnr)
+
+  if reg then
+    return reg.registerOptions
+  end
+end
+
 --- @private
 --- Handles a notification sent by an LSP server by invoking the
 --- corresponding handler.
@@ -1061,7 +1200,7 @@ function Client:_add_workspace_folder(dir)
     end
   end
 
-  local wf = assert(get_workspace_folders(nil, dir))
+  local wf = assert(get_workspace_folders(dir))
 
   self:_notify(ms.workspace_didChangeWorkspaceFolders, {
     event = { added = wf, removed = {} },
@@ -1076,7 +1215,7 @@ end
 --- Remove a directory to the workspace folders.
 --- @param dir string?
 function Client:_remove_workspace_folder(dir)
-  local wf = assert(get_workspace_folders(nil, dir))
+  local wf = assert(get_workspace_folders(dir))
 
   self:_notify(ms.workspace_didChangeWorkspaceFolders, {
     event = { added = {}, removed = wf },
diff --git a/runtime/lua/vim/lsp/handlers.lua b/runtime/lua/vim/lsp/handlers.lua
index 2b7aefe0e1..5c28d88b38 100644
--- a/runtime/lua/vim/lsp/handlers.lua
+++ b/runtime/lua/vim/lsp/handlers.lua
@@ -122,46 +122,19 @@ end
 --- @see # https://microsoft.github.io/language-server-protocol/specifications/specification-current/#client_registerCapability
 --- @param params lsp.RegistrationParams
 RSC[ms.client_registerCapability] = function(_, params, ctx)
-  local client_id = ctx.client_id
-  local client = assert(vim.lsp.get_client_by_id(client_id))
-
-  client.dynamic_capabilities:register(params.registrations)
-  for bufnr, _ in pairs(client.attached_buffers) do
+  local client = assert(vim.lsp.get_client_by_id(ctx.client_id))
+  client:_register(params.registrations)
+  for bufnr in pairs(client.attached_buffers) do
     vim.lsp._set_defaults(client, bufnr)
   end
-
-  ---@type string[]
-  local unsupported = {}
-  for _, reg in ipairs(params.registrations) do
-    if reg.method == ms.workspace_didChangeWatchedFiles then
-      vim.lsp._watchfiles.register(reg, ctx)
-    elseif not client.dynamic_capabilities:supports_registration(reg.method) then
-      unsupported[#unsupported + 1] = reg.method
-    end
-  end
-  if #unsupported > 0 then
-    local warning_tpl = 'The language server %s triggers a registerCapability '
-      .. 'handler for %s despite dynamicRegistration set to false. '
-      .. 'Report upstream, this warning is harmless'
-    local client_name = client and client.name or string.format('id=%d', client_id)
-    local warning = string.format(warning_tpl, client_name, table.concat(unsupported, ', '))
-    log.warn(warning)
-  end
   return vim.NIL
 end
 
 --- @see # https://microsoft.github.io/language-server-protocol/specifications/specification-current/#client_unregisterCapability
 --- @param params lsp.UnregistrationParams
 RSC[ms.client_unregisterCapability] = function(_, params, ctx)
-  local client_id = ctx.client_id
-  local client = assert(vim.lsp.get_client_by_id(client_id))
-  client.dynamic_capabilities:unregister(params.unregisterations)
-
-  for _, unreg in ipairs(params.unregisterations) do
-    if unreg.method == ms.workspace_didChangeWatchedFiles then
-      vim.lsp._watchfiles.unregister(unreg, ctx)
-    end
-  end
+  local client = assert(vim.lsp.get_client_by_id(ctx.client_id))
+  client:_unregister(params.unregisterations)
   return vim.NIL
 end
 
@@ -173,8 +146,7 @@ RSC[ms.workspace_applyEdit] = function(_, params, ctx)
     'workspace/applyEdit must be called with `ApplyWorkspaceEditParams`. Server is violating the specification'
   )
   -- TODO(ashkan) Do something more with label?
-  local client_id = ctx.client_id
-  local client = assert(vim.lsp.get_client_by_id(client_id))
+  local client = assert(vim.lsp.get_client_by_id(ctx.client_id))
   if params.label then
     print('Workspace edit', params.label)
   end
@@ -196,12 +168,11 @@ end
 --- @see # https://microsoft.github.io/language-server-protocol/specifications/specification-current/#workspace_configuration
 --- @param params lsp.ConfigurationParams
 RSC[ms.workspace_configuration] = function(_, params, ctx)
-  local client_id = ctx.client_id
-  local client = vim.lsp.get_client_by_id(client_id)
+  local client = vim.lsp.get_client_by_id(ctx.client_id)
   if not client then
     err_message(
       'LSP[',
-      client_id,
+      ctx.client_id,
       '] client has shut down after sending a workspace/configuration request'
     )
     return
@@ -229,10 +200,9 @@ end
 
 --- @see # https://microsoft.github.io/language-server-protocol/specifications/specification-current/#workspace_workspaceFolders
 RSC[ms.workspace_workspaceFolders] = function(_, _, ctx)
-  local client_id = ctx.client_id
-  local client = vim.lsp.get_client_by_id(client_id)
+  local client = vim.lsp.get_client_by_id(ctx.client_id)
   if not client then
-    err_message('LSP[id=', client_id, '] client has shut down after sending the message')
+    err_message('LSP[id=', ctx.client_id, '] client has shut down after sending the message')
     return
   end
   return client.workspace_folders or vim.NIL