refactor(lsp): utility functions for enable()/is_enabled()

runtime/lua/vim/lsp/_folding_range.lua

@@ -87,7 +87,7 @@ 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
+  if not api.nvim_buf_is_loaded(bufnr) or not vim.b[bufnr]._lsp_enable_folding_range then
     return
   end
   for _, winid in ipairs(vim.fn.win_findbuf(bufnr)) do
@@ -159,7 +159,7 @@ 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_folding_range_enabled then
+  if not vim.b._lsp_enable_folding_range then
     return
   end
 
@@ -252,7 +252,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_folding_range_enabled = nil
+        vim.b[bufnr]._lsp_enable_folding_range = nil
       end
     end,
   })
@@ -349,8 +349,8 @@ end
 function M.foldexpr(lnum)
   local bufnr = api.nvim_get_current_buf()
   local state = State.active[bufnr]
-  if not vim.b[bufnr]._lsp_folding_range_enabled then
+  if not vim.b[bufnr]._lsp_enable_folding_range then
-    vim.b[bufnr]._lsp_folding_range_enabled = true
+    vim.b[bufnr]._lsp_enable_folding_range = true
     if state then
       state:refresh()
     end

runtime/lua/vim/lsp/util.lua

@@ -2340,6 +2340,85 @@ function M._refresh(method, opts)
   end
 end
 
+---@param feature string
+---@param client_id? integer
+local function make_enable_var(feature, client_id)
+  return ('_lsp_enabled_%s%s'):format(feature, client_id and ('_client_%d'):format(client_id) or '')
+end
+
+---@class vim.lsp.enable.Filter
+---@inlinedoc
+---
+--- Buffer number, or 0 for current buffer, or nil for all.
+---@field bufnr? integer
+---
+--- Client ID, or nil for all
+---@field client_id? integer
+
+---@param feature string
+---@param filter? vim.lsp.enable.Filter
+function M._is_enabled(feature, filter)
+  vim.validate('feature', feature, 'string')
+  vim.validate('filter', filter, 'table', true)
+
+  filter = filter or {}
+  local bufnr = filter.bufnr
+  local client_id = filter.client_id
+
+  local var = make_enable_var(feature)
+  local client_var = make_enable_var(feature, client_id)
+  return vim.F.if_nil(client_id and vim.g[client_var], vim.g[var])
+    and vim.F.if_nil(bufnr and vim.b[bufnr][var], vim.g[var])
+end
+
+---@param feature 'semantic_tokens'
+---@param enable? boolean
+---@param filter? vim.lsp.enable.Filter
+function M._enable(feature, enable, filter)
+  vim.validate('feature', feature, 'string')
+  vim.validate('enable', enable, 'boolean', true)
+  vim.validate('filter', filter, 'table', true)
+
+  enable = enable == nil or enable
+  filter = filter or {}
+  local bufnr = filter.bufnr
+  local client_id = filter.client_id
+  assert(
+    not (bufnr and client_id),
+    'Only one of `bufnr` or `client_id` filters can be specified at a time.'
+  )
+
+  local var = make_enable_var(feature)
+  local client_var = make_enable_var(feature, client_id)
+
+  if client_id then
+    if enable == vim.g[var] then
+      vim.g[client_var] = nil
+    else
+      vim.g[client_var] = enable
+    end
+  elseif bufnr then
+    if enable == vim.g[var] then
+      vim.b[bufnr][var] = nil
+    else
+      vim.b[bufnr][var] = enable
+    end
+  else
+    vim.g[var] = enable
+    for _, it_bufnr in ipairs(api.nvim_list_bufs()) do
+      if api.nvim_buf_is_loaded(it_bufnr) and vim.b[it_bufnr][var] == enable then
+        vim.b[it_bufnr][var] = nil
+      end
+    end
+    for _, it_client in ipairs(vim.lsp.get_clients()) do
+      local it_client_var = make_enable_var(feature, it_client.id)
+      if vim.g[it_client_var] and vim.g[it_client_var] == enable then
+        vim.g[it_client_var] = nil
+      end
+    end
+  end
+end
+
 M._get_line_byte_from_position = get_line_byte_from_position
 
 ---@nodoc

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_folding_range_enabled`', function()
+    it('controls the value of `b:_lsp_enable_folding_range`', function()
       eq(
         true,
         exec_lua(function()
-          return vim.b._lsp_folding_range_enabled
+          return vim.b._lsp_enable_folding_range
         end)
       )
       command [[setlocal foldexpr=]]
       eq(
         nil,
         exec_lua(function()
-          return vim.b._lsp_folding_range_enabled
+          return vim.b._lsp_enable_folding_range
         end)
       )
       command([[set foldexpr=v:lua.vim.lsp.foldexpr()]])
       eq(
         true,
         exec_lua(function()
-          return vim.b._lsp_folding_range_enabled
+          return vim.b._lsp_enable_folding_range
         end)
       )
     end)