refactor(lsp): redesign LSP folding state #34469

runtime/lua/vim/lsp/_folding_range.lua

@@ -12,8 +12,11 @@ local supported_fold_kinds = {
 
 local M = {}
 
----@class (private) vim.lsp.folding_range.BufState
+---@class (private) vim.lsp.folding_range.State
 ---
+---@field active table<integer, vim.lsp.folding_range.State?>
+---@field bufnr integer
+---@field augroup integer
 ---@field version? integer
 ---
 --- Never use this directly, `renew()` the cached foldinfo
@@ -30,15 +33,10 @@ local M = {}
 ---
 --- Index in the form of start_row -> collapsed_text
 ---@field row_text table<integer, string?>
-
+local State = { active = {} }
----@type table<integer, vim.lsp.folding_range.BufState?>
-local bufstates = {}
 
 --- Renew the cached foldinfo in the buffer.
----@param bufnr integer
+function State:renew()
-local function renew(bufnr)
-  local bufstate = assert(bufstates[bufnr])
-
   ---@type table<integer, [integer, ">" | "<"?]?>
   local row_level = {}
   ---@type table<integer, table<lsp.FoldingRangeKind, true?>?>>
@@ -46,7 +44,7 @@ local function renew(bufnr)
   ---@type table<integer, string?>
   local row_text = {}
 
-  for client_id, ranges in pairs(bufstate.client_ranges) do
+  for client_id, ranges in pairs(self.client_ranges) do
     for _, range in ipairs(ranges) do
       local start_row = range.startLine
       local end_row = range.endLine
@@ -77,16 +75,14 @@ local function renew(bufnr)
     end
   end
 
-  bufstate.row_level = row_level
+  self.row_level = row_level
-  bufstate.row_kinds = row_kinds
+  self.row_kinds = row_kinds
-  bufstate.row_text = row_text
+  self.row_text = row_text
 end
 
---- Renew the cached foldinfo then force `foldexpr()` to be re-evaluated,
+--- Force `foldexpr()` to be re-evaluated, without opening folds.
---- without opening folds.
 ---@param bufnr integer
 local function foldupdate(bufnr)
-  renew(bufnr)
   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
@@ -120,119 +116,110 @@ local function schedule_foldupdate(bufnr)
 end
 
 ---@param results table<integer,{err: lsp.ResponseError?, result: lsp.FoldingRange[]?}>
----@type lsp.MultiHandler
+---@param ctx lsp.HandlerContext
-local function multi_handler(results, ctx)
+function State:multi_handler(results, ctx)
-  local bufnr = assert(ctx.bufnr)
   -- Handling responses from outdated buffer only causes performance overhead.
-  if util.buf_versions[bufnr] ~= ctx.version then
+  if util.buf_versions[self.bufnr] ~= ctx.version then
     return
   end
 
-  local bufstate = assert(bufstates[bufnr])
   for client_id, result in pairs(results) do
     if result.err then
       log.error(result.err)
     else
-      bufstate.client_ranges[client_id] = result.result
+      self.client_ranges[client_id] = result.result
     end
   end
-  bufstate.version = ctx.version
+  self.version = ctx.version
 
+  self:renew()
   if api.nvim_get_mode().mode:match('^i') then
     -- `foldUpdate()` is guarded in insert mode.
-    schedule_foldupdate(bufnr)
+    schedule_foldupdate(self.bufnr)
   else
-    foldupdate(bufnr)
+    foldupdate(self.bufnr)
   end
 end
 
+---@param err lsp.ResponseError?
 ---@param result lsp.FoldingRange[]?
----@type lsp.Handler
+---@param ctx lsp.HandlerContext, config?: table
-local function handler(err, result, ctx)
+function State:handler(err, result, ctx)
-  multi_handler({ [ctx.client_id] = { err = err, result = result } }, ctx)
+  self:multi_handler({ [ctx.client_id] = { err = err, result = result } }, ctx)
 end
 
 --- Request `textDocument/foldingRange` from the server.
 --- `foldupdate()` is scheduled once after the request is completed.
----@param bufnr integer
 ---@param client? vim.lsp.Client The client whose server supports `foldingRange`.
-local function request(bufnr, client)
+function State:request(client)
   ---@type lsp.FoldingRangeParams
-  local params = { textDocument = util.make_text_document_params(bufnr) }
+  local params = { textDocument = util.make_text_document_params(self.bufnr) }
 
   if client then
-    client:request(ms.textDocument_foldingRange, params, handler, bufnr)
+    client:request(ms.textDocument_foldingRange, params, function(...)
+      self:handler(...)
+    end, self.bufnr)
     return
   end
 
-  if not next(vim.lsp.get_clients({ bufnr = bufnr, method = ms.textDocument_foldingRange })) then
+  if
+    not next(vim.lsp.get_clients({ bufnr = self.bufnr, method = ms.textDocument_foldingRange }))
+  then
     return
   end
 
-  vim.lsp.buf_request_all(bufnr, ms.textDocument_foldingRange, params, multi_handler)
+  vim.lsp.buf_request_all(self.bufnr, ms.textDocument_foldingRange, params, function(...)
+    self:multi_handler(...)
+  end)
 end
 
--- NOTE:
+function State:reset()
--- `bufstate` and event hooks are interdependent:
+  self.client_ranges = {}
--- * `bufstate` needs event hooks for correctness.
+  self.row_level = {}
--- * event hooks require the previous `bufstate` for updates.
+  self.row_kinds = {}
--- Since they are manually created and destroyed,
+  self.row_text = {}
--- we ensure their lifecycles are always synchronized.
+end
---
--- TODO(ofseed):
--- 1. Implement clearing `bufstate` and event hooks
---    when no clients in the buffer support the corresponding method.
--- 2. Then generalize this state management to other LSP modules.
-local augroup_setup = api.nvim_create_augroup('nvim.lsp.folding_range.setup', {})
 
---- Initialize `bufstate` and event hooks, then request folding ranges.
+--- Initialize `state` and event hooks, then request folding ranges.
---- Manage their lifecycle within this function.
 ---@param bufnr integer
----@return vim.lsp.folding_range.BufState?
+---@return vim.lsp.folding_range.State
-local function setup(bufnr)
+function State.new(bufnr)
-  if not api.nvim_buf_is_loaded(bufnr) then
+  local self = setmetatable({}, { __index = State })
-    return
+  self.bufnr = bufnr
-  end
+  self.augroup = api.nvim_create_augroup('nvim.lsp.folding_range:' .. bufnr, { clear = true })
+  self:reset()
 
-  -- Register the new `bufstate`.
+  State.active[bufnr] = self
-  bufstates[bufnr] = {
-    client_ranges = {},
-    row_level = {},
-    row_kinds = {},
-    row_text = {},
-  }
 
-  -- Event hooks from `buf_attach` can't be removed externally.
-  -- Hooks and `bufstate` share the same lifecycle;
-  -- they should self-destroy if `bufstate == nil`.
   api.nvim_buf_attach(bufnr, false, {
     -- `on_detach` also runs on buffer reload (`:e`).
-    -- Ensure `bufstate` and hooks are cleared to avoid duplication or leftover states.
+    -- 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',
       })
-      bufstates[bufnr] = nil
+      local state = State.active[bufnr]
-      api.nvim_clear_autocmds({ buffer = bufnr, group = augroup_setup })
+      if state then
+        state:destroy()
+      end
     end,
     -- Reset `bufstate` and request folding ranges.
     on_reload = function()
-      bufstates[bufnr] = {
+      local state = State.active[bufnr]
-        client_ranges = {},
+      if state then
-        row_level = {},
+        state:reset()
-        row_kinds = {},
+        state:request()
-        row_text = {},
+      end
-      }
-      request(bufnr)
     end,
     --- Sync changed rows with their previous foldlevels before applying new ones.
     on_bytes = function(_, _, _, start_row, _, _, old_row, _, _, new_row, _, _)
-      if bufstates[bufnr] == nil then
+      local state = State.active[bufnr]
+      if state == nil then
         return true
       end
-      local row_level = bufstates[bufnr].row_level
+      local row_level = state.row_level
       if next(row_level) == nil then
         return
       end
@@ -252,7 +239,7 @@ local function setup(bufnr)
     end,
   })
   api.nvim_create_autocmd('LspDetach', {
-    group = augroup_setup,
+    group = self.augroup,
     buffer = bufnr,
     callback = function(args)
       if not api.nvim_buf_is_loaded(bufnr) then
@@ -261,7 +248,7 @@ local function setup(bufnr)
 
       ---@type integer
       local client_id = args.data.client_id
-      bufstates[bufnr].client_ranges[client_id] = nil
+      self.client_ranges[client_id] = nil
 
       ---@type vim.lsp.Client[]
       local clients = vim
@@ -272,29 +259,25 @@ local function setup(bufnr)
         end)
         :totable()
       if #clients == 0 then
-        bufstates[bufnr] = {
+        self:reset()
-          client_ranges = {},
-          row_level = {},
-          row_kinds = {},
-          row_text = {},
-        }
       end
 
+      self:renew()
       foldupdate(bufnr)
     end,
   })
   api.nvim_create_autocmd('LspAttach', {
-    group = augroup_setup,
+    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
-        request(bufnr, client)
+        self:request(client)
       end
     end,
   })
   api.nvim_create_autocmd('LspNotify', {
-    group = augroup_setup,
+    group = self.augroup,
     buffer = bufnr,
     callback = function(args)
       local client = assert(vim.lsp.get_client_by_id(args.data.client_id))
@@ -305,22 +288,39 @@ local function setup(bufnr)
           or args.data.method == ms.textDocument_didOpen
         )
       then
-        request(bufnr, client)
+        self:request(client)
       end
     end,
   })
 
