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

1 files changed, 88 insertions, 43 deletions

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