refactor(lsp): stateful data abstraction, vim.lsp.Capability #34639

Problem:
Closes #31453

Solution:
Introduce `vim.lsp.Capability`, which may serve as the base class for
all LSP features that require caching data. it
- was created if there is at least one client that supports the specific method;
- was destroyed if all clients that support the method were detached.

- Apply the refactor for `folding_range.lua` and `semantic_tokens.lua`.
- Show active features in :checkhealth.

Future:
I found that these features that are expected to be refactored by
`vim.lsp.Capability` have one characteristic in common: they all send
LSP requests once the document is modified. The following code is
different, but they are all for this purpose.

- semantic tokens:
fb8dba413f/runtime/lua/vim/lsp/semantic_tokens.lua (L192-L198)
- inlay hints, folding ranges, document color
fb8dba413f/runtime/lua/vim/lsp/inlay_hint.lua (L250-L266)

I think I can sum up this characteristic as the need to keep certain
data synchronized with the latest version computed by the server.
I believe we can handle this at the `vim.lsp.Capability` level, and
I think it will be very useful.

Therefore, my next step is to implement LSP request sending and data
synchronization on `vim.lsp.Capability`, rather than limiting it to the
current create/destroy data approach.

runtime/doc/news.txt

@@ -195,6 +195,7 @@ LSP
 • The function form of `cmd` in a vim.lsp.Config or vim.lsp.ClientConfig
   receives the resolved config as the second arg: `cmd(dispatchers, config)`.
 • Support for annotated text edits.
+• `:checkhealth vim.lsp` is now available to check which buffers the active LSP features are attached to.
 
 LUA
 

runtime/lua/vim/lsp.lua