-  request(bufnr)
+  return self
+end
 
-  return bufstates[bufnr]
+function State:destroy()
+  api.nvim_del_augroup_by_id(self.augroup)
+  State.active[self.bufnr] = nil
+end
+
+local function setup(bufnr)
+  if not api.nvim_buf_is_loaded(bufnr) then
+    return
+  end
+
+  local state = State.active[bufnr]
+  if not state then
+    state = State.new(bufnr)
+  end
+
+  state:request()
+  return state
 end
 
 ---@param kind lsp.FoldingRangeKind
 ---@param winid integer
-local function foldclose(kind, winid)
+function State:foldclose(kind, winid)
   vim._with({ win = winid }, function()
     local bufnr = api.nvim_win_get_buf(winid)
-    local row_kinds = bufstates[bufnr].row_kinds
+    local row_kinds = State.active[bufnr].row_kinds
     -- Reverse traverse to ensure that the smallest ranges are closed first.
     for row = api.nvim_buf_line_count(bufnr) - 1, 0, -1 do
       local kinds = row_kinds[row]
@@ -339,13 +339,13 @@ function M.foldclose(kind, winid)
 
   winid = winid or api.nvim_get_current_win()
   local bufnr = api.nvim_win_get_buf(winid)
-  local bufstate = bufstates[bufnr]
+  local state = State.active[bufnr]
-  if not bufstate then
+  if not state then
     return
   end
 
-  if bufstate.version == util.buf_versions[bufnr] then
+  if state.version == util.buf_versions[bufnr] then
-    foldclose(kind, winid)
+    state:foldclose(kind, winid)
     return
   end
   -- Schedule `foldclose()` if the buffer is not up-to-date.
@@ -356,10 +356,10 @@ function M.foldclose(kind, winid)
   ---@type lsp.FoldingRangeParams
   local params = { textDocument = util.make_text_document_params(bufnr) }
   vim.lsp.buf_request_all(bufnr, ms.textDocument_foldingRange, params, function(...)
-    multi_handler(...)
+    state:multi_handler(...)
     -- Ensure this buffer stays as the current buffer after the async request
     if api.nvim_win_get_buf(winid) == bufnr then
-      foldclose(kind, winid)
+      state:foldclose(kind, winid)
     end
   end)
 end
@@ -369,9 +369,9 @@ function M.foldtext()
   local bufnr = api.nvim_get_current_buf()
   local lnum = vim.v.foldstart
   local row = lnum - 1
-  local bufstate = bufstates[bufnr]
+  local state = State.active[bufnr]
-  if bufstate and bufstate.row_text[row] then
+  if state and state.row_text[row] then
-    return bufstate.row_text[row]
+    return state.row_text[row]
   end
   return vim.fn.getline(lnum)
 end
@@ -380,13 +380,13 @@ end
 ---@return string level
 function M.foldexpr(lnum)
   local bufnr = api.nvim_get_current_buf()
-  local bufstate = bufstates[bufnr] or setup(bufnr)
+  local state = State.active[bufnr] or setup(bufnr)
-  if not bufstate then
+  if not state then
     return '0'
   end
 
   local row = (lnum or vim.v.lnum) - 1
-  local level = bufstate.row_level[row]
+  local level = state.row_level[row]
   return level and (level[2] or '') .. (level[1] or '0') or '0'
 end