path: root/runtime/lua/vim/lsp/semantic_tokens.lua

local api = vim.api
local bit = require('bit')
local handlers = require('vim.lsp.handlers')
local ms = require('vim.lsp.protocol').Methods
local util = require('vim.lsp.util')
local uv = vim.uv

--- @class STTokenRange
--- @field line integer line number 0-based
--- @field start_col integer start column 0-based
--- @field end_col integer end column 0-based
--- @field type string token type as string
--- @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? integer document version associated with this result
--- @field result_id? string resultId from the server; used with delta requests
--- @field highlights? STTokenRange[] cache of highlight ranges for this document version
--- @field tokens? integer[] raw token array as received by the server. used for calculating delta responses
--- @field namespace_cleared? boolean whether the namespace was cleared for this result yet
---
--- @class STActiveRequest
--- @field request_id integer the LSP request ID of the most recent request sent to the server
--- @field version integer the document version associated with the most recent request
---
--- @class STClientState
--- @field namespace integer
--- @field active_request STActiveRequest
--- @field current_result STCurrentResult

---@class STHighlighter
---@field active table<integer, STHighlighter>
---@field bufnr integer
---@field augroup integer augroup for buffer events
---@field debounce integer milliseconds to debounce requests for new tokens
---@field timer table uv_timer for debouncing requests for new tokens
---@field client_state table<integer, 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.
local function lower_bound(tokens, line, lo, hi)
  while lo < hi do
    local mid = bit.rshift(lo + hi, 1) -- Equivalent to floor((lo + hi) / 2).
    if tokens[mid].line < line then
      lo = mid + 1
    else
      hi = mid
    end
  end
  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.
local function upper_bound(tokens, line, lo, hi)
  while lo < hi do
    local mid = bit.rshift(lo + hi, 1) -- Equivalent to 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
---
---@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_table[idx]] = true
    end
    x = bit.rshift(x, 1)
    idx = idx + 1
  end

  return modifiers
end

