refactor(lsp): move `util.enable` to `capability.enable`

1 files changed, 85 insertions, 0 deletions

diff --git a/runtime/lua/vim/lsp/_capability.lua b/runtime/lua/vim/lsp/_capability.lua
index 5b38f93e30..5acd9bf067 100644
--- a/runtime/lua/vim/lsp/_capability.lua
+++ b/runtime/lua/vim/lsp/_capability.lua
@@ -3,6 +3,7 @@ local api = vim.api
 ---@alias vim.lsp.capability.Name
 ---| 'semantic_tokens'
 ---| 'folding_range'
+---| 'linked_editing_range'
 
 --- Tracks all supported capabilities, all of which derive from `vim.lsp.Capability`.
 --- Returns capability *prototypes*, not their instances.
@@ -99,6 +100,90 @@ function M:on_detach(client_id)
   self.client_state[client_id] = nil
 end
 
+---@param name vim.lsp.capability.Name
+local function make_enable_var(name)
+  return ('_lsp_enabled_%s'):format(name)
+end
+
+--- Optional filters |kwargs|,
+---@class vim.lsp.capability.enable.Filter
+---@inlinedoc
+---
+--- Buffer number, or 0 for current buffer, or nil for all.
+--- (default: all)
+---@field bufnr? integer
+---
+--- Client ID, or nil for all.
+--- (default: all)
+---@field client_id? integer
+
+---@param name vim.lsp.capability.Name
+---@param enable? boolean
+---@param filter? vim.lsp.capability.enable.Filter
+function M.enable(name, enable, filter)
+  vim.validate('name', name, '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 and vim._resolve_bufnr(filter.bufnr)
+  local client_id = filter.client_id
+  assert(not (bufnr and client_id), '`bufnr` and `client_id` are mutually exclusive.')
+
+  local var = make_enable_var(name)
+  local client = client_id and vim.lsp.get_client_by_id(client_id)
+
+  -- 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.
+  if client then
+    if enable == vim.g[var] then
+      client._enabled_capabilities[name] = nil
+    else
+      client._enabled_capabilities[name] = 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
+      if it_client._enabled_capabilities[name] == enable then
+        it_client._enabled_capabilities[name] = nil
+      end
+    end
+  end
+end
+
+---@param name vim.lsp.capability.Name
+---@param filter? vim.lsp.capability.enable.Filter
+function M.is_enabled(name, filter)
+  vim.validate('name', name, 'string')
+  vim.validate('filter', filter, 'table', true)
+
+  filter = filter or {}
+  local bufnr = filter.bufnr and vim._resolve_bufnr(filter.bufnr)
+  local client_id = filter.client_id
+
+  local var = make_enable_var(name)
+  local client = client_id and vim.lsp.get_client_by_id(client_id)
+
+  -- As a fallback when not explicitly enabled or disabled:
+  -- Clients are treated as "enabled" since their capabilities can control behavior.
+  -- Buffers are treated as "disabled" to allow users to enable them as needed.
+  return vim.F.if_nil(client and client._enabled_capabilities[name], vim.g[var], true)
+    and vim.F.if_nil(bufnr and vim.b[bufnr][var], vim.g[var], false)
+end
+
 M.all = all_capabilities
 
 return M