feat(lsp): use the meta model to generate server capability map

4 files changed, 153 insertions, 94 deletions

diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index a45f9adeb6..cbec06df02 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -38,55 +38,6 @@ lsp._resolve_to_request = {
   [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' },
-  [ms.callHierarchy_outgoingCalls] = { 'callHierarchyProvider' },
-  [ms.codeAction_resolve] = { 'codeActionProvider', 'resolveProvider' },
-  [ms.codeLens_resolve] = { 'codeLensProvider', 'resolveProvider' },
-  [ms.documentLink_resolve] = { 'documentLinkProvider', 'resolveProvider' },
-  [ms.inlayHint_resolve] = { 'inlayHintProvider', 'resolveProvider' },
-  [ms.textDocument_codeAction] = { 'codeActionProvider' },
-  [ms.textDocument_codeLens] = { 'codeLensProvider' },
-  [ms.textDocument_completion] = { 'completionProvider' },
-  [ms.textDocument_declaration] = { 'declarationProvider' },
-  [ms.textDocument_definition] = { 'definitionProvider' },
-  [ms.textDocument_diagnostic] = { 'diagnosticProvider' },
-  [ms.textDocument_didClose] = { 'textDocumentSync', 'openClose' },
-  [ms.textDocument_didOpen] = { 'textDocumentSync', 'openClose' },
-  [ms.textDocument_documentColor] = { 'colorProvider' },
-  [ms.textDocument_documentHighlight] = { 'documentHighlightProvider' },
-  [ms.textDocument_documentLink] = { 'documentLinkProvider' },
-  [ms.textDocument_documentSymbol] = { 'documentSymbolProvider' },
-  [ms.textDocument_foldingRange] = { 'foldingRangeProvider' },
-  [ms.textDocument_formatting] = { 'documentFormattingProvider' },
-  [ms.textDocument_hover] = { 'hoverProvider' },
-  [ms.textDocument_implementation] = { 'implementationProvider' },
-  [ms.textDocument_inlayHint] = { 'inlayHintProvider' },
-  [ms.textDocument_inlineValue] = { 'inlineValueProvider' },
-  [ms.textDocument_linkedEditingRange] = { 'linkedEditingRangeProvider' },
-  [ms.textDocument_moniker] = { 'monikerProvider' },
-  [ms.textDocument_onTypeFormatting] = { 'documentOnTypeFormattingProvider' },
-  [ms.textDocument_prepareCallHierarchy] = { 'callHierarchyProvider' },
-  [ms.textDocument_prepareRename] = { 'renameProvider', 'prepareProvider' },
-  [ms.textDocument_prepareTypeHierarchy] = { 'typeHierarchyProvider' },
-  [ms.textDocument_rangeFormatting] = { 'documentRangeFormattingProvider' },
-  [ms.textDocument_rangesFormatting] = { 'documentRangeFormattingProvider', 'rangesSupport' },
-  [ms.textDocument_references] = { 'referencesProvider' },
-  [ms.textDocument_rename] = { 'renameProvider' },
-  [ms.textDocument_selectionRange] = { 'selectionRangeProvider' },
-  [ms.textDocument_semanticTokens_full] = { 'semanticTokensProvider' },
-  [ms.textDocument_semanticTokens_full_delta] = { 'semanticTokensProvider' },
-  [ms.textDocument_signatureHelp] = { 'signatureHelpProvider' },
-  [ms.textDocument_typeDefinition] = { 'typeDefinitionProvider' },
-  [ms.textDocument_willSaveWaitUntil] = { 'textDocumentSync', 'willSaveWaitUntil' },
-  [ms.textDocument_willSave] = { 'textDocumentSync', 'willSave' },
-  [ms.typeHierarchy_subtypes] = { 'typeHierarchyProvider' },
-  [ms.typeHierarchy_supertypes] = { 'typeHierarchyProvider' },
-  [ms.workspace_executeCommand] = { 'executeCommandProvider' },
-  [ms.workspace_symbol] = { 'workspaceSymbolProvider' },
-}
-
 -- TODO improve handling of scratch buffers with LSP attached.
 
 ---@private
diff --git a/runtime/lua/vim/lsp/client.lua b/runtime/lua/vim/lsp/client.lua
index 253ccc48f4..79a91c3d5f 100644
--- a/runtime/lua/vim/lsp/client.lua
+++ b/runtime/lua/vim/lsp/client.lua
@@ -1056,7 +1056,7 @@ function Client:supports_method(method, bufnr)
     --- @diagnostic disable-next-line:no-unknown
     bufnr = bufnr.bufnr
   end
-  local required_capability = lsp._request_name_to_capability[method]
+  local required_capability = lsp.protocol._request_name_to_capability[method]
   -- if we don't know about the method, assume that the client supports it.
   if not required_capability then
     return true
diff --git a/runtime/lua/vim/lsp/protocol.lua b/runtime/lua/vim/lsp/protocol.lua
index 6895098126..7975006f9d 100644
--- a/runtime/lua/vim/lsp/protocol.lua
+++ b/runtime/lua/vim/lsp/protocol.lua
@@ -725,7 +725,6 @@ end
 --- | vim.lsp.protocol.Method.ServerToClient
 
 -- Generated by gen_lsp.lua, keep at end of file.
----
 --- @enum vim.lsp.protocol.Methods
 --- @see https://microsoft.github.io/language-server-protocol/specification/#metaModel
 --- LSP method names.
@@ -1080,4 +1079,68 @@ protocol.Methods = {
   workspace_workspaceFolders = 'workspace/workspaceFolders',
 }
 
+-- stylua: ignore start
+-- Generated by gen_lsp.lua, keep at end of file.
+--- Maps method names to the required server capability
+protocol._request_name_to_capability = {
+  ['codeAction/resolve'] = { 'codeActionProvider', 'resolveProvider' },
+  ['codeLens/resolve'] = { 'codeLensProvider', 'resolveProvider' },
+  ['completionItem/resolve'] = { 'completionProvider', 'resolveProvider' },
+  ['documentLink/resolve'] = { 'documentLinkProvider', 'resolveProvider' },
+  ['inlayHint/resolve'] = { 'inlayHintProvider', 'resolveProvider' },
+  ['textDocument/codeAction'] = { 'codeActionProvider' },
+  ['textDocument/codeLens'] = { 'codeLensProvider' },
+  ['textDocument/completion'] = { 'completionProvider' },
+  ['textDocument/declaration'] = { 'declarationProvider' },
+  ['textDocument/definition'] = { 'definitionProvider' },
+  ['textDocument/diagnostic'] = { 'diagnosticProvider' },
+  ['textDocument/didChange'] = { 'textDocumentSync' },
+  ['textDocument/didClose'] = { 'textDocumentSync', 'openClose' },
+  ['textDocument/didOpen'] = { 'textDocumentSync', 'openClose' },
+  ['textDocument/didSave'] = { 'textDocumentSync', 'save' },
+  ['textDocument/documentColor'] = { 'colorProvider' },
+  ['textDocument/documentHighlight'] = { 'documentHighlightProvider' },
+  ['textDocument/documentLink'] = { 'documentLinkProvider' },
+  ['textDocument/documentSymbol'] = { 'documentSymbolProvider' },
+  ['textDocument/foldingRange'] = { 'foldingRangeProvider' },
+  ['textDocument/formatting'] = { 'documentFormattingProvider' },
+  ['textDocument/hover'] = { 'hoverProvider' },
+  ['textDocument/implementation'] = { 'implementationProvider' },
+  ['textDocument/inlayHint'] = { 'inlayHintProvider' },
+  ['textDocument/inlineCompletion'] = { 'inlineCompletionProvider' },
+  ['textDocument/inlineValue'] = { 'inlineValueProvider' },
+  ['textDocument/linkedEditingRange'] = { 'linkedEditingRangeProvider' },
+  ['textDocument/moniker'] = { 'monikerProvider' },
+  ['textDocument/onTypeFormatting'] = { 'documentOnTypeFormattingProvider' },
+  ['textDocument/prepareCallHierarchy'] = { 'callHierarchyProvider' },
+  ['textDocument/prepareRename'] = { 'renameProvider', 'prepareProvider' },
+  ['textDocument/prepareTypeHierarchy'] = { 'typeHierarchyProvider' },
+  ['textDocument/rangeFormatting'] = { 'documentRangeFormattingProvider' },
+  ['textDocument/rangesFormatting'] = { 'documentRangeFormattingProvider', 'rangesSupport' },
+  ['textDocument/references'] = { 'referencesProvider' },
+  ['textDocument/rename'] = { 'renameProvider' },
+  ['textDocument/selectionRange'] = { 'selectionRangeProvider' },
+  ['textDocument/semanticTokens/full'] = { 'semanticTokensProvider' },
+  ['textDocument/semanticTokens/full/delta'] = { 'semanticTokensProvider', 'full', 'delta' },
+  ['textDocument/semanticTokens/range'] = { 'semanticTokensProvider', 'range' },
+  ['textDocument/signatureHelp'] = { 'signatureHelpProvider' },
+  ['textDocument/typeDefinition'] = { 'typeDefinitionProvider' },
+  ['textDocument/willSave'] = { 'textDocumentSync', 'willSave' },
+  ['textDocument/willSaveWaitUntil'] = { 'textDocumentSync', 'willSaveWaitUntil' },
+  ['workspaceSymbol/resolve'] = { 'workspaceSymbolProvider', 'resolveProvider' },
+  ['workspace/diagnostic'] = { 'diagnosticProvider', 'workspaceDiagnostics' },
+  ['workspace/didChangeWorkspaceFolders'] = { 'workspace', 'workspaceFolders', 'changeNotifications' },
+  ['workspace/didCreateFiles'] = { 'workspace', 'fileOperations', 'didCreate' },
+  ['workspace/didDeleteFiles'] = { 'workspace', 'fileOperations', 'didDelete' },
+  ['workspace/didRenameFiles'] = { 'workspace', 'fileOperations', 'didRename' },
+  ['workspace/executeCommand'] = { 'executeCommandProvider' },
+  ['workspace/symbol'] = { 'workspaceSymbolProvider' },
+  ['workspace/textDocumentContent'] = { 'workspace', 'textDocumentContent' },
+  ['workspace/willCreateFiles'] = { 'workspace', 'fileOperations', 'willCreate' },
+  ['workspace/willDeleteFiles'] = { 'workspace', 'fileOperations', 'willDelete' },
+  ['workspace/willRenameFiles'] = { 'workspace', 'fileOperations', 'willRename' },
+  ['workspace/workspaceFolders'] = { 'workspace', 'workspaceFolders' },
+}
+-- stylua: ignore end
+
 return protocol
diff --git a/scripts/gen_lsp.lua b/scripts/gen_lsp.lua
index 8afbae239a..3e419c7d59 100644
--- a/scripts/gen_lsp.lua
+++ b/scripts/gen_lsp.lua
@@ -6,7 +6,7 @@ Generates lua-ls annotations for lsp.
 USAGE:
 nvim -l scripts/gen_lsp.lua gen  # by default, this will overwrite runtime/lua/vim/lsp/_meta/protocol.lua
 nvim -l scripts/gen_lsp.lua gen --version 3.18 --out runtime/lua/vim/lsp/_meta/protocol.lua
-nvim -l scripts/gen_lsp.lua gen --version 3.18 --methods
+nvim -l scripts/gen_lsp.lua gen --version 3.18 --methods --capabilities
 ]]
 
 local DEFAULT_LSP_VERSION = '3.18'
