refactor(lsp): use vim.lsp._capability.enable internally

runtime/lua/vim/lsp/_capability.lua

@@ -19,6 +19,9 @@ local all_capabilities = {}
 --- Static field as the identifier of the LSP capability it supports.
 ---@field name vim.lsp.capability.Name
 ---
+--- Static field records the method this capability requires.
+---@field method vim.lsp.protocol.Method.ClientToServer
+---
 --- Static field for retrieving the instance associated with a specific `bufnr`.
 ---
 --- Index in the form of `bufnr` -> `capability`
@@ -123,6 +126,38 @@ function M.enable(name, enable, filter)
   local var = make_enable_var(name)
   local client = client_id and vim.lsp.get_client_by_id(client_id)
 
+  -- Attach or detach the client and its capability
+  -- based on the user’s latest marker value.
+  for _, it_client in ipairs(client and { client } or vim.lsp.get_clients()) do
+    for _, it_bufnr in
+      ipairs(
+        bufnr and { it_client.attached_buffers[bufnr] and bufnr }
+          or vim.lsp.get_buffers_by_client_id(it_client.id)
+      )
+    do
+      if enable ~= M.is_enabled(name, { bufnr = it_bufnr, client_id = it_client.id }) then
+        local Capability = all_capabilities[name]
+
+        if enable then
+          if it_client:supports_method(Capability.method) then
+            local capability = Capability.active[bufnr] or Capability:new(it_bufnr)
+            if not capability.client_state[it_client.id] then
+              capability:on_attach(it_client.id)
+            end
+          end
+        else
+          local capability = Capability.active[it_bufnr]
+          if capability then
+            capability:on_detach(it_client.id)
+            if not next(capability.client_state) then
+              capability:destroy()
+            end
+          end
+        end
+      end
+    end
+  end
+
   -- Updates the marker value.
   -- If local marker matches the global marker, set it to nil
   -- so that `is_enable` falls back to the global marker.

runtime/lua/vim/lsp/_folding_range.lua

@@ -35,9 +35,14 @@ local Capability = require('vim.lsp._capability')
 ---
 --- Index in the form of start_row -> collapsed_text
 ---@field row_text table<integer, string?>
-local State = { name = 'folding_range', active = {} }
+local State = {
+  name = 'folding_range',
+  method = ms.textDocument_foldingRange,
+  active = {},
+}
 State.__index = State
 setmetatable(State, Capability)
+Capability.all[State.name] = State
 
 --- Re-evaluate the cached foldinfo in the buffer.
 function State:evaluate()
@@ -87,9 +92,6 @@ end
 --- Force `foldexpr()` to be re-evaluated, without opening folds.
 ---@param bufnr integer
 local function foldupdate(bufnr)
-  if not api.nvim_buf_is_loaded(bufnr) or not vim.b[bufnr]._lsp_enabled_folding_range then
-    return
-  end
   for _, winid in ipairs(vim.fn.win_findbuf(bufnr)) do
     local wininfo = vim.fn.getwininfo(winid)[1]
     if wininfo and wininfo.tabnr == vim.fn.tabpagenr() then
@@ -159,10 +161,6 @@ end
 --- `foldupdate()` is scheduled once after the request is completed.
 ---@param client? vim.lsp.Client The client whose server supports `foldingRange`.
 function State:refresh(client)
-  if not vim.b._lsp_enabled_folding_range then
-    return
-  end
-
   ---@type lsp.FoldingRangeParams
   local params = { textDocument = util.make_text_document_params(self.bufnr) }
 
@@ -252,7 +250,7 @@ function State:new(bufnr)
     pattern = 'foldexpr',
     callback = function()
       if vim.v.option_type == 'global' or vim.api.nvim_get_current_buf() == bufnr then
-        vim.b[bufnr]._lsp_enabled_folding_range = nil
+        vim.lsp._capability.enable('folding_range', false, { bufnr = bufnr })
       end
     end,
   })
@@ -267,6 +265,7 @@ end
 
 ---@param client_id integer
 function State:on_attach(client_id)
+  self.client_state = {}
   self:refresh(vim.lsp.get_client_by_id(client_id))
 end
 
@@ -277,16 +276,6 @@ function State:on_detach(client_id)
   foldupdate(self.bufnr)
 end
 
----@param bufnr integer
----@param client_id integer
-function M._setup(bufnr, client_id)
-  local state = State.active[bufnr]
-  if not state then
-    state = State:new(bufnr)
-  end
-  state:on_attach(client_id)
-end
-
 ---@param kind lsp.FoldingRangeKind
 ---@param winid integer
 function State:foldclose(kind, winid)
@@ -352,14 +341,14 @@ end
 ---@return string level
 function M.foldexpr(lnum)
   local bufnr = api.nvim_get_current_buf()
-  local state = State.active[bufnr]
-  if not vim.b[bufnr]._lsp_enabled_folding_range then
-    vim.b[bufnr]._lsp_enabled_folding_range = true
-    if state then
-      state:refresh()
-    end
+  if not vim.lsp._capability.is_enabled('folding_range', { bufnr = bufnr }) then
+    -- `foldexpr` lead to a textlock, so any further operations need to be scheduled.
+    vim.schedule(function()
+      vim.lsp._capability.enable('folding_range', true, { bufnr = bufnr })
+    end)
   end
 
+  local state = State.active[bufnr]
   if not state then
     return '0'
   end

runtime/lua/vim/lsp/client.lua

@@ -511,6 +511,10 @@ function Client:initialize()
     root_path = vim.uri_to_fname(root_uri)
   end
 
