feat(lsp): more annotations

runtime/lua/vim/lsp/semantic_tokens.lua

@@ -10,7 +10,7 @@ local uv = vim.uv
 --- @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 modifiers table<string,boolean> token modifiers as a set. E.g., { static = true, readonly = true }
 --- @field marked boolean whether this token has had extmarks applied
 ---
 --- @class STCurrentResult
@@ -21,8 +21,8 @@ local uv = vim.uv
 --- @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
+--- @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
@@ -72,9 +72,11 @@ end
 
 --- Extracts modifier strings from the encoded number in the token array
 ---
+---@param x integer
+---@param modifiers_table table<integer,string>
 ---@return table<string, boolean>
 local function modifiers_from_number(x, modifiers_table)
-  local modifiers = {}
+  local modifiers = {} ---@type table<string,boolean>
   local idx = 1
   while x > 0 do
     if bit.band(x, 1) == 1 then
@@ -89,20 +91,24 @@ end
 
 --- Converts a raw token list to a list of highlight ranges used by the on_win callback
 ---
+---@param data integer[]
+---@param bufnr integer
+---@param client lsp.Client
+---@param request STActiveRequest
 ---@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 ranges = {} ---@type STTokenRange[]
 
   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 line ---@type integer?
   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
@@ -167,6 +173,7 @@ end
 ---
 ---@private
 ---@param bufnr integer
+---@return STHighlighter
 function STHighlighter.new(bufnr)
   local self = setmetatable({}, { __index = STHighlighter })
 
@@ -221,7 +228,7 @@ function STHighlighter.new(bufnr)
   return self
 end
 
----@private
+---@package
 function STHighlighter:destroy()
   for client_id, _ in pairs(self.client_state) do
     self:detach(client_id)
@@ -231,7 +238,7 @@ function STHighlighter:destroy()
   STHighlighter.active[self.bufnr] = nil
 end
 
----@private
+---@package
 function STHighlighter:attach(client_id)
   local state = self.client_state[client_id]
   if not state then
@@ -244,7 +251,7 @@ function STHighlighter:attach(client_id)
   end
 end
 
----@private
+---@package
 function STHighlighter:detach(client_id)
   local state = self.client_state[client_id]
   if state then
@@ -267,7 +274,7 @@ end
 --- Finally, if the request was successful, the requestId and document version
 --- are saved to facilitate document synchronization in the response.
 ---
----@private
+---@package
 function STHighlighter:send_request()
   local version = util.buf_versions[self.bufnr]
 
@@ -303,7 +310,8 @@ function STHighlighter:send_request()
         -- 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]
+        local bufnr = assert(ctx.bufnr)
+        local highlighter = STHighlighter.active[bufnr]
         if not err and c and highlighter then
           coroutine.wrap(STHighlighter.process_response)(highlighter, response, c, version)
         end
@@ -328,6 +336,7 @@ end
 --- Finally, a redraw command is issued to force nvim to redraw the screen to
 --- pick up changed highlight tokens.
 ---
+---@param response lsp.SemanticTokens|lsp.SemanticTokensDelta
 ---@private
 function STHighlighter:process_response(response, client, version)
   local state = self.client_state[client.id]
@@ -348,15 +357,15 @@ function STHighlighter:process_response(response, client, version)
 
   -- if we have a response to a delta request, update the state of our tokens
   -- appropriately. if it's a full response, just use that
-  local tokens
+  local tokens ---@type integer[]
   local token_edits = response.edits
   if token_edits then
     table.sort(token_edits, function(a, b)
       return a.start < b.start
     end)
 
-    tokens = {}
-    local old_tokens = state.current_result.tokens
+    tokens = {} --- @type integer[]
+    local old_tokens = assert(state.current_result.tokens)
     local idx = 1
     for _, token_edit in ipairs(token_edits) do
       vim.list_extend(tokens, old_tokens, idx, token_edit.start)
@@ -404,7 +413,9 @@ end
 --- handler to avoid the "blink" that occurs due to the timing between the
 --- response handler and the actual redraw.
 ---
----@private
+---@package
+---@param topline integer
+---@param botline integer
 function STHighlighter:on_win(topline, botline)
   for client_id, state in pairs(self.client_state) do
     local current_result = state.current_result
@@ -450,7 +461,7 @@ function STHighlighter:on_win(topline, botline)
       end
 
       local ft = vim.bo[self.bufnr].filetype
-      local highlights = current_result.highlights
+      local highlights = assert(current_result.highlights)
       local first = lower_bound(highlights, topline, 1, #highlights + 1)
       local last = upper_bound(highlights, botline, first, #highlights + 1) - 1
 
@@ -480,7 +491,7 @@ end
 
 --- Reset the buffer's highlighting state and clears the extmark highlights.
 ---
----@private
+---@package
 function STHighlighter:reset()
   for client_id, state in pairs(self.client_state) do
     api.nvim_buf_clear_namespace(self.bufnr, state.namespace, 0, -1)
@@ -499,7 +510,7 @@ end
 --- in the on_win callback. The rest of the current results are saved
 --- in case the server supports delta requests.
 ---
----@private
+---@package
 ---@param client_id integer
 function STHighlighter:mark_dirty(client_id)
   local state = self.client_state[client_id]
@@ -521,7 +532,7 @@ function STHighlighter:mark_dirty(client_id)
   end
 end
 
----@private
+---@package
 function STHighlighter:on_change()
   self:reset_timer()
   if self.debounce > 0 then
@@ -636,6 +647,9 @@ function M.stop(bufnr, client_id)
   end
 end
 
+--- @class STTokenRangeInspect : STTokenRange
+--- @field client_id integer
+
 --- Return the semantic token(s) at the given position.
 --- If called without arguments, returns the token under the cursor.
 ---
@@ -643,13 +657,14 @@ end
 ---@param row integer|nil Position row (default cursor position)
 ---@param col integer|nil Position column (default cursor position)
 ---
----@return table|nil (table|nil) List of tokens at position. Each token has
+---@return STTokenRangeInspect[]|nil (table|nil) List of tokens at position. Each token has
 ---        the following fields:
 ---        - line (integer) line number, 0-based
 ---        - start_col (integer) start column, 0-based
 ---        - end_col (integer) 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 }
+---        - client_id (integer)
 function M.get_at_pos(bufnr, row, col)
   if bufnr == nil or bufnr == 0 then
     bufnr = api.nvim_get_current_buf()
@@ -665,13 +680,14 @@ function M.get_at_pos(bufnr, row, col)
     row, col = cursor[1] - 1, cursor[2]
   end
 
-  local tokens = {}
+  local tokens = {} --- @type STTokenRangeInspect[]
   for client_id, client in pairs(highlighter.client_state) do
     local highlights = client.current_result.highlights
     if highlights then
       local idx = lower_bound(highlights, row, 1, #highlights + 1)
       for i = idx, #highlights do
         local token = highlights[i]
+        --- @cast token STTokenRangeInspect
 
         if token.line > row then
           break