@@ -2,6 +2,7 @@ local api = vim.api
 local validate = vim.validate
 
 local lsp = vim._defer_require('vim.lsp', {
+  _capability = ..., --- @module 'vim.lsp._capability'
   _changetracking = ..., --- @module 'vim.lsp._changetracking'
   _folding_range = ..., --- @module 'vim.lsp._folding_range'
   _snippet_grammar = ..., --- @module 'vim.lsp._snippet_grammar'

runtime/lua/vim/lsp/_capability.lua

@@ -0,0 +1,77 @@
+local api = vim.api
+
+--- `vim.lsp.Capability` is expected to be created one-to-one with a buffer
+--- when there is at least one supported client attached to that buffer,
+--- and will be destroyed when all supporting clients are detached.
+---@class vim.lsp.Capability
+---
+--- Static field for retrieving the instance associated with a specific `bufnr`.
+---
+--- Index inthe form of `bufnr` -> `capability`
+---@field active table<integer, vim.lsp.Capability?>
+---
+--- The LSP feature it supports.
+---@field name string
+---
+--- Buffer number it associated with.
+---@field bufnr integer
+---
+--- The augroup owned by this instance, which will be cleared upon destruction.
+---@field augroup integer
+---
+--- Per-client state data, scoped to the lifetime of the attached client.
+---@field client_state table<integer, table>
+local M = {}
+M.__index = M
+
+---@generic T : vim.lsp.Capability
+---@param self T
+---@param bufnr integer
+---@return T
+function M:new(bufnr)
+  -- `self` in the `new()` function refers to the concrete type (i.e., the metatable).
+  -- `Class` may be a subtype of `Capability`, as it supports inheritance.
+  ---@type vim.lsp.Capability
+  local Class = self
+  assert(Class.name and Class.active, 'Do not instantiate the abstract class')
+
+  ---@type vim.lsp.Capability
+  self = setmetatable({}, Class)
+  self.bufnr = bufnr
+  self.augroup = api.nvim_create_augroup(
+    string.format('nvim.lsp.%s:%s', self.name:gsub('%s+', '_'):lower(), bufnr),
+    { clear = true }
+  )
+  self.client_state = {}
+
+  api.nvim_create_autocmd('LspDetach', {
+    group = self.augroup,
+    buffer = bufnr,
+    callback = function(args)
+      self:on_detach(args.data.client_id)
+      if next(self.client_state) == nil then
+        self:destroy()
+      end
+    end,
+  })
+
+  Class.active[bufnr] = self
+  return self
+end
+
+function M:destroy()
+  -- In case the function is called before all the clients detached.
+  for client_id, _ in pairs(self.client_state) do
+    self:on_detach(client_id)
+  end
+
+  api.nvim_del_augroup_by_id(self.augroup)
+  self.active[self.bufnr] = nil
+end
+
+---@param client_id integer
+function M:on_detach(client_id)
+  self.client_state[client_id] = nil
+end
+
+return M

runtime/lua/vim/lsp/_folding_range.lua

@@ -12,18 +12,20 @@ local supported_fold_kinds = {
 
 local M = {}
 
----@class (private) vim.lsp.folding_range.State
+local Capability = require('vim.lsp._capability')
+
+---@class (private) vim.lsp.folding_range.State : vim.lsp.Capability
 ---
 ---@field active table<integer, vim.lsp.folding_range.State?>
----@field bufnr integer
+---
----@field augroup integer
+--- `TextDocument` version this `state` corresponds to.
 ---@field version? integer
 ---
---- Never use this directly, `renew()` the cached foldinfo
+--- Never use this directly, `evaluate()` the cached foldinfo
 --- then use on demand via `row_*` fields.
 ---
 --- Index In the form of client_id -> ranges
----@field client_ranges table<integer, lsp.FoldingRange[]?>
+---@field client_state table<integer, lsp.FoldingRange[]?>
 ---
 --- Index in the form of row -> [foldlevel, mark]
 ---@field row_level table<integer, [integer, ">" | "<"?]?>
@@ -33,10 +35,12 @@ local M = {}
 ---
 --- Index in the form of start_row -> collapsed_text
 ---@field row_text table<integer, string?>
-local State = { active = {} }
+local State = { name = 'Folding Range', active = {} }
+State.__index = State
+setmetatable(State, Capability)
 
---- Renew the cached foldinfo in the buffer.
+--- Re-evaluate the cached foldinfo in the buffer.
-function State:renew()
+function State:evaluate()
   ---@type table<integer, [integer, ">" | "<"?]?>
   local row_level = {}
   ---@type table<integer, table<lsp.FoldingRangeKind, true?>?>>
@@ -44,7 +48,7 @@ function State:renew()
   ---@type table<integer, string?>
   local row_text = {}
 
-  for client_id, ranges in pairs(self.client_ranges) do
+  for client_id, ranges in pairs(self.client_state) do
     for _, range in ipairs(ranges) do
       local start_row = range.startLine
       local end_row = range.endLine
@@ -83,6 +87,9 @@ 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_folding_range_enabled 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
@@ -127,12 +134,12 @@ function State:multi_handler(results, ctx)
     if result.err then
       log.error(result.err)
     else
-      self.client_ranges[client_id] = result.result
+      self.client_state[client_id] = result.result
     end
   end
   self.version = ctx.version
 
-  self:renew()
+  self:evaluate()
   if api.nvim_get_mode().mode:match('^i') then
     -- `foldUpdate()` is guarded in insert mode.
     schedule_foldupdate(self.bufnr)
@@ -151,7 +158,11 @@ end
 --- Request `textDocument/foldingRange` from the server.
 --- `foldupdate()` is scheduled once after the request is completed.
 ---@param client? vim.lsp.Client The client whose server supports `foldingRange`.
-function State:request(client)
+function State:refresh(client)
+  if not vim.b._lsp_folding_range_enabled then
+    return
+  end
+
   ---@type lsp.FoldingRangeParams
   local params = { textDocument = util.make_text_document_params(self.bufnr) }
 
@@ -174,7 +185,6 @@ function State:request(client)
 end
 
 function State:reset()
-  self.client_ranges = {}
   self.row_level = {}
   self.row_kinds = {}
   self.row_text = {}
@@ -183,34 +193,17 @@ end
 --- Initialize `state` and event hooks, then request folding ranges.
 ---@param bufnr integer
 ---@return vim.lsp.folding_range.State
-function State.new(bufnr)
+function State:new(bufnr)
-  local self = setmetatable({}, { __index = State })
+  self = Capability.new(self, bufnr)
-  self.bufnr = bufnr
-  self.augroup = api.nvim_create_augroup('nvim.lsp.folding_range:' .. bufnr, { clear = true })
   self:reset()
 
-  State.active[bufnr] = self
-
   api.nvim_buf_attach(bufnr, false, {
-    -- `on_detach` also runs on buffer reload (`:e`).
-    -- Ensure `state` and hooks are cleared to avoid duplication or leftover states.
-    on_detach = function()
-      util._cancel_requests({
-        bufnr = bufnr,
-        method = ms.textDocument_foldingRange,
-        type = 'pending',
-      })
-      local state = State.active[bufnr]
-      if state then
-        state:destroy()
-      end
-    end,
     -- Reset `bufstate` and request folding ranges.
     on_reload = function()
       local state = State.active[bufnr]
       if state then
         state:reset()
-        state:request()
+        state:refresh()
       end
     end,
     --- Sync changed rows with their previous foldlevels before applying new ones.
@@ -238,44 +231,6 @@ function State.new(bufnr)
       end
     end,
   })
-  api.nvim_create_autocmd('LspDetach', {
-    group = self.augroup,
-    buffer = bufnr,
-    callback = function(args)
-      if not api.nvim_buf_is_loaded(bufnr) then
-        return
-      end
-
-      ---@type integer
-      local client_id = args.data.client_id
-      self.client_ranges[client_id] = nil
-
-      ---@type vim.lsp.Client[]
-      local clients = vim
-        .iter(vim.lsp.get_clients({ bufnr = bufnr, method = ms.textDocument_foldingRange }))
-        ---@param client vim.lsp.Client
-        :filter(function(client)
-          return client.id ~= client_id
-        end)
-        :totable()
-      if #clients == 0 then
-        self:reset()
-      end
-
-      self:renew()
-      foldupdate(bufnr)
-    end,
-  })
-  api.nvim_create_autocmd('LspAttach', {
-    group = self.augroup,
-    buffer = bufnr,
-    callback = function(args)
-      local client = assert(vim.lsp.get_client_by_id(args.data.client_id))
-      if client:supports_method(vim.lsp.protocol.Methods.textDocument_foldingRange, bufnr) then
-        self:request(client)
-      end
-    end,
-  })
   api.nvim_create_autocmd('LspNotify', {
     group = self.augroup,
     buffer = bufnr,
@@ -288,7 +243,16 @@ function State.new(bufnr)
           or args.data.method == ms.textDocument_didOpen
         )
       then
-        self:request(client)
+        self:refresh(client)
+      end
+    end,
+  })
+  api.nvim_create_autocmd('OptionSet', {
+    group = self.augroup,
+    pattern = 'foldexpr',
+    callback = function()
+      if vim.v.option_type == 'global' or vim.api.nvim_get_current_buf() == bufnr then
+        vim.b[bufnr]._lsp_folding_range_enabled = nil
       end
     end,
   })
@@ -301,18 +265,22 @@ function State:destroy()
   State.active[self.bufnr] = nil
 end
 
-local function setup(bufnr)
+---@params client_id integer
-  if not api.nvim_buf_is_loaded(bufnr) then
+function State:on_detach(client_id)
-    return
+  self.client_state[client_id] = nil
-  end
+  self:evaluate()
+  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)
+    state = State:new(bufnr)
   end
 
-  state:request()
+  state:refresh(client_id and vim.lsp.get_client_by_id(client_id))
-  return state
 end
 
 ---@param kind lsp.FoldingRangeKind
@@ -344,11 +312,11 @@ function M.foldclose(kind, winid)
     return
   end
 
+  -- Schedule `foldclose()` if the buffer is not up-to-date.
   if state.version == util.buf_versions[bufnr] then
     state:foldclose(kind, winid)
     return
   end
-  -- Schedule `foldclose()` if the buffer is not up-to-date.
 
   if not next(vim.lsp.get_clients({ bufnr = bufnr, method = ms.textDocument_foldingRange })) then
     return
@@ -380,14 +348,22 @@ end
 ---@return string level
 function M.foldexpr(lnum)
   local bufnr = api.nvim_get_current_buf()
-  local state = State.active[bufnr] or setup(bufnr)
+  local state = State.active[bufnr]
+  if not vim.b[bufnr]._lsp_folding_range_enabled then
+    vim.b[bufnr]._lsp_folding_range_enabled = true
+    if state then
+      state:refresh()
+    end
+  end
+
   if not state then
     return '0'
   end
-
   local row = (lnum or vim.v.lnum) - 1
   local level = state.row_level[row]
   return level and (level[2] or '') .. (level[1] or '0') or '0'
 end
 
+M.__FoldEvaluator = State
+
 return M

runtime/lua/vim/lsp/client.lua

@@ -1082,6 +1082,9 @@ function Client:on_attach(bufnr)
     if vim.tbl_get(self.server_capabilities, 'semanticTokensProvider', 'full') then
       lsp.semantic_tokens.start(bufnr, self.id)
     end
+    if vim.tbl_get(self.server_capabilities, 'foldingRangeProvider') then
+      lsp._folding_range._setup(bufnr)
+    end
   end)
 
   self.attached_buffers[bufnr] = true