+  -- HACK: Capability modules must be loaded
+  require('vim.lsp.semantic_tokens')
+  require('vim.lsp._folding_range')
+
   local init_params = {
     -- The process Id of the parent process that started the server. Is null if
     -- the process has not been started by another process.  If the parent
@@ -1086,16 +1090,21 @@ function Client:on_attach(bufnr)
   })
 
   self:_run_callbacks(self._on_attach_cbs, lsp.client_errors.ON_ATTACH_ERROR, self, bufnr)
-
-  -- schedule the initialization of semantic tokens to give the above
+  -- schedule the initialization of capabilities to give the above
   -- on_attach and LspAttach callbacks the ability to schedule wrap the
   -- opt-out (deleting the semanticTokensProvider from capabilities)
   vim.schedule(function()
-    if self:supports_method(ms.textDocument_semanticTokens_full) then
-      lsp.semantic_tokens._start(bufnr, self.id)
-    end
-    if self:supports_method(ms.textDocument_foldingRange) then
-      lsp._folding_range._setup(bufnr, self.id)
+    for _, Capability in pairs(vim.lsp._capability.all) do
+      if
+        self:supports_method(Capability.method)
+        and vim.lsp._capability.is_enabled(Capability.name, {
+          bufnr = bufnr,
+          client_id = self.id,
+        })
+      then
+        local capability = Capability.active[bufnr] or Capability:new(bufnr)
+        capability:on_attach(self.id)
+      end
     end
   end)
 

runtime/lua/vim/lsp/semantic_tokens.lua

@@ -41,9 +41,14 @@ local M = {}
 ---@field debounce integer milliseconds to debounce requests for new tokens
 ---@field timer table uv_timer for debouncing requests for new tokens
 ---@field client_state table<integer, STClientState>
-local STHighlighter = { name = 'semantic_tokens', active = {} }
+local STHighlighter = {
+  name = 'semantic_tokens',
+  method = ms.textDocument_semanticTokens_full,
+  active = {},
+}
 STHighlighter.__index = STHighlighter
 setmetatable(STHighlighter, Capability)
+Capability.all[STHighlighter.name] = STHighlighter
 
 --- Extracts modifier strings from the encoded number in the token array
 ---
@@ -156,6 +161,7 @@ end
 ---@param bufnr integer
 ---@return STHighlighter
 function STHighlighter:new(bufnr)
+  self.debounce = 200
   self = Capability.new(self, bufnr)
 
   api.nvim_buf_attach(bufnr, false, {
@@ -164,13 +170,11 @@ function STHighlighter:new(bufnr)
       if not highlighter then
         return true
       end
-      if M.is_enabled({ bufnr = buf }) then
-        highlighter:on_change()
-      end
+      highlighter:on_change()
     end,
     on_reload = function(_, buf)
       local highlighter = STHighlighter.active[buf]
-      if highlighter and M.is_enabled({ bufnr = bufnr }) then
+      if highlighter then
         highlighter:reset()
         highlighter:send_request()
       end
@@ -181,9 +185,7 @@ function STHighlighter:new(bufnr)
     buffer = self.bufnr,
     group = self.augroup,
     callback = function()
-      if M.is_enabled({ bufnr = bufnr }) then
-        self:send_request()
-      end
+      self:send_request()
     end,
   })
 
@@ -201,9 +203,7 @@ function STHighlighter:on_attach(client_id)
     }
     self.client_state[client_id] = state
   end
-  if M.is_enabled({ bufnr = self.bufnr }) then
-    self:send_request()
-  end
+  self:send_request()
 end
 
 ---@package
@@ -687,17 +687,6 @@ end
 ---@param filter? vim.lsp.capability.enable.Filter
 function M.enable(enable, filter)
   vim.lsp._capability.enable('semantic_tokens', enable, filter)
-
-  for _, bufnr in ipairs(api.nvim_list_bufs()) do
-    local highlighter = STHighlighter.active[bufnr]
-    if highlighter then
-      if M.is_enabled({ bufnr = bufnr }) then
-        highlighter:send_request()
-      else
-        highlighter:reset()
-      end
-    end
-  end
 end
 
 --- @nodoc
@@ -779,7 +768,7 @@ function M.force_refresh(bufnr)
 
   for _, buffer in ipairs(buffers) do
     local highlighter = STHighlighter.active[buffer]
-    if highlighter and M.is_enabled({ bufnr = bufnr }) then
+    if highlighter then
       highlighter:reset()
       highlighter:send_request()
     end

test/functional/plugin/lsp/folding_range_spec.lua

@@ -135,25 +135,25 @@ static int foldLevel(linenr_T lnum)
       command([[split]])
     end)
 
-    it('controls the value of `b:_lsp_enabled_folding_range`', function()
+    it('controls whether folding range is enabled', function()
       eq(
         true,
         exec_lua(function()
-          return vim.b._lsp_enabled_folding_range
+          return vim.lsp._capability.is_enabled('folding_range', { bufnr = 0 })
         end)
       )
       command [[setlocal foldexpr=]]
       eq(
-        nil,
+        false,
         exec_lua(function()
-          return vim.b._lsp_enabled_folding_range
+          return vim.lsp._capability.is_enabled('folding_range', { bufnr = 0 })
         end)
       )
       command([[set foldexpr=v:lua.vim.lsp.foldexpr()]])
       eq(
         true,
         exec_lua(function()
-          return vim.b._lsp_enabled_folding_range
+          return vim.lsp._capability.is_enabled('folding_range', { bufnr = 0 })
         end)
       )
     end)