feat(lsp): support textDocument/documentColor

test(lsp): add tests form `vim.lsp.document_color`

runtime/lua/vim/lsp.lua

@@ -12,6 +12,7 @@ local lsp = vim._defer_require('vim.lsp', {
   codelens = ..., --- @module 'vim.lsp.codelens'
   completion = ..., --- @module 'vim.lsp.completion'
   diagnostic = ..., --- @module 'vim.lsp.diagnostic'
+  document_color = ..., --- @module 'vim.lsp.document_color'
   handlers = ..., --- @module 'vim.lsp.handlers'
   inlay_hint = ..., --- @module 'vim.lsp.inlay_hint'
   log = ..., --- @module 'vim.lsp.log'

runtime/lua/vim/lsp/document_color.lua

@@ -0,0 +1,350 @@
+local api = vim.api
+local lsp = vim.lsp
+local util = lsp.util
+local ms = lsp.protocol.Methods
+
+local document_color_ns = api.nvim_create_namespace('nvim.lsp.document_color')
+local document_color_augroup = api.nvim_create_augroup('nvim.lsp.document_color', {})
+
+local M = {}
+
+--- @class (private) vim.lsp.document_color.HighlightInfo
+--- @field hex_code string Resolved HEX color
+--- @field range Range4 Range of the highlight
+--- @field hl_group? string Highlight group name. Won't be present if the style is a custom function.
+
+--- @class (private) vim.lsp.document_color.BufState
+--- @field enabled boolean Whether document_color is enabled for the current buffer
+--- @field buf_version? integer Buffer version for which the color ranges correspond to
+--- @field applied_version table<integer, integer?> (client_id -> buffer version) Last buffer version for which we applied color ranges
+--- @field hl_info table<integer, vim.lsp.document_color.HighlightInfo[]?> (client_id -> color highlights) Processed highlight information
+
+--- @type table<integer, vim.lsp.document_color.BufState?>
+local bufstates = {}
+
+--- @inlinedoc
+--- @class vim.lsp.document_color.enable.Opts
+---
+--- Highlight style. It can be one of the pre-defined styles, a string to be used as virtual text, or a
+--- function that receives the buffer handle, the range (start line, start col, end line, end col) and
+--- the resolved hex color. (default: `'background'`)
+--- @field style? 'background'|'foreground'|'virtual'|string|fun(bufnr: integer, range: Range4, hex_code: string)
+
+-- Default options.
+--- @type vim.lsp.document_color.enable.Opts
+local document_color_opts = { style = 'background' }
+
+--- @param color string
+local function get_contrast_color(color)
+  local r_s, g_s, b_s = color:match('^#(%x%x)(%x%x)(%x%x)$')
+  local r, g, b = tonumber(r_s, 16), tonumber(g_s, 16), tonumber(b_s, 16)
+
+  -- Source: https://www.w3.org/TR/WCAG21/#dfn-relative-luminance
+  -- Using power 2.2 is a close approximation to full piecewise transform
+  local R, G, B = (r / 255) ^ 2.2, (g / 255) ^ 2.2, (b / 255) ^ 2.2
+  local is_bright = (0.2126 * R + 0.7152 * G + 0.0722 * B) > 0.5
+  return is_bright and '#000000' or '#ffffff'
+end
+
+--- Returns the hex string representing the given LSP color.
+--- @param color lsp.Color
+--- @return string
+local function get_hex_code(color)
+  -- The RGB values in lsp.Color are in the [0-1] range, but we want them to be in the [0-255] range instead.
+  --- @param n number
+  color = vim.tbl_map(function(n)
+    return math.floor((n * 255) + 0.5)
+  end, color)
+
+  return ('#%02x%02x%02x'):format(color.red, color.green, color.blue):lower()
+end
+
+--- Cache of the highlight groups that we've already created.
+--- @type table<string, true>
+local color_cache = {}
+
+--- Gets or creates the highlight group for the given LSP color information.
+---
+--- @param hex_code string
+--- @param style string
+--- @return string
+local function get_hl_group(hex_code, style)
+  if style ~= 'background' then
+    style = 'foreground'
+  end
+
+  local hl_name = ('LspDocumentColor_%s_%s'):format(hex_code:sub(2), style)
+
+  if not color_cache[hl_name] then
+    if style == 'background' then
+      api.nvim_set_hl(0, hl_name, { bg = hex_code, fg = get_contrast_color(hex_code) })
+    else
+      api.nvim_set_hl(0, hl_name, { fg = hex_code })
+    end
+
+    color_cache[hl_name] = true
+  end
+
+  return hl_name
+end
+
+--- @param bufnr integer
+--- @param enabled boolean
+local function reset_bufstate(bufnr, enabled)
+  bufstates[bufnr] = { enabled = enabled, applied_version = {}, hl_info = {} }
+end
+
+--- |lsp-handler| for the `textDocument/documentColor` method.
+---
+--- @param err? lsp.ResponseError
+--- @param result? lsp.ColorInformation[]
+--- @param ctx lsp.HandlerContext
+local function on_document_color(err, result, ctx)
+  if err then
+    lsp.log.error('document_color', err)
+    return
+  end
+
+  local bufnr = assert(ctx.bufnr)
+  local bufstate = assert(bufstates[bufnr])
+  local client_id = ctx.client_id
+
+  if
+    util.buf_versions[bufnr] ~= ctx.version
+    or not result
+    or not api.nvim_buf_is_loaded(bufnr)
+    or not bufstate.enabled
+  then
+    return
+  end
+
+  local hl_infos = {} --- @type vim.lsp.document_color.HighlightInfo[]
+  local style = document_color_opts.style
+  for _, res in ipairs(result) do
+    local range = {
+      res.range.start.line,
+      res.range.start.character,
+      res.range['end'].line,
+      res.range['end'].character,
+    }
+    local hex_code = get_hex_code(res.color)
+    local hl_info = { range = range, hex_code = hex_code }
+
+    if type(style) == 'string' then
+      hl_info.hl_group = get_hl_group(hex_code, style)
+    end
+
+    table.insert(hl_infos, hl_info)
+  end
+  bufstate.hl_info[client_id] = hl_infos
+
+  bufstate.buf_version = ctx.version
+  api.nvim__redraw({ buf = bufnr, valid = true, flush = false })
+end
+
+--- @param bufnr integer
+local function buf_clear(bufnr)
+  local bufstate = assert(bufstates[bufnr])
+  local client_ids = vim.tbl_keys(bufstate.hl_info) --- @type integer[]
+
+  for _, id in ipairs(client_ids) do
+    bufstate.hl_info[id] = {}
+  end
+
+  api.nvim_buf_clear_namespace(bufnr, document_color_ns, 0, -1)
+  api.nvim__redraw({ buf = bufnr, valid = true, flush = false })
+end
+
+--- @param bufnr integer
+--- @param client_id? integer
+local function buf_refresh(bufnr, client_id)
+  util._refresh(ms.textDocument_documentColor, {
+    bufnr = bufnr,
+    handler = on_document_color,
+    client_id = client_id,
+  })
+end
+
+--- @param bufnr integer
+local function buf_disable(bufnr)
+  buf_clear(bufnr)
+  reset_bufstate(bufnr, false)
+end
+
+--- @param bufnr integer
+local function buf_enable(bufnr)
+  reset_bufstate(bufnr, true)
+
+  api.nvim_create_autocmd('LspNotify', {
+    buffer = bufnr,
+    group = document_color_augroup,
+    desc = 'Refresh document_color on document changes',
+    callback = function(args)
+      local method = args.data.method --- @type string
+
+      if
+        (method == ms.textDocument_didChange or method == ms.textDocument_didOpen)
+        and bufstates[args.buf].enabled
+      then
+        buf_refresh(args.buf, args.data.client_id)
+      end
+    end,
+  })
+
+  api.nvim_create_autocmd('LspAttach', {
+    buffer = bufnr,
+    group = document_color_augroup,
+    desc = 'Enable document_color when LSP client attaches',
+    callback = function(args)
+      api.nvim_buf_attach(args.buf, false, {
+        on_reload = function(_, buf)
+          buf_clear(buf)
+          if bufstates[buf].enabled then
+            buf_refresh(buf)
+          end
+        end,
+        on_detach = function(_, buf)
+          buf_disable(buf)
+        end,
+      })
+    end,
+  })
+
+  api.nvim_create_autocmd('LspDetach', {
+    buffer = bufnr,
+    group = document_color_augroup,
+    desc = 'Disable document_color if all supporting clients detach',
+    callback = function(args)
+      local clients = lsp.get_clients({ bufnr = args.buf, method = ms.textDocument_documentColor })
+
+      if
+        not vim.iter(clients):any(function(c)
+          return c.id ~= args.data.client_id
+        end)
+      then
+        -- There are no clients left in the buffer that support document color, so turn it off.
+        buf_disable(args.buf)
+      end
+    end,
+  })
+
+  buf_refresh(bufnr)
+end
+
+--- Query whether document colors are enabled in the given buffer.
+---
+--- @param bufnr? integer Buffer handle, or 0 for current. (default: 0)
+--- @return boolean
+function M.is_enabled(bufnr)
+  vim.validate('bufnr', bufnr, 'number', true)
+
+  return bufstates[vim._resolve_bufnr(bufnr)].enabled
+end
+
+--- Enables document highlighting from the given language client in the given buffer.
+---
+--- You can enable document highlighting from a supporting client as follows:
+--- ```lua
+--- vim.api.nvim_create_autocmd('LspAttach', {
+---   callback = function(args)
+---     local client = vim.lsp.get_client_by_id(args.data.client_id)
+---
+---     if client:supports_method('textDocument/documentColor')
+---       vim.lsp.document_color.enable(true, args.buf)
+---     end
+---   end
+--- })
+--- ```
+---
+--- To "toggle", pass the inverse of `is_enabled()`:
+---
+--- ```lua
+--- vim.lsp.document_color.enable(not vim.lsp.document_color.is_enabled())
+--- ```
+---
+--- @param enable? boolean True to enable, false to disable. (default: `true`)
+--- @param bufnr? integer Buffer handle, or 0 for current. (default: 0)
+--- @param opts? vim.lsp.document_color.enable.Opts
+function M.enable(enable, bufnr, opts)
+  vim.validate('enable', enable, 'boolean', true)
+  vim.validate('bufnr', bufnr, 'number', true)
+  vim.validate('opts', opts, 'table', true)
+
+  enable = enable == nil or enable
+  bufnr = vim._resolve_bufnr(bufnr)
+  document_color_opts = vim.tbl_extend('keep', opts or {}, document_color_opts)
+
+  if enable then
+    buf_enable(bufnr)
+  else
+    buf_disable(bufnr)
+  end
+end
+
+api.nvim_create_autocmd('ColorScheme', {
+  pattern = '*',
+  group = document_color_augroup,
+  desc = 'Refresh document_color',
+  callback = function()
+    color_cache = {}
+
+    for _, bufnr in ipairs(api.nvim_list_bufs()) do
+      buf_clear(bufnr)
+      if api.nvim_buf_is_loaded(bufnr) and bufstates[bufnr].enabled then
+        buf_refresh(bufnr)
+      else
+        reset_bufstate(bufnr, false)
+      end
+    end
+  end,
+})
+
+api.nvim_set_decoration_provider(document_color_ns, {
+  on_win = function(_, _, bufnr)
+    if not bufstates[bufnr] then
+      reset_bufstate(bufnr, false)
+    end
+    local bufstate = assert(bufstates[bufnr])
+
+    local all_applied = #bufstate.applied_version > 0
+      and vim.iter(pairs(bufstate.applied_version)):all(function(_, buf_version)
+        return buf_version == bufstate.buf_version
+      end)
+
+    if bufstate.buf_version ~= util.buf_versions[bufnr] or all_applied then
+      return
+    end
+
+    api.nvim_buf_clear_namespace(bufnr, document_color_ns, 0, -1)
+
+    local style = document_color_opts.style
+
+    for client_id, client_hls in pairs(bufstate.hl_info) do
+      if bufstate.applied_version[client_id] ~= bufstate.buf_version then
+        for _, hl in ipairs(client_hls) do
+          if type(style) == 'function' then
+            style(bufnr, hl.range, hl.hex_code)
+          elseif style == 'foreground' or style == 'background' then
+            api.nvim_buf_set_extmark(bufnr, document_color_ns, hl.range[1], hl.range[2], {
+              end_row = hl.range[3],
+              end_col = hl.range[4],
+              hl_group = hl.hl_group,
+              strict = false,
+            })
+          else
+            -- Default swatch: \uf0c8
+            local swatch = style == 'virtual' and ' ' or style
+            api.nvim_buf_set_extmark(bufnr, document_color_ns, hl.range[1], hl.range[2], {
+              virt_text = { { swatch, hl.hl_group } },
+              virt_text_pos = 'inline',
+            })
+          end
+        end
+
+        bufstate.applied_version[client_id] = bufstate.buf_version
+      end
+    end
+  end,
+})
+
+return M

runtime/lua/vim/lsp/protocol.lua

@@ -532,6 +532,9 @@ function protocol.make_client_capabilities()
       callHierarchy = {
         dynamicRegistration = false,
       },
+      colorProvider = {
+        dynamicRegistration = false,
+      },
     },
     workspace = {
       symbol = {