feat(lsp)!: add rule-based sem token highlighting (#22022)

feat(lsp)!: change semantic token highlighting

Change the default highlights used, and add more highlights per token.

Add an LspTokenUpdate event and a highlight_token function.

:Inspect now shows any highlights applied by token highlighting rules,
default or user-defined.

BREAKING CHANGE: change the default highlight groups used by semantic
token highlighting.

2 files changed, 178 insertions, 102 deletions

diff --git a/runtime/lua/vim/_inspector.lua b/runtime/lua/vim/_inspector.lua
index 9e91597192..92d380b08c 100644
--- a/runtime/lua/vim/_inspector.lua
+++ b/runtime/lua/vim/_inspector.lua
@@ -2,7 +2,7 @@
 ---@field syntax boolean include syntax based highlight groups (defaults to true)
 ---@field treesitter boolean include treesitter based highlight groups (defaults to true)
 ---@field extmarks boolean|"all" include extmarks. When `all`, then extmarks without a `hl_group` will also be included (defaults to true)
----@field semantic_tokens boolean include semantic tokens (defaults to true)
+---@field semantic_tokens boolean include semantic token highlights (defaults to true)
 local defaults = {
   syntax = true,
   treesitter = true,
@@ -81,47 +81,54 @@ function vim.inspect_pos(bufnr, row, col, filter)
     end
   end
 
-  -- semantic tokens
-  if filter.semantic_tokens then
-    for _, token in ipairs(vim.lsp.semantic_tokens.get_at_pos(bufnr, row, col) or {}) do
-      token.hl_groups = {
-        type = resolve_hl({ hl_group = '@' .. token.type }),
-        modifiers = vim.tbl_map(function(modifier)
-          return resolve_hl({ hl_group = '@' .. modifier })
-        end, token.modifiers or {}),
-      }
-      table.insert(results.semantic_tokens, token)
+  --- Convert an extmark tuple into a map-like table
+  --- @private
+  local function to_map(extmark)
+    extmark = {
+      id = extmark[1],
+      row = extmark[2],
+      col = extmark[3],
+      opts = resolve_hl(extmark[4]),
+    }
+    extmark.end_row = extmark.opts.end_row or extmark.row -- inclusive
+    extmark.end_col = extmark.opts.end_col or (extmark.col + 1) -- exclusive
+    return extmark
+  end
+
+  --- Check if an extmark overlaps this position
+  --- @private
+  local function is_here(extmark)
+    return (row >= extmark.row and row <= extmark.end_row) -- within the rows of the extmark
+      and (row > extmark.row or col >= extmark.col) -- either not the first row, or in range of the col
+      and (row < extmark.end_row or col < extmark.end_col) -- either not in the last row or in range of the col
+  end
+
+  -- all extmarks at this position
+  local extmarks = {}
+  for ns, nsid in pairs(vim.api.nvim_get_namespaces()) do
+    local ns_marks = vim.api.nvim_buf_get_extmarks(bufnr, nsid, 0, -1, { details = true })
+    ns_marks = vim.tbl_map(to_map, ns_marks)
+    ns_marks = vim.tbl_filter(is_here, ns_marks)
+    for _, mark in ipairs(ns_marks) do
+      mark.ns_id = nsid
+      mark.ns = ns
     end
+    vim.list_extend(extmarks, ns_marks)
+  end
+
+  if filter.semantic_tokens then
+    results.semantic_tokens = vim.tbl_filter(function(extmark)
+      return extmark.ns:find('vim_lsp_semantic_tokens') == 1
+    end, extmarks)
   end
 
-  -- extmarks
   if filter.extmarks then
-    for ns, nsid in pairs(vim.api.nvim_get_namespaces()) do
-      if ns:find('vim_lsp_semantic_tokens') ~= 1 then
-        local extmarks = vim.api.nvim_buf_get_extmarks(bufnr, nsid, 0, -1, { details = true })
-        for _, extmark in ipairs(extmarks) do
-          extmark = {
-            ns_id = nsid,
-            ns = ns,
-            id = extmark[1],
-            row = extmark[2],
-            col = extmark[3],
-            opts = resolve_hl(extmark[4]),
-          }
-          local end_row = extmark.opts.end_row or extmark.row -- inclusive
-          local end_col = extmark.opts.end_col or (extmark.col + 1) -- exclusive
-          if
-            (filter.extmarks == 'all' or extmark.opts.hl_group) -- filter hl_group
-            and (row >= extmark.row and row <= end_row) -- within the rows of the extmark
-            and (row > extmark.row or col >= extmark.col) -- either not the first row, or in range of the col
-            and (row < end_row or col < end_col) -- either not in the last row or in range of the col
-          then
-            table.insert(results.extmarks, extmark)
-          end
-        end
-      end
-    end
+    results.extmarks = vim.tbl_filter(function(extmark)
+      return extmark.ns:find('vim_lsp_semantic_tokens') ~= 1
+        and (filter.extmarks == 'all' or extmark.opts.hl_group)
+    end, extmarks)
   end
+
   return results
 end
 
@@ -174,16 +181,17 @@ function vim.show_pos(bufnr, row, col, filter)
     nl()
   end
 
+  -- semantic tokens
   if #items.semantic_tokens > 0 then
     append('Semantic Tokens', 'Title')
     nl()
-    for _, token in ipairs(items.semantic_tokens) do
-      local client = vim.lsp.get_client_by_id(token.client_id)
-      client = client and (' (' .. client.name .. ')') or ''
-      item(token.hl_groups.type, 'type' .. client)
-      for _, modifier in ipairs(token.hl_groups.modifiers) do
-        item(modifier, 'modifier' .. client)
-      end
+    local sorted_marks = vim.fn.sort(items.semantic_tokens, function(left, right)
+      local left_first = left.opts.priority < right.opts.priority
+        or left.opts.priority == right.opts.priority and left.opts.hl_group < right.opts.hl_group
+      return left_first and -1 or 1
+    end)
+    for _, extmark in ipairs(sorted_marks) do
+      item(extmark.opts, 'priority: ' .. extmark.opts.priority)
     end
     nl()
   end
@@ -197,6 +205,7 @@ function vim.show_pos(bufnr, row, col, filter)
     end
     nl()
   end
+
   -- extmarks
   if #items.extmarks > 0 then
     append('Extmarks', 'Title')
diff --git a/runtime/lua/vim/lsp/semantic_tokens.lua b/runtime/lua/vim/lsp/semantic_tokens.lua
index 24b5c6c24e..7983d066b8 100644
--- a/runtime/lua/vim/lsp/semantic_tokens.lua
+++ b/runtime/lua/vim/lsp/semantic_tokens.lua
@@ -8,8 +8,8 @@ local bit = require('bit')
 --- @field start_col number start column 0-based
 --- @field end_col number end column 0-based
 --- @field type string token type as string
---- @field modifiers string[] token modifiers as strings
---- @field extmark_added boolean whether this extmark has been added to the buffer yet
+--- @field modifiers table token modifiers as a set. E.g., { static = true, readonly = true }
+--- @field marked boolean whether this token has had extmarks applied
 ---
 --- @class STCurrentResult
 --- @field version number document version associated with this result
@@ -36,10 +36,13 @@ local bit = require('bit')
 ---@field client_state table<number, STClientState>
 local STHighlighter = { active = {} }
 
+--- Do a binary search of the tokens in the half-open range [lo, hi).
+---
+--- Return the index i in range such that tokens[j].line < line for all j < i, and
+--- tokens[j].line >= line for all j >= i, or return hi if no such index is found.
+---
 ---@private
-local function binary_search(tokens, line)
-  local lo = 1
-  local hi = #tokens
+local function lower_bound(tokens, line, lo, hi)
   while lo < hi do
     local mid = math.floor((lo + hi) / 2)
     if tokens[mid].line < line then
@@ -51,16 +54,34 @@ local function binary_search(tokens, line)
   return lo
 end
 
+--- Do a binary search of the tokens in the half-open range [lo, hi).
+---
+--- Return the index i in range such that tokens[j].line <= line for all j < i, and
+--- tokens[j].line > line for all j >= i, or return hi if no such index is found.
+---
+---@private
+local function upper_bound(tokens, line, lo, hi)
+  while lo < hi do
+    local mid = math.floor((lo + hi) / 2)
+    if line < tokens[mid].line then
+      hi = mid
+    else
+      lo = mid + 1
+    end
+  end
+  return lo
+end
+
 --- Extracts modifier strings from the encoded number in the token array
 ---
 ---@private
----@return string[]
+---@return table<string, boolean>
 local function modifiers_from_number(x, modifiers_table)
   local modifiers = {}
   local idx = 1
   while x > 0 do
     if bit.band(x, 1) == 1 then
-      modifiers[#modifiers + 1] = modifiers_table[idx]
+      modifiers[modifiers_table[idx]] = true
     end
     x = bit.rshift(x, 1)
     idx = idx + 1
@@ -109,7 +130,7 @@ local function tokens_to_ranges(data, bufnr, client)
         end_col = end_col,
         type = token_type,
         modifiers = modifiers,
-        extmark_added = false,
+        marked = false,
       }
     end
   end
@@ -355,7 +376,7 @@ end
 ---
 ---@private
 function STHighlighter:on_win(topline, botline)
-  for _, state in pairs(self.client_state) do
+  for client_id, state in pairs(self.client_state) do
     local current_result = state.current_result
     if current_result.version and current_result.version == util.buf_versions[self.bufnr] then
       if not current_result.namespace_cleared then
@@ -372,52 +393,55 @@ function STHighlighter:on_win(topline, botline)
       --
       -- Instead, we have to use normal extmarks that can attach to locations
       -- in the buffer and are persisted between redraws.
+      --
+      -- `strict = false` is necessary here for the 1% of cases where the
+      -- current result doesn't actually match the buffer contents. Some
+      -- LSP servers can respond with stale tokens on requests if they are
+      -- still processing changes from a didChange notification.
+      --
+      -- LSP servers that do this _should_ follow up known stale responses
+      -- with a refresh notification once they've finished processing the
+      -- didChange notification, which would re-synchronize the tokens from
+      -- our end.
+      --
+      -- The server I know of that does this is clangd when the preamble of
+      -- a file changes and the token request is processed with a stale
+      -- preamble while the new one is still being built. Once the preamble
+      -- finishes, clangd sends a refresh request which lets the client
+      -- re-synchronize the tokens.
+
+      local set_mark = function(token, hl_group, delta)
+        vim.api.nvim_buf_set_extmark(self.bufnr, state.namespace, token.line, token.start_col, {
+          hl_group = hl_group,
+          end_col = token.end_col,
+          priority = vim.highlight.priorities.semantic_tokens + delta,
+          strict = false,
+        })
+      end
+
+      local ft = vim.bo[self.bufnr].filetype
       local highlights = current_result.highlights
-      local idx = binary_search(highlights, topline)
+      local first = lower_bound(highlights, topline, 1, #highlights + 1)
+      local last = upper_bound(highlights, botline, first, #highlights + 1) - 1
 
-      for i = idx, #highlights do
+      for i = first, last do
         local token = highlights[i]
-
-        if token.line > botline then
-          break
-        end
-
-        if not token.extmark_added then
-          -- `strict = false` is necessary here for the 1% of cases where the
-          -- current result doesn't actually match the buffer contents. Some
-          -- LSP servers can respond with stale tokens on requests if they are
-          -- still processing changes from a didChange notification.
-          --
-          -- LSP servers that do this _should_ follow up known stale responses
-          -- with a refresh notification once they've finished processing the
-          -- didChange notification, which would re-synchronize the tokens from
-          -- our end.
-          --
-          -- The server I know of that does this is clangd when the preamble of
-          -- a file changes and the token request is processed with a stale
-          -- preamble while the new one is still being built. Once the preamble
-          -- finishes, clangd sends a refresh request which lets the client
-          -- re-synchronize the tokens.
-          api.nvim_buf_set_extmark(self.bufnr, state.namespace, token.line, token.start_col, {
-            hl_group = '@' .. token.type,
-            end_col = token.end_col,
-            priority = vim.highlight.priorities.semantic_tokens,
-            strict = false,
-          })
-
-          -- TODO(bfredl) use single extmark when hl_group supports table
-          if #token.modifiers > 0 then
-            for _, modifier in pairs(token.modifiers) do
-              api.nvim_buf_set_extmark(self.bufnr, state.namespace, token.line, token.start_col, {
-                hl_group = '@' .. modifier,
-                end_col = token.end_col,
-                priority = vim.highlight.priorities.semantic_tokens + 1,
-                strict = false,
-              })
-            end
+        if not token.marked then
+          set_mark(token, string.format('@lsp.type.%s.%s', token.type, ft), 0)
+          for modifier, _ in pairs(token.modifiers) do
+            set_mark(token, string.format('@lsp.mod.%s.%s', modifier, ft), 1)
+            set_mark(token, string.format('@lsp.typemod.%s.%s.%s', token.type, modifier, ft), 2)
           end
-
-          token.extmark_added = true
+          token.marked = true
+
+          api.nvim_exec_autocmds('LspTokenUpdate', {
+            pattern = vim.api.nvim_buf_get_name(self.bufnr),
+            modeline = false,
+            data = {
+              token = token,
+              client_id = client_id,
+            },
+          })
         end
       end
     end
@@ -588,7 +612,13 @@ end
 ---@param row number|nil Position row (default cursor position)
 ---@param col number|nil Position column (default cursor position)
 ---
----@return table|nil (table|nil) List of tokens at position
+---@return table|nil (table|nil) List of tokens at position. Each token has
+---        the following fields:
+---        - line (number) line number, 0-based
+---        - start_col (number) start column, 0-based
+---        - end_col (number) end column, 0-based
+---        - type (string) token type as string, e.g. "variable"
+---        - modifiers (table) token modifiers as a set. E.g., { static = true, readonly = true }
 function M.get_at_pos(bufnr, row, col)
   if bufnr == nil or bufnr == 0 then
     bufnr = api.nvim_get_current_buf()
@@ -608,7 +638,7 @@ function M.get_at_pos(bufnr, row, col)
   for client_id, client in pairs(highlighter.client_state) do
     local highlights = client.current_result.highlights
     if highlights then
-      local idx = binary_search(highlights, row)
+      local idx = lower_bound(highlights, row, 1, #highlights + 1)
       for i = idx, #highlights do
         local token = highlights[i]
 
@@ -631,23 +661,60 @@ end
 --- Only has an effect if the buffer is currently active for semantic token
 --- highlighting (|vim.lsp.semantic_tokens.start()| has been called for it)
 ---
----@param bufnr (nil|number) default: current buffer
+---@param bufnr (number|nil) filter by buffer. All buffers if nil, current
+---       buffer if 0
 function M.force_refresh(bufnr)
   vim.validate({
     bufnr = { bufnr, 'n', true },
   })
 
-  if bufnr == nil or bufnr == 0 then
-    bufnr = api.nvim_get_current_buf()
+  local buffers = bufnr == nil and vim.tbl_keys(STHighlighter.active)
+    or bufnr == 0 and { api.nvim_get_current_buf() }
+    or { bufnr }
+
+  for _, buffer in ipairs(buffers) do
+    local highlighter = STHighlighter.active[buffer]
+    if highlighter then
+      highlighter:reset()
+      highlighter:send_request()
+    end
   end
+end
 
+--- Highlight a semantic token.
+---
+--- Apply an extmark with a given highlight group for a semantic token. The
+--- mark will be deleted by the semantic token engine when appropriate; for
+--- example, when the LSP sends updated tokens. This function is intended for
+--- use inside |LspTokenUpdate| callbacks.
+---@param token (table) a semantic token, found as `args.data.token` in
+---       |LspTokenUpdate|.
+---@param bufnr (number) the buffer to highlight
+---@param client_id (number) The ID of the |vim.lsp.client|
+---@param hl_group (string) Highlight group name
+---@param opts (table|nil) Optional parameters.
+---       - priority: (number|nil) Priority for the applied extmark. Defaults
+---         to `vim.highlight.priorities.semantic_tokens + 3`
+function M.highlight_token(token, bufnr, client_id, hl_group, opts)
   local highlighter = STHighlighter.active[bufnr]
   if not highlighter then
     return
   end
 
-  highlighter:reset()
-  highlighter:send_request()
+  local state = highlighter.client_state[client_id]
+  if not state then
+    return
+  end
+
+  opts = opts or {}
+  local priority = opts.priority or vim.highlight.priorities.semantic_tokens + 3
+
+  vim.api.nvim_buf_set_extmark(bufnr, state.namespace, token.line, token.start_col, {
+    hl_group = hl_group,
+    end_col = token.end_col,
+    priority = priority,
+    strict = false,
+  })
 end
 
 --- |lsp-handler| for the method `workspace/semanticTokens/refresh`