runtime/lua/vim/lsp/health.lua

@@ -28,6 +28,43 @@ local function check_log()
   report_fn(string.format('Log size: %d KB', log_size / 1000))
 end
 
+local function check_active_features()
+  vim.health.start('vim.lsp: Active Features')
+  ---@type vim.lsp.Capability[]
+  local features = {
+    require('vim.lsp.semantic_tokens').__STHighlighter,
+    require('vim.lsp._folding_range').__FoldEvaluator,
+  }
+  for _, feature in ipairs(features) do
+    ---@type string[]
+    local buf_infos = {}
+    for bufnr, instance in pairs(feature.active) do
+      local client_info = vim
+        .iter(pairs(instance.client_state))
+        :map(function(client_id)
+          local client = vim.lsp.get_client_by_id(client_id)
+          if client then
+            return string.format('%s (id: %d)', client.name, client.id)
+          else
+            return string.format('unknow (id: %d)', client_id)
+          end
+        end)
+        :join(', ')
+      if client_info == '' then
+        client_info = 'No supported client attached'
+      end
+
+      buf_infos[#buf_infos + 1] = string.format('    [%d]: %s', bufnr, client_info)
+    end
+
+    report_info(table.concat({
+      feature.name,
+      '- Active buffers:',
+      string.format(table.concat(buf_infos, '\n')),
+    }, '\n'))
+  end
+end
+
 --- @param f function
 --- @return string
 local function func_tostring(f)
@@ -223,6 +260,7 @@ end
 --- Performs a healthcheck for LSP
 function M.check()
   check_log()
+  check_active_features()
   check_active_clients()
   check_enabled_configs()
   check_watcher()

runtime/lua/vim/lsp/semantic_tokens.lua

@@ -5,6 +5,8 @@ local util = require('vim.lsp.util')
 local Range = require('vim.treesitter._range')
 local uv = vim.uv
 
+local Capability = require('vim.lsp._capability')
+
 --- @class (private) STTokenRange
 --- @field line integer line number 0-based
 --- @field start_col integer start column 0-based
@@ -30,14 +32,16 @@ local uv = vim.uv
 --- @field active_request STActiveRequest
 --- @field current_result STCurrentResult
 
----@class (private) STHighlighter
+---@class (private) STHighlighter : vim.lsp.Capability
 ---@field active table<integer, STHighlighter>
 ---@field bufnr integer
 ---@field augroup integer augroup for buffer events
 ---@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 = { active = {} }
+local STHighlighter = { name = 'Semantic Tokens', active = {} }
+STHighlighter.__index = STHighlighter
+setmetatable(STHighlighter, Capability)
 
 --- Do a binary search of the tokens in the half-open range [lo, hi).
 ---
@@ -179,14 +183,8 @@ end
 ---@private
 ---@param bufnr integer
 ---@return STHighlighter
-function STHighlighter.new(bufnr)
+function STHighlighter:new(bufnr)
-  local self = setmetatable({}, { __index = STHighlighter })
+  self = Capability.new(self, bufnr)
-
-  self.bufnr = bufnr
-  self.augroup = api.nvim_create_augroup('nvim.lsp.semantic_tokens:' .. bufnr, { clear = true })
-  self.client_state = {}
-
-  STHighlighter.active[bufnr] = self
 
   api.nvim_buf_attach(bufnr, false, {
     on_lines = function(_, buf)
@@ -213,32 +211,11 @@ function STHighlighter.new(bufnr)
     end,
   })
 
-  api.nvim_create_autocmd('LspDetach', {
-    buffer = self.bufnr,
-    group = self.augroup,
-    callback = function(args)
-      self:detach(args.data.client_id)
-      if vim.tbl_isempty(self.client_state) then
-        self:destroy()
-      end
-    end,
-  })
-
   return self
 end
 
 ---@package
-function STHighlighter:destroy()
+function STHighlighter:on_attach(client_id)
-  for client_id, _ in pairs(self.client_state) do
-    self:detach(client_id)
-  end
-
-  api.nvim_del_augroup_by_id(self.augroup)
-  STHighlighter.active[self.bufnr] = nil
-end
-
----@package
-function STHighlighter:attach(client_id)
   local state = self.client_state[client_id]
   if not state then
     state = {
@@ -251,7 +228,7 @@ function STHighlighter:attach(client_id)
 end
 
 ---@package
-function STHighlighter:detach(client_id)
+function STHighlighter:on_detach(client_id)
   local state = self.client_state[client_id]
   if state then
     --TODO: delete namespace if/when that becomes possible
@@ -657,13 +634,13 @@ function M.start(bufnr, client_id, opts)
   local highlighter = STHighlighter.active[bufnr]
 
   if not highlighter then
-    highlighter = STHighlighter.new(bufnr)
+    highlighter = STHighlighter:new(bufnr)
     highlighter.debounce = opts.debounce or 200
   else
     highlighter.debounce = math.max(highlighter.debounce, opts.debounce or 200)
   end
 
-  highlighter:attach(client_id)
+  highlighter:on_attach(client_id)
   highlighter:send_request()
 end
 
@@ -687,7 +664,7 @@ function M.stop(bufnr, client_id)
     return
   end
 
-  highlighter:detach(client_id)
+  highlighter:on_detach(client_id)
 
   if vim.tbl_isempty(highlighter.client_state) then
     highlighter:destroy()

test/functional/plugin/lsp/folding_range_spec.lua

@@ -4,7 +4,6 @@ local Screen = require('test.functional.ui.screen')
 local t_lsp = require('test.functional.plugin.lsp.testutil')
 
 local eq = t.eq
-local tempname = t.tmpname
 
 local clear_notrace = t_lsp.clear_notrace
 local create_server_definition = t_lsp.create_server_definition
@@ -121,52 +120,6 @@ static int foldLevel(linenr_T lnum)
     api.nvim_exec_autocmds('VimLeavePre', { modeline = false })
   end)
 
-  describe('setup()', function()
-    ---@type integer
-    local bufnr_set_expr
-    ---@type integer
-    local bufnr_never_set_expr
-
-    local function buf_autocmd_num(bufnr_to_check)
-      return exec_lua(function()
-        return #vim.api.nvim_get_autocmds({ buffer = bufnr_to_check, event = 'LspNotify' })
-      end)
-    end
-
-    before_each(function()
-      command([[setlocal foldexpr=v:lua.vim.lsp.foldexpr()]])
-      exec_lua(function()
-        bufnr_set_expr = vim.api.nvim_create_buf(true, false)
-        vim.api.nvim_set_current_buf(bufnr_set_expr)
-      end)
-      insert(text)
-      command('write ' .. tempname(false))
-      command([[setlocal foldexpr=v:lua.vim.lsp.foldexpr()]])
-      exec_lua(function()
-        bufnr_never_set_expr = vim.api.nvim_create_buf(true, false)
-        vim.api.nvim_set_current_buf(bufnr_never_set_expr)
-      end)
-      insert(text)
-      api.nvim_win_set_buf(0, bufnr_set_expr)
-    end)
-
-    it('only create event hooks where foldexpr has been set', function()
-      eq(1, buf_autocmd_num(bufnr))
-      eq(1, buf_autocmd_num(bufnr_set_expr))
-      eq(0, buf_autocmd_num(bufnr_never_set_expr))
-    end)
-
-    it('does not create duplicate event hooks after reloaded', function()
-      command('edit')
-      eq(1, buf_autocmd_num(bufnr_set_expr))
-    end)
-
-    it('cleans up event hooks when buffer is unloaded', function()
-      command('bdelete')
-      eq(0, buf_autocmd_num(bufnr_set_expr))
-    end)
-  end)
-
   describe('expr()', function()
     --- @type test.functional.ui.screen
     local screen
@@ -182,6 +135,29 @@ static int foldLevel(linenr_T lnum)
       command([[split]])
     end)
 
+    it('controls the value of `b:_lsp_folding_range_enabled`', function()
+      eq(
+        true,
+        exec_lua(function()
+          return vim.b._lsp_folding_range_enabled
+        end)
+      )
+      command [[setlocal foldexpr=]]
+      eq(
+        nil,
+        exec_lua(function()
+          return vim.b._lsp_folding_range_enabled
+        end)
+      )
+      command([[set foldexpr=v:lua.vim.lsp.foldexpr()]])
+      eq(
+        true,
+        exec_lua(function()
+          return vim.b._lsp_folding_range_enabled
+        end)
+      )
+    end)
+
     it('can compute fold levels', function()
       ---@type table<integer, string>
       local foldlevels = {}