--- Converts a raw token list to a list of highlight ranges used by the on_win callback
---
---@return STTokenRange[]
local function tokens_to_ranges(data, bufnr, client, request)
  local legend = client.server_capabilities.semanticTokensProvider.legend
  local token_types = legend.tokenTypes
  local token_modifiers = legend.tokenModifiers
  local lines = api.nvim_buf_get_lines(bufnr, 0, -1, false)
  local ranges = {}

  local start = uv.hrtime()
  local ms_to_ns = 1000 * 1000
  local yield_interval_ns = 5 * ms_to_ns
  local co, is_main = coroutine.running()

  local line
  local start_char = 0
  for i = 1, #data, 5 do
    -- if this function is called from the main coroutine, let it run to completion with no yield
    if not is_main then
      local elapsed_ns = uv.hrtime() - start

      if elapsed_ns > yield_interval_ns then
        vim.schedule(function()
          coroutine.resume(co, util.buf_versions[bufnr])
        end)
        if request.version ~= coroutine.yield() then
          -- request became stale since the last time the coroutine ran.
          -- abandon it by yielding without a way to resume
          coroutine.yield()
        end

        start = uv.hrtime()
      end
    end

    local delta_line = data[i]
    line = line and line + delta_line or delta_line
    local delta_start = data[i + 1]
    start_char = delta_line == 0 and start_char + delta_start or delta_start

    -- data[i+3] +1 because Lua tables are 1-indexed
    local token_type = token_types[data[i + 3] + 1]
    local modifiers = modifiers_from_number(data[i + 4], token_modifiers)

    local function _get_byte_pos(col)
      if col > 0 then
        local buf_line = lines[line + 1] or ''
        local ok, result
        ok, result = pcall(util._str_byteindex_enc, buf_line, col, client.offset_encoding)
        if ok then
          return result
        end
        return math.min(#buf_line, col)
      end
      return col
    end

    local start_col = _get_byte_pos(start_char)
    local end_col = _get_byte_pos(start_char + data[i + 2])

    if token_type then
      ranges[#ranges + 1] = {
        line = line,
        start_col = start_col,
        end_col = end_col,
        type = token_type,
        modifiers = modifiers,
        marked = false,
      }
    end
  end

  return ranges
end

--- Construct a new STHighlighter for the buffer
---
---@private
---@param bufnr integer
function STHighlighter.new(bufnr)
  local self = setmetatable({}, { __index = STHighlighter })

  self.bufnr = bufnr
  self.augroup = api.nvim_create_augroup('vim_lsp_semantic_tokens:' .. bufnr, { clear = true })
  self.client_state = {}

  STHighlighter.active[bufnr] = self

  api.nvim_buf_attach(bufnr, false, {
    on_lines = function(_, buf)
      local highlighter = STHighlighter.active[buf]
      if not highlighter then
        return true
      end
      highlighter:on_change()
    end,
    on_reload = function(_, buf)
      local highlighter = STHighlighter.active[buf]
      if highlighter then
        highlighter:reset()
        highlighter:send_request()
      end
    end,
    on_detach = function(_, buf)
      local highlighter = STHighlighter.active[buf]
      if highlighter then
        highlighter:destroy()
      end
    end,
  })

  api.nvim_create_autocmd({ 'BufWinEnter', 'InsertLeave' }, {
    buffer = self.bufnr,
    group = self.augroup,
    callback = function()
      self:send_request()
    end,
  })

  api.nvim_create_autocmd('LspDetach', {
    buffer = self.bufnr,
    group = self.augroup,
    callback = function(args)
      self:detach(args.data.client_id)
      if vim.tbl_isempty(self.client_state) then
        self:destroy()
      end
    end,
  })

  return self
end

---@private
function STHighlighter:destroy()
  for client_id, _ in pairs(self.client_state) do
    self:detach(client_id)
  end

  api.nvim_del_augroup_by_id(self.augroup)
  STHighlighter.active[self.bufnr] = nil
end

---@private
function STHighlighter:attach(client_id)
  local state = self.client_state[client_id]
  if not state then
    state = {
      namespace = api.nvim_create_namespace('vim_lsp_semantic_tokens:' .. client_id),
      active_request = {},
      current_result = {},
    }
    self.client_state[client_id] = state
  end
end

---@private
function STHighlighter:detach(client_id)
  local state = self.client_state[client_id]
  if state then
    --TODO: delete namespace if/when that becomes possible
    api.nvim_buf_clear_namespace(self.bufnr, state.namespace, 0, -1)
    self.client_state[client_id] = nil
  end
end

--- This is the entry point for getting all the tokens in a buffer.
---
--- For the given clients (or all attached, if not provided), this sends a request
--- to ask for semantic tokens. If the server supports delta requests, that will
--- be prioritized if we have a previous requestId and token array.
---
--- This function will skip servers where there is an already an active request in
--- flight for the same version. If there is a stale request in flight, that is
--- cancelled prior to sending a new one.
---
--- Finally, if the request was successful, the requestId and document version
--- are saved to facilitate document synchronization in the response.
---
---@private
function STHighlighter:send_request()
  local version = util.buf_versions[self.bufnr]

  self:reset_timer()

  for client_id, state in pairs(self.client_state) do
    local client = vim.lsp.get_client_by_id(client_id)

    local current_result = state.current_result
    local active_request = state.active_request

    -- Only send a request for this client if the current result is out of date and
    -- there isn't a current a request in flight for this version
    if client and current_result.version ~= version and active_request.version ~= version then
      -- cancel stale in-flight request
      if active_request.request_id then
        client.cancel_request(active_request.request_id)
        active_request = {}
        state.active_request = active_request
      end

      local spec = client.server_capabilities.semanticTokensProvider.full
      local hasEditProvider = type(spec) == 'table' and spec.delta

      local params = { textDocument = util.make_text_document_params(self.bufnr) }
      local method = ms.textDocument_semanticTokens_full

      if hasEditProvider and current_result.result_id then
        method = method .. '/delta'
        params.previousResultId = current_result.result_id
      end
      local success, request_id = client.request(method, params, function(err, response, ctx)
        -- look client up again using ctx.client_id instead of using a captured
        -- client object
        local c = vim.lsp.get_client_by_id(ctx.client_id)
        local highlighter = STHighlighter.active[ctx.bufnr]
        if not err and c and highlighter then
          coroutine.wrap(STHighlighter.process_response)(highlighter, response, c, version)
        end
      end, self.bufnr)

      if success then
        active_request.request_id = request_id
        active_request.version = version
      end
    end
  end
end

--- This function will parse the semantic token responses and set up the cache
--- (current_result). It also performs document synchronization by checking the
--- version of the document associated with the resulting request_id and only
--- performing work if the response is not out-of-date.
---
--- Delta edits are applied if necessary, and new highlight ranges are calculated
--- and stored in the buffer state.
---
--- Finally, a redraw command is issued to force nvim to redraw the screen to
--- pick up changed highlight tokens.
---
---@private
function S