@@ -57,13 +57,21 @@ local function to_luaname(s)
 end
 
 ---@param protocol vim._gen_lsp.Protocol
-local function gen_methods(protocol)
+---@param gen_methods boolean
+---@param gen_capabilities boolean
+local function write_to_protocol(protocol, gen_methods, gen_capabilities)
+  if not gen_methods and not gen_capabilities then
+    return
+  end
+
   local indent = (' '):rep(2)
 
   --- @class vim._gen_lsp.Request
   --- @field deprecated? string
   --- @field documentation? string
   --- @field messageDirection string
+  --- @field clientCapability? string
+  --- @field serverCapability? string
   --- @field method string
   --- @field params? any
   --- @field proposed? boolean
@@ -76,6 +84,8 @@ local function gen_methods(protocol)
   --- @field documentation? string
   --- @field errorData? any
   --- @field messageDirection string
+  --- @field clientCapability? string
+  --- @field serverCapability? string
   --- @field method string
   --- @field params? any[]
   --- @field partialResult? any
@@ -91,52 +101,85 @@ local function gen_methods(protocol)
     return to_luaname(a.method) < to_luaname(b.method)
   end)
 
-  local output = {
-    '-- Generated by gen_lsp.lua, keep at end of file.',
-    '--- @alias vim.lsp.protocol.Method.ClientToServer',
-  }
-  for _, item in ipairs(all) do
-    if item.method and item.messageDirection == 'clientToServer' then
-      output[#output + 1] = ("--- | '%s',"):format(item.method)
+  local output = { '-- Generated by gen_lsp.lua, keep at end of file.' }
+
+  if gen_methods then
+    output[#output + 1] = '--- @alias vim.lsp.protocol.Method.ClientToServer'
+
+    for _, item in ipairs(all) do
+      if item.method and item.messageDirection == 'clientToServer' then
+        output[#output + 1] = ("--- | '%s',"):format(item.method)
+      end
     end
-  end
 
-  vim.list_extend(output, {
-    '',
-    '--- @alias vim.lsp.protocol.Method.ServerToClient',
-  })
-  for _, item in ipairs(all) do
-    if item.method and item.messageDirection == 'serverToClient' then
-      output[#output + 1] = ("--- | '%s',"):format(item.method)
+    vim.list_extend(output, {
+      '',
+      '--- @alias vim.lsp.protocol.Method.ServerToClient',
+    })
+    for _, item in ipairs(all) do
+      if item.method and item.messageDirection == 'serverToClient' then
+        output[#output + 1] = ("--- | '%s',"):format(item.method)
+      end
     end
-  end
 
-  vim.list_extend(output, {
-    '',
-    '--- @alias vim.lsp.protocol.Method',
-    '--- | vim.lsp.protocol.Method.ClientToServer',
-    '--- | vim.lsp.protocol.Method.ServerToClient',
-    '',
-    '-- Generated by gen_lsp.lua, keep at end of file.',
-    '---',
-    '--- @enum vim.lsp.protocol.Methods',
-    '--- @see https://microsoft.github.io/language-server-protocol/specification/#metaModel',
-    '--- LSP method names.',
-    'protocol.Methods = {',
-  })
-
-  for _, item in ipairs(all) do
-    if item.method then
-      if item.documentation then
-        local document = vim.split(item.documentation, '\n?\n', { trimempty = true })
-        for _, docstring in ipairs(document) do
-          output[#output + 1] = indent .. '--- ' .. docstring
+    vim.list_extend(output, {
+      '',
+      '--- @alias vim.lsp.protocol.Method',
+      '--- | vim.lsp.protocol.Method.ClientToServer',
+      '--- | vim.lsp.protocol.Method.ServerToClient',
+      '',
+      '-- Generated by gen_lsp.lua, keep at end of file.',
+      '--- @enum vim.lsp.protocol.Methods',
+      '--- @see https://microsoft.github.io/language-server-protocol/specification/#metaModel',
+      '--- LSP method names.',
+      'protocol.Methods = {',
+    })
+
+    for _, item in ipairs(all) do
+      if item.method then
+        if item.documentation then
+          local document = vim.split(item.documentation, '\n?\n', { trimempty = true })
+          for _, docstring in ipairs(document) do
+            output[#output + 1] = indent .. '--- ' .. docstring
+          end
         end
+        output[#output + 1] = ("%s%s = '%s',"):format(indent, to_luaname(item.method), item.method)
+      end
+    end
+    output[#output + 1] = '}'
+  end
+
+  if gen_capabilities then
+    vim.list_extend(output, {
+      '',
+      '-- stylua: ignore start',
+      '-- Generated by gen_lsp.lua, keep at end of file.',
+      '--- Maps method names to the required server capability',
+      'protocol._request_name_to_capability = {',
+    })
+
+    for _, item in ipairs(all) do
+      if item.serverCapability then
+        output[#output + 1] = ("%s['%s'] = { %s },"):format(
+          indent,
+          item.method,
+          table.concat(
+            vim
+              .iter(vim.split(item.serverCapability, '.', { plain = true }))
+              :map(function(segment)
+                return "'" .. segment .. "'"
+              end)
+              :totable(),
+            ', '
+          )
+        )
       end
-      output[#output + 1] = ("%s%s = '%s',"):format(indent, to_luaname(item.method), item.method)
     end
+
+    output[#output + 1] = '}'
+    output[#output + 1] = '-- stylua: ignore end'
   end
-  output[#output + 1] = '}'
+
   output[#output + 1] = ''
   output[#output + 1] = 'return protocol'
 
@@ -157,15 +200,14 @@ end
 ---@field output_file string
 ---@field version string
 ---@field methods boolean
+---@field capabilities boolean
 
 ---@param opt vim._gen_lsp.opt
 function M.gen(opt)
   --- @type vim._gen_lsp.Protocol
   local protocol = read_json(opt)
 
-  if opt.methods then
-    gen_methods(protocol)
-  end
+  write_to_protocol(protocol, opt.methods, opt.capabilities)
 
   local output = {
     '--' .. '[[',
@@ -433,6 +475,7 @@ local opt = {
   output_file = 'runtime/lua/vim/lsp/_meta/protocol.lua',
   version = DEFAULT_LSP_VERSION,
   methods = false,
+  capabilities = false,
 }
 
 local command = nil
@@ -446,6 +489,8 @@ while i <= #_G.arg do
     i = i + 1
   elseif _G.arg[i] == '--methods' then
     opt.methods = true
+  elseif _G.arg[i] == '--capabilities' then
+    opt.capabilities = true
   elseif vim.startswith(_G.arg[i], '-') then
     error('Unrecognized args: ' .. _G.arg[i])
   else