fix(treesitter): fix most diagnostics

6 files changed, 331 insertions, 125 deletions

diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index 582922ecb6..96b1e24ba9 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -3,8 +3,11 @@ local query = require('vim.treesitter.query')
 local language = require('vim.treesitter.language')
 local LanguageTree = require('vim.treesitter.languagetree')
 
+---@type table<integer,LanguageTree>
 local parsers = setmetatable({}, { __mode = 'v' })
 
+---@class TreesitterModule
+---@field highlighter TSHighlighter
 local M = vim.tbl_extend('error', query, language)
 
 M.language_version = vim._ts_get_language_version()
@@ -12,6 +15,7 @@ M.minimum_language_version = vim._ts_get_minimum_language_version()
 
 setmetatable(M, {
   __index = function(t, k)
+    ---@diagnostic disable:no-unknown
     if k == 'highlighter' then
       t[k] = require('vim.treesitter.highlighter')
       return t[k]
@@ -29,21 +33,23 @@ setmetatable(M, {
 ---
 --- It is not recommended to use this; use |get_parser()| instead.
 ---
----@param bufnr string Buffer the parser will be tied to (0 for current buffer)
+---@param bufnr integer Buffer the parser will be tied to (0 for current buffer)
 ---@param lang string Language of the parser
 ---@param opts (table|nil) Options to pass to the created language tree
 ---
----@return LanguageTree |LanguageTree| object to use for parsing
+---@return LanguageTree object to use for parsing
 function M._create_parser(bufnr, lang, opts)
   language.require_language(lang)
   if bufnr == 0 then
-    bufnr = a.nvim_get_current_buf()
+    bufnr = vim.api.nvim_get_current_buf()
   end
 
   vim.fn.bufload(bufnr)
 
   local self = LanguageTree.new(bufnr, lang, opts)
 
+  ---@diagnostic disable:invisible
+
   ---@private
   local function bytes_cb(_, ...)
     self:_on_bytes(...)
@@ -58,12 +64,14 @@ function M._create_parser(bufnr, lang, opts)
   end
 
   ---@private
-  local function reload_cb(_, ...)
-    self:_on_reload(...)
+  local function reload_cb(_)
+    self:_on_reload()
   end
 
+  local source = self:source() --[[@as integer]]
+
   a.nvim_buf_attach(
-    self:source(),
+    source,
     false,
     { on_bytes = bytes_cb, on_detach = detach_cb, on_reload = reload_cb, preview = true }
   )
@@ -77,11 +85,11 @@ end
 ---
 --- If needed, this will create the parser.
 ---
----@param bufnr (number|nil) Buffer the parser should be tied to (default: current buffer)
+---@param bufnr (integer|nil) Buffer the parser should be tied to (default: current buffer)
 ---@param lang (string|nil) Filetype of this parser (default: buffer filetype)
 ---@param opts (table|nil) Options to pass to the created language tree
 ---
----@return LanguageTree |LanguageTree| object to use for parsing
+---@return LanguageTree object to use for parsing
 function M.get_parser(bufnr, lang, opts)
   opts = opts or {}
 
@@ -107,7 +115,7 @@ end
 ---@param lang string Language of this string
 ---@param opts (table|nil) Options to pass to the created language tree
 ---
----@return LanguageTree |LanguageTree| object to use for parsing
+---@return LanguageTree object to use for parsing
 function M.get_string_parser(str, lang, opts)
   vim.validate({
     str = { str, 'string' },
@@ -120,8 +128,8 @@ end
 
 --- Determines whether a node is the ancestor of another
 ---
----@param dest userdata Possible ancestor |tsnode|
----@param source userdata Possible descendant |tsnode|
+---@param dest TSNode Possible ancestor
+---@param source TSNode Possible descendant
 ---
 ---@return boolean True if {dest} is an ancestor of {source}
 function M.is_ancestor(dest, source)
@@ -143,9 +151,12 @@ end
 
 --- Returns the node's range or an unpacked range table
 ---
----@param node_or_range (userdata|table) |tsnode| or table of positions
+---@param node_or_range (TSNode|table) Node or table of positions
 ---
----@return table `{ start_row, start_col, end_row, end_col }`
+---@return integer start_row
+---@return integer start_col
+---@return integer end_row
+---@return integer end_col
 function M.get_node_range(node_or_range)
   if type(node_or_range) == 'table' then
     return unpack(node_or_range)
@@ -156,9 +167,9 @@ end
 
 --- Determines whether (line, col) position is in node range
 ---
----@param node userdata |tsnode| defining the range
----@param line number Line (0-based)
----@param col number Column (0-based)
+---@param node TSNode defining the range
+---@param line integer Line (0-based)
+---@param col integer Column (0-based)
 ---
 ---@return boolean True if the position is in node range
 function M.is_in_node_range(node, line, col)
@@ -180,7 +191,7 @@ end
 
 --- Determines if a node contains a range
 ---
----@param node userdata |tsnode|
+---@param node TSNode
 ---@param range table
 ---
 ---@return boolean True if the {node} contains the {range}
@@ -197,9 +208,9 @@ end
 --- Each capture is represented by a table containing the capture name as a string as
 --- well as a table of metadata (`priority`, `conceal`, ...; empty if none are defined).
 ---
----@param bufnr number Buffer number (0 for current buffer)
----@param row number Position row
----@param col number Position column
+---@param bufnr integer Buffer number (0 for current buffer)
+---@param row integer Position row
+---@param col integer Position column
 ---
 ---@return table[] List of captures `{ capture = "capture name", metadata = { ... } }`
 function M.get_captures_at_pos(bufnr, row, col)
@@ -250,7 +261,7 @@ end
 
 --- Returns a list of highlight capture names under the cursor
 ---
----@param winnr (number|nil) Window handle or 0 for current window (default)
+---@param winnr (integer|nil) Window handle or 0 for current window (default)
 ---
 ---@return string[] List of capture names
 function M.get_captures_at_cursor(winnr)
@@ -271,14 +282,14 @@ end
 
 --- Returns the smallest named node at the given position
 ---
----@param bufnr number Buffer number (0 for current buffer)
----@param row number Position row
----@param col number Position column
+---@param bufnr integer Buffer number (0 for current buffer)
+---@param row integer Position row
+---@param col integer Position column
 ---@param opts table Optional keyword arguments:
 ---             - lang string|nil Parser language
 ---             - ignore_injections boolean Ignore injected languages (default true)
 ---
----@return userdata|nil |tsnode| under the cursor
+---@return TSNode|nil under the cursor
 function M.get_node_at_pos(bufnr, row, col, opts)
   if bufnr == 0 then
     bufnr = a.nvim_get_current_buf()
@@ -295,7 +306,7 @@ end
 
 --- Returns the smallest named node under the cursor
 ---
----@param winnr (number|nil) Window handle or 0 for current window (default)
+---@param winnr (integer|nil) Window handle or 0 for current window (default)
 ---
 ---@return string Name of node under the cursor
 function M.get_node_at_cursor(winnr)
@@ -323,7 +334,7 @@ end
 --- })
 --- </pre>
 ---
----@param bufnr (number|nil) Buffer to be highlighted (default: current buffer)
+---@param bufnr (integer|nil) Buffer to be highlighted (default: current buffer)
 ---@param lang (string|nil) Language of the parser (default: buffer filetype)
 function M.start(bufnr, lang)
   bufnr = bufnr or a.nvim_get_current_buf()
@@ -333,7 +344,7 @@ end
 
 --- Stops treesitter highlighting for a buffer
 ---
----@param bufnr (number|nil) Buffer to stop highlighting (default: current buffer)
+---@param bufnr (integer|nil) Buffer to stop highlighting (default: current buffer)
 function M.stop(bufnr)
   bufnr = bufnr or a.nvim_get_current_buf()
 
@@ -351,13 +362,13 @@ end
 ---@param opts table|nil Optional options table with the following possible keys:
 ---                      - lang (string|nil): The language of the source buffer. If omitted, the
 ---                        filetype of the source buffer is used.
----                      - bufnr (number|nil): Buffer to draw the tree into. If omitted, a new
+---                      - bufnr (integer|nil): Buffer to draw the tree into. If omitted, a new
 ---                        buffer is created.
----                      - winid (number|nil): Window id to display the tree buffer in. If omitted,
+---                      - winid (integer|nil): Window id to display the tree buffer in. If omitted,
 ---                        a new window is created with {command}.
 ---                      - command (string|nil): Vimscript command to create the window. Default
 ---                        value is "topleft 60vnew". Only used when {winid} is nil.
----                      - title (string|fun(bufnr:number):string|nil): Title of the window. If a
+---                      - title (string|fun(bufnr:integer):string|nil): Title of the window. If a
 ---                        function, it accepts the buffer number of the source buffer as its only
 ---                        argument and should return a string.
 function M.show_tree(opts)
diff --git a/runtime/lua/vim/treesitter/_meta.lua b/runtime/lua/vim/treesitter/_meta.lua
new file mode 100644
index 0000000000..87b4560798
--- /dev/null
+++ b/runtime/lua/vim/treesitter/_meta.lua
@@ -0,0 +1,60 @@
+---@meta
+
+---@class TSNode
+---@field id fun(self: TSNode): integer
+---@field range fun(self: TSNode): integer, integer, integer, integer
+---@field start fun(self: TSNode): integer, integer, integer
+---@field end_ fun(self: TSNode): integer, integer, integer
+---@field type fun(self: TSNode): string
+---@field symbol fun(self: TSNode): integer
+---@field named fun(self: TSNode): boolean
+---@field missing fun(self: TSNode): boolean
+---@field child_count fun(self: TSNode): integer
+---@field named_child_count fun(self: TSNode): integer
+---@field child fun(self: TSNode, integer): TSNode
+---@field name_child fun(self: TSNode, integer): TSNode
+---@field descendant_for_range fun(self: TSNode, integer, integer, integer, integer): TSNode
+---@field named_descendant_for_range fun(self: TSNode, integer, integer, integer, integer): TSNode
+---@field parent fun(self: TSNode): TSNode
+---@field next_sibling fun(self: TSNode): TSNode
+---@field prev_sibling fun(self: TSNode): TSNode
+---@field next_named_sibling fun(self: TSNode): TSNode
+---@field prev_named_sibling fun(self: TSNode): TSNode
+---@field named_children fun(self: TSNode): TSNode[]
+---@field has_error fun(self: TSNode): boolean
+---@field iter_children fun(self: TSNode): fun(): TSNode, string
+local TSNode = {}
+
+---@param query userdata
+---@param captures true
+---@param start integer
+---@param end_ integer
+---@return fun(): integer, TSNode, any
+function TSNode:_rawquery(query, captures, start, end_) end
+
+---@param query userdata
+---@param captures false
+---@param start integer
+---@param end_ integer
+---@return fun(): string, any
+function TSNode:_rawquery(query, captures, start, end_) end
+
+---@class TSParser
+---@field parse fun(self: TSParser, tree, source: integer|string): TSTree, integer[]
+---@field included_ranges fun(self: TSParser): integer[]
+---@field set_included_ranges fun(self: TSParser, ranges: integer[][])
+
+---@class TSTree
+---@field root fun(self: TSTree): TSNode
+---@field edit fun(self: TSTree, _: integer, _: integer, _: integer, _: integer, _: integer, _: integer, _: integer, _: integer, _:integer)
+---@field copy fun(self: TSTree): TSTree
+
+---@return integer
+vim._ts_get_language_version = function() end
+
+---@return integer
+vim._ts_get_minimum_language_version = function() end
+
+---@param lang string
+---@return TSParser
+vim._create_ts_parser = function(lang) end
diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index d77a0d0d03..8adaa4ef2f 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -1,13 +1,27 @@
 local a = vim.api
 local query = require('vim.treesitter.query')
 
--- support reload for quick experimentation
+---@alias TSHlIter fun(): integer, TSNode, TSMetadata
+
+---@class TSHighlightState
+---@field next_row integer
+---@field iter TSHlIter|nil
+
 ---@class TSHighlighter
+---@field active table<integer,TSHighlighter>
+---@field bufnr integer
+---@field orig_spelloptions string
+---@field _highlight_states table<TSTree,TSHighlightState>
+---@field _queries table<string,TSHighlighterQuery>
+---@field tree LanguageTree
 local TSHighlighter = rawget(vim.treesitter, 'TSHighlighter') or {}
 TSHighlighter.__index = TSHighlighter
 
 TSHighlighter.active = TSHighlighter.active or {}
 
+---@class TSHighlighterQuery
+---@field _query Query|nil
+---@field hl_cache table<integer,integer>
 local TSHighlighterQuery = {}
 TSHighlighterQuery.__index = TSHighlighterQuery
 
@@ -46,7 +60,7 @@ end
 
 --- Creates a new highlighter using @param tree
 ---
----@param tree LanguageTree |LanguageTree| parser object to use for highlighting
+---@param tree LanguageTree parser object to use for highlighting
 ---@param opts (table|nil) Configuration of the highlighter:
 ---           - queries table overwrite queries used by the highlighter
 ---@return TSHighlighter Created highlighter object
@@ -57,9 +71,10 @@ function TSHighlighter.new(tree, opts)
     error('TSHighlighter can not be used with a string parser source.')
   end
 
-  opts = opts or {}
+  opts = opts or {} ---@type { queries: table<string,string> }
   self.tree = tree
   tree:register_cbs({
+    ---@diagnostic disable:invisible
     on_changedtree = function(...)
       self:on_changedtree(...)
     end,
@@ -67,17 +82,20 @@ function TSHighlighter.new(tree, opts)
       self:on_bytes(...)
     end,
     on_detach = function(...)
+      ---@diagnostic disable-next-line:redundant-parameter
       self:on_detach(...)
     end,
   })
 
-  self.bufnr = tree:source()
+  self.bufnr = tree:source() --[[@as integer]]
   self.edit_count = 0
   self.redraw_count = 0
   self.line_count = {}
   -- A map of highlight states.
   -- This state is kept during rendering across each line update.
   self._highlight_states = {}
+
+  ---@type table<string,TSHighlighterQuery>
   self._queries = {}
 
   -- Queries for a specific language can be overridden by a custom
@@ -128,6 +146,8 @@ function TSHighlighter:destroy()
 end
 
 ---@private
+---@param tstree TSTree
+---@return TSHighlightState
 function TSHighlighter:get_highlight_state(tstree)
   if not self._highlight_states[tstree] then
     self._highlight_states[tstree] = {
@@ -145,6 +165,8 @@ function TSHighlighter:reset_highlight_state()
 end
 
 ---@private
+---@param start_row integer
+---@param new_end integer
 function TSHighlighter:on_bytes(_, _, start_row, _, _, _, _, _, new_end)
   a.nvim__buf_redraw_range(self.bufnr, start_row, start_row + new_end + 1)
 end
@@ -155,6 +177,7 @@ function TSHighlighter:on_detach()
 end
 
 ---@private
+---@param changes integer[][]?
 function TSHighlighter:on_changedtree(changes)
   for _, ch in ipairs(changes or {}) do
     a.nvim__buf_redraw_range(self.bufnr, ch[1], ch[3] + 1)
@@ -165,7 +188,7 @@ end
 --
 ---@private
 ---@param lang string Language used by the highlighter.
----@return Query
+---@return TSHighlighterQuery
 function TSHighlighter:get_query(lang)
   if not self._queries[lang] then
     self._queries[lang] = TSHighlighterQuery.new(lang)
@@ -175,7 +198,12 @@ function TSHighlighter:get_query(lang)
 end
 
 ---@private
+---@param self TSHighlighter
+---@param buf integer
+---@param line integer
+---@param is_spell_nav boolean
 local function on_line_impl(self, buf, line, is_spell_nav)
+  ---@diagnostic disable:invisible
   self.tree:for_each_tree(function(tstree, tree)
     if not tstree then
       return
@@ -213,7 +241,7 @@ local function on_line_impl(self, buf, line, is_spell_nav)
       local hl = highlighter_query.hl_cache[capture]
 
       local capture_name = highlighter_query:query().captures[capture]
-      local spell = nil
+      local spell = nil ---@type boolean?
       if capture_name == 'spell' then
         spell = true
       elseif capture_name == 'nospell' then
@@ -242,6 +270,9 @@ local function on_line_impl(self, buf, line, is_spell_nav)
 end
 
 ---@private
+---@param _win integer
+---@param buf integer
+---@param line integer
 function TSHighlighter._on_line(_, _win, buf, line, _)
   local self = TSHighlighter.active[buf]
   if not self then
@@ -252,6 +283,9 @@ function TSHighlighter._on_line(_, _win, buf, line, _)
 end
 
 ---@private
+---@param buf integer
+---@param srow integer
+---@param erow integer
 function TSHighlighter._on_spell_nav(_, _, buf, srow, _, erow, _)
   local self = TSHighlighter.active[buf]
   if not self then
@@ -266,6 +300,7 @@ function TSHighlighter._on_spell_nav(_, _, buf, srow, _, erow, _)
 end
 
 ---@private
+---@param buf integer
 function TSHighlighter._on_buf(_, buf)
   local self = TSHighlighter.active[buf]
   if self then
@@ -274,6 +309,9 @@ function TSHighlighter._on_buf(_, buf)
 end
 
 ---@private
+---@param _win integer
+---@param buf integer
+---@param _topline integer
 function TSHighlighter._on_win(_, _win, buf, _topline)
   local self = TSHighlighter.active[buf]
   if not self then
diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index a1e96f8ef2..89aac3ae26 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -2,20 +2,39 @@ local a = vim.api
 local query = require('vim.treesitter.query')
 local language = require('vim.treesitter.language')
 
----@class LanguageTree
----@field _callbacks function[] Callback handlers
----@field _children LanguageTree[] Injected languages
----@field _injection_query table Queries defining injected languages
----@field _opts table Options
----@field _parser userdata Parser for language
----@field _regions table List of regions this tree should manage and parse
----@field _lang string Language name
----@field _regions table
----@field _source (number|string) Buffer or string to parse
----@field _trees userdata[] Reference to parsed |tstree| (one for each language)
----@field _valid boolean If the parsed tree is valid
+---@alias Range {[1]: integer, [2]: integer, [3]: integer, [4]: integer}
+--
+---@alias TSCallbackName
+---| 'changedtree'
+---| 'bytes'
+---| 'detach'
+---| 'child_added'
+---| 'child_removed'
+
+---@alias TSCallbackNameOn
+---| 'on_changedtree'
+---| 'on_bytes'
+---| 'on_detach'
+---| 'on_child_added'
+---| 'on_child_removed'
 
+---@class LanguageTree
+---@field private _callbacks table<TSCallbackName,function[]> Callback handlers
+---@field private _children table<string,LanguageTree> Injected languages
+---@field private _injection_query Query Queries defining injected languages
+---@field private _opts table Options
+---@field private _parser TSParser Parser for language
+---@field private _regions Range[][] List of regions this tree should manage and parse
+---@field private _lang string Language name
+---@field private _source (integer|string) Buffer or string to parse
+---@field private _trees TSTree[] Reference to parsed tree (one for each language)
+---@field private _valid boolean If the parsed tree is valid
 local LanguageTree = {}
+
+---@class LanguageTreeOpts
+---@field queries table<string,string>  -- Deprecated
+---@field injections table<string,string>
+
 LanguageTree.__index = LanguageTree
 
 --- A |LanguageTree| holds the treesitter parser for a given language {lang} used
@@ -23,16 +42,17 @@ LanguageTree.__index = LanguageTree
 --- needs to store parsers for these child languages as well (which in turn may contain
 --- child languages themselves, hence the name).
 ---
----@param source (number|string) Buffer or a string of text to parse
+---@param source (integer|string) Buffer or a string of text to parse
 ---@param lang string Root language this tree represents
 ---@param opts (table|nil) Optional keyword arguments:
 ---             - injections table Mapping language to injection query strings.
 ---                                This is useful for overriding the built-in
 ---                                runtime file searching for the injection language
 ---                                query per language.
----@return LanguageTree |LanguageTree| parser object
+---@return LanguageTree parser object
 function LanguageTree.new(source, lang, opts)
   language.require_language(lang)
+  ---@type LanguageTreeOpts
   opts = opts or {}
 
   if opts.queries then
@@ -65,6 +85,7 @@ function LanguageTree.new(source, lang, opts)
 end
 
 --- Invalidates this parser and all its children
+---@param reload boolean|nil
 function LanguageTree:invalidate(reload)
   self._valid = false
 
@@ -73,7 +94,7 @@ function LanguageTree:invalidate(reload)
     self._trees = {}
   end
 
-  for _, child in ipairs(self._children) do
+  for _, child in pairs(self._children) do
     child:invalidate(reload)
   end
 end
@@ -111,8 +132,8 @@ end
 --- This will run the injection query for this language to
 --- determine if any child languages should be created.
 ---
----@return userdata[] Table of parsed |tstree|
----@return table Change list
+---@return TSTree[]
+---@return table|nil Change list
 function LanguageTree:parse()
   if self._valid then
     return self._trees
@@ -146,7 +167,7 @@ function LanguageTree:parse()
   end
 
   local injections_by_lang = self:_get_injections()
-  local seen_langs = {}
+  local seen_langs = {} ---@type table<string,boolean>
 
   for lang, injection_ranges in pairs(injections_by_lang) do
     local has_lang = language.require_language(lang, nil, true)
@@ -188,8 +209,8 @@ end
 
 --- Invokes the callback for each |LanguageTree| and its children recursively
 ---
----@param fn function(tree: LanguageTree, lang: string)
----@param include_self boolean Whether to include the invoking tree in the results
+---@param fn fun(tree: LanguageTree, lang: string)
+---@param include_self boolean|nil Whether to include the invoking tree in the results
 function LanguageTree:for_each_child(fn, include_self)
   if include_self then
     fn(self, self._lang)
@@ -204,7 +225,7 @@ end
 ---
 --- Note: This includes the invoking tree's child trees as well.
 ---
----@param fn function(tree: TSTree, languageTree: LanguageTree)
+---@param fn fun(tree: TSTree, ltree: LanguageTree)
 function LanguageTree:for_each_tree(fn)
   for _, tree in ipairs(self._trees) do
     fn(tree, self)
@@ -221,7 +242,7 @@ end
 ---
 ---@private
 ---@param lang string Language to add.
----@return LanguageTree Injected |LanguageTree|
+---@return LanguageTree injected
 function LanguageTree:add_child(lang)
   if self._children[lang] then
     self:remove_child(lang)
@@ -258,7 +279,7 @@ end
 --- `remove_child` must be called on the parent to remove it.
 function LanguageTree:destroy()
   -- Cleanup here
-  for _, child in ipairs(self._children) do
+  for _, child in pairs(self._children) do
     child:destroy()
   end
 end
@@ -280,20 +301,22 @@ end
 --- Note: This call invalidates the tree and requires it to be parsed again.
 ---
 ---@private
----@param regions table List of regions this tree should manage and parse.
+---@param regions integer[][][] List of regions this tree should manage and parse.
 function LanguageTree:set_included_regions(regions)
   -- Transform the tables from 4 element long to 6 element long (with byte offset)
   for _, region in ipairs(regions) do
     for i, range in ipairs(region) do
       if type(range) == 'table' and #range == 4 then
+        ---@diagnostic disable-next-line:no-unknown
         local start_row, start_col, end_row, end_col = unpack(range)
         local start_byte = 0
         local end_byte = 0
+        local source = self._source
         -- TODO(vigoux): proper byte computation here, and account for EOL ?
-        if type(self._source) == 'number' then
+        if type(source) == 'number' then
           -- Easy case, this is a buffer parser
-          start_byte = a.nvim_buf_get_offset(self._source, start_row) + start_col
-          end_byte = a.nvim_buf_get_offset(self._source, end_row) + end_col
+          start_byte = a.nvim_buf_get_offset(source, start_row) + start_col
+          end_byte = a.nvim_buf_get_offset(source, end_row) + end_col
         elseif type(self._source) == 'string' then
           -- string parser, single `\n` delimited string
           start_byte = vim.fn.byteidx(self._source, start_col)
@@ -320,9 +343,13 @@ function LanguageTree:included_regions()
 end
 
 ---@private
+---@param node TSNode
+---@param id integer
+---@param metadata TSMetadata
+---@return Range
 local function get_range_from_metadata(node, id, metadata)
   if metadata[id] and metadata[id].range then
-    return metadata[id].range
+    return metadata[id].range --[[@as Range]]
   end
   return { node:range() }
 end
@@ -334,11 +361,13 @@ end
 --- TODO: Allow for an offset predicate to tailor the injection range
 ---       instead of using the entire nodes range.
 ---@private
+---@return table<string, integer[][]>
 function LanguageTree:_get_injections()
   if not self._injection_query then
     return {}
   end
 
+  ---@type table<integer,table<string,table<integer,table>>>
   local injections = {}
 
   for tree_index, tree in ipairs(self._trees) do
@@ -348,14 +377,14 @@ function LanguageTree:_get_injections()
     for pattern, match, metadata in
       self._injection_query:iter_matches(root_node, self._source, start_line, end_line + 1)
     do
-      local lang = nil
-      local ranges = {}
-      local combined = metadata.combined
+      local lang = nil ---@type string
+      local ranges = {} ---@type Range[]
+      local combined = metadata.combined ---@type boolean
 
       -- Directives can configure how injections are captured as well as actual node captures.
       -- This allows more advanced processing for determining ranges and language resolution.
       if metadata.content then
-        local content = metadata.content
+        local content = metadata.content ---@type any
 
         -- Allow for captured nodes to be used
         if type(content) == 'number' then
@@ -368,7 +397,7 @@ function LanguageTree:_get_injections()
       end
 
       if metadata.language then
-        lang = metadata.language
+        lang = metadata.language ---@type string
       end
 
       -- You can specify the content and language together
@@ -379,7 +408,7 @@ function LanguageTree:_get_injections()
 
         -- Lang should override any other language tag
         if name == 'language' and not lang then
-          lang = query.get_node_text(node, self._source)
+          lang = query.get_node_text(node, self._source) --[[@as string]]
         elseif name == 'combined' then
           combined = true
         elseif name == 'content' and #ranges == 0 then
@@ -417,6 +446,7 @@ function LanguageTree:_get_injections()
     end
   end
 
+  ---@type table<string,Range[][]>
   local result = {}
 
   -- Generate a map by lang of node lists.
@@ -429,11 +459,13 @@ function LanguageTree:_get_injections()
 
       for _, entry in pairs(patterns) do
         if entry.combined then
+          ---@diagnostic disable-next-line:no-unknown
           local regions = vim.tbl_map(function(e)
             return vim.tbl_flatten(e)
           end, entry.regions)
           table.insert(result[lang], regions)
         else
+          ---@diagnostic disable-next-line:no-unknown
           for _, ranges in ipairs(entry.regions) do
             table.insert(result[lang], ranges)
           end
@@ -446,6 +478,7 @@ function LanguageTree:_get_injections()
 end
 
 ---@private
+---@param cb_name TSCallbackName
 function LanguageTree:_do_callback(cb_name, ...)
   for _, cb in ipairs(self._callbacks[cb_name]) do
     cb(...)
@@ -453,6 +486,17 @@ function LanguageTree:_do_callback(cb_name, ...)
 end
 
 ---@private
+---@param bufnr integer
+---@param changed_tick integer
+---@param start_row integer
+---@param start_col integer
+---@param start_byte integer
+---@param old_row integer
+---@param old_col integer
+---@param old_byte integer
+---@param new_row integer
+---@param new_col integer
+---@param new_byte integer
 function LanguageTree:_on_bytes(
   bufnr,
   changed_tick,
@@ -523,6 +567,7 @@ end
 ---           - `on_child_added` : emitted when a child is added to the tree.
 ---           - `on_child_removed` : emitted when a child is removed from the tree.
 function LanguageTree:register_cbs(cbs)
+  ---@cast cbs table<TSCallbackNameOn,function>
   if not cbs then
     return
   end
@@ -549,6 +594,9 @@ function LanguageTree:register_cbs(cbs)
 end
 
 ---@private
+---@param tree TSTree
+---@param range Range
+---@return boolean
 local function tree_contains(tree, range)
   local start_row, start_col, end_row, end_col = tree:root():range()
   local start_fits = start_row < range[1] or (start_row == range[1] and start_col <= range[2])
@@ -559,7 +607,7 @@ end
 
 --- Determines whether {range} is contained in the |LanguageTree|.
 ---
----@param range table `{ start_line, start_col, end_line, end_col }`
+---@param range Range `{ start_line, start_col, end_line, end_col }`
 ---@return boolean
 function LanguageTree:contains(range)
   for _, tree in pairs(self._trees) do
@@ -573,10 +621,10 @@ end
 
 --- Gets the tree that contains {range}.
 ---
----@param range table `{ start_line, start_col, end_line, end_col }`
+---@param range Range `{ start_line, start_col, end_line, end_col }`
 ---@param opts table|nil Optional keyword arguments:
 ---             - ignore_injections boolean Ignore injected languages (default true)
----@return userdata|nil Contained |tstree|
+---@return TSTree|nil
 function LanguageTree:tree_for_range(range, opts)
   opts = opts or {}
   local ignore = vim.F.if_nil(opts.ignore_injections, true)
@@ -602,10 +650,10 @@ end
 
 --- Gets the smallest named node that contains {range}.
 ---
----@param range table `{ start_line, start_col, end_line, end_col }`
+---@param range Range `{ start_line, start_col, end_line, end_col }`
 ---@param opts table|nil Optional keyword arguments:
 ---             - ignore_injections boolean Ignore injected languages (default true)
----@return userdata|nil Found |tsnode|
+---@return TSNode|nil Found node
 function LanguageTree:named_node_for_range(range, opts)
   local tree = self:tree_for_range(range, opts)
   if tree then
@@ -615,7 +663,7 @@ end
 
 --- Gets the appropriate language that contains {range}.
 ---
----@param range table `{ start_line, start_col, end_line, end_col }`
+---@param range Range `{ start_line, start_col, end_line, end_col }`
 ---@return LanguageTree Managing {range}
 function LanguageTree:language_for_range(range)
   for _, child in pairs(self._children) do
diff --git a/runtime/lua/vim/treesitter/playground.lua b/runtime/lua/vim/treesitter/playground.lua
index bb073290c6..be7764e6f0 100644
--- a/runtime/lua/vim/treesitter/playground.lua
+++ b/runtime/lua/vim/treesitter/playground.lua
@@ -1,12 +1,13 @@
 local api = vim.api
 
-local M = {}
-
----@class Playground
+---@class TSPlayground
 ---@field ns number API namespace
 ---@field opts table Options table with the following keys:
 ---                  - anon (boolean): If true, display anonymous nodes
 ---                  - lang (boolean): If true, display the language alongside each node
+---@field nodes Node[]
+---@field named Node[]
+local TSPlayground = {}
 ---
 ---@class Node
 ---@field id number Node id
@@ -18,6 +19,7 @@ local M = {}
 ---@field end_lnum number Final line number of this node in the source buffer
 ---@field end_col number Final column number of this node in the source buffer
 ---@field lang string Source language of this node
+---@field root TSNode
 
 --- Traverse all child nodes starting at {node}.
 ---
@@ -31,10 +33,10 @@ local M = {}
 --- node of each of these trees is contained within a node in the primary tree. The {injections}
 --- table maps nodes in the primary tree to root nodes of injected trees.
 ---
----@param node userdata Starting node to begin traversal |tsnode|
+---@param node TSNode Starting node to begin traversal |tsnode|
 ---@param depth number Current recursion depth
 ---@param lang string Language of the tree currently being traversed
----@param injections table Mapping of node ids to root nodes of injected language trees (see
+---@param injections table<integer,Node> Mapping of node ids to root nodes of injected language trees (see
 ---                        explanation above)
 ---@param tree Node[] Output table containing a list of tables each representing a node in the tree
 ---@private
@@ -48,7 +50,7 @@ local function traverse(node, depth, lang, injections, tree)
     local type = child:type()
     local lnum, col, end_lnum, end_col = child:range()
     local named = child:named()
-    local text
+    local text ---@type string
     if named then
       if field then
         text = string.format('%s: (%s)', field, type)
@@ -79,14 +81,14 @@ end
 
 --- Create a new Playground object.
 ---
----@param bufnr number Source buffer number
+---@param bufnr integer Source buffer number
 ---@param lang string|nil Language of source buffer
 ---
----@return Playground|nil
+---@return TSPlayground|nil
 ---@return string|nil Error message, if any
 ---
 ---@private
-function M.new(self, bufnr, lang)
+function TSPlayground:new(bufnr, lang)
   local ok, parser = pcall(vim.treesitter.get_parser, bufnr or 0, lang)
   if not ok then
     return nil, 'No parser available for the given buffer'
@@ -96,7 +98,7 @@ function M.new(self, bufnr, lang)
   -- the primary tree that contains that root. Add a mapping from the node in the primary tree to
   -- the root in the child tree to the {injections} table.
   local root = parser:parse()[1]:root()
-  local injections = {}
+  local injections = {} ---@type table<integer,table>
   parser:for_each_child(function(child, lang_)
     child:for_each_tree(function(tree)
       local r = tree:root()
@@ -112,7 +114,7 @@ function M.new(self, bufnr, lang)
 
   local nodes = traverse(root, 0, parser:lang(), injections, {})
 
-  local named = {}
+  local named = {} ---@type Node[]
   for _, v in ipairs(nodes) do
     if v.named then
       named[#named + 1] = v
@@ -138,9 +140,9 @@ end
 ---
 ---@param bufnr number Buffer number to write into.
 ---@private
-function M.draw(self, bufnr)
+function TSPlayground:draw(bufnr)
   vim.bo[bufnr].modifiable = true
-  local lines = {}
+  local lines = {} ---@type string[]
   for _, item in self:iter() do
     lines[#lines + 1] = table.concat({
       string.rep(' ', item.depth),
@@ -168,19 +170,19 @@ end
 ---@param i number Node number to get
 ---@return Node
 ---@private
-function M.get(self, i)
+function TSPlayground:get(i)
   local t = self.opts.anon and self.nodes or self.named
   return t[i]
 end
 
 --- Iterate over all of the nodes in this Playground object.
 ---
----@return function Iterator over all nodes in this Playground
+---@return (fun(): integer, Node) Iterator over all nodes in this Playground
 ---@return table
 ---@return number
 ---@private
-function M.iter(self)
+function TSPlayground:iter()
   return ipairs(self.opts.anon and self.nodes or self.named)
 end
 
-return M
+return TSPlayground
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index dbf134573d..84ed2667b9 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -1,21 +1,25 @@
 local a = vim.api
 local language = require('vim.treesitter.language')
 
--- query: pattern matching on trees
--- predicate matching is implemented in lua
---
 ---@class Query
 ---@field captures string[] List of captures used in query
----@field info table Contains used queries, predicates, directives
+---@field info TSQueryInfo Contains used queries, predicates, directives
 ---@field query userdata Parsed query
 local Query = {}
 Query.__index = Query
 
+---@class TSQueryInfo
+---@field captures table
+---@field patterns table<string,any[][]>
+
 local M = {}
 
 ---@private
+---@param files string[]
+---@return string[]
 local function dedupe_files(files)
   local result = {}
+  ---@type table<string,boolean>
   local seen = {}
 
   for _, path in ipairs(files) do
@@ -65,10 +69,10 @@ function M.get_query_files(lang, query_name, is_included)
     return {}
   end
 
-  local base_query = nil
+  local base_query = nil ---@type string?
   local extensions = {}
 
-  local base_langs = {}
+  local base_langs = {} ---@type string[]
 
   -- Now get the base languages by looking at the first line of every file
   -- The syntax is the following :
@@ -87,6 +91,7 @@ function M.get_query_files(lang, query_name, is_included)
     local extension = false
 
     for modeline in
+      ---@return string
       function()
         return file:read('*l')
       end
@@ -97,6 +102,7 @@ function M.get_query_files(lang, query_name, is_included)
 
       local langlist = modeline:match(MODELINE_FORMAT)
       if langlist then
+        ---@diagnostic disable-next-line:param-type-mismatch
         for _, incllang in ipairs(vim.split(langlist, ',', true)) do
           local is_optional = incllang:match('%(.*%)')
 
@@ -137,6 +143,8 @@ function M.get_query_files(lang, query_name, is_included)
 end
 
 ---@private
+---@param filenames string[]
+---@return string
 local function read_query_files(filenames)
   local contents = {}
 
@@ -147,7 +155,8 @@ local function read_query_files(filenames)
   return table.concat(contents, '')
 end
 
---- The explicitly set queries from |vim.treesitter.query.set_query()|
+-- The explicitly set queries from |vim.treesitter.query.set_query()|
+---@type table<string,table<string,Query>>
 local explicit_queries = setmetatable({}, {
   __index = function(t, k)
     local lang_queries = {}
@@ -174,7 +183,7 @@ end
 ---@param lang string Language to use for the query
 ---@param query_name string Name of the query (e.g. "highlights")
 ---
----@return Query Parsed query
+---@return Query|nil Parsed query
 function M.get_query(lang, query_name)
   if explicit_queries[lang][query_name] then
     return explicit_queries[lang][query_name]
@@ -188,6 +197,7 @@ function M.get_query(lang, query_name)
   end
 end
 
+---@type {[string]: {[string]: Query}}
 local query_cache = vim.defaulttable(function()
   return setmetatable({}, { __mode = 'v' })
 end)
@@ -226,11 +236,11 @@ end
 
 --- Gets the text corresponding to a given node
 ---
----@param node userdata |tsnode|
+---@param node TSNode
 ---@param source (number|string) Buffer or string from which the {node} is extracted
 ---@param opts (table|nil) Optional parameters.
 ---          - concat: (boolean) Concatenate result in a string (default true)
----@return (string[]|string)
+---@return (string[]|string|nil)
 function M.get_node_text(node, source, opts)
   opts = opts or {}
   local concat = vim.F.if_nil(opts.concat, true)
@@ -239,12 +249,12 @@ function M.get_node_text(node, source, opts)
   local end_row, end_col, end_byte = node:end_()
 
   if type(source) == 'number' then
-    local lines
     local eof_row = a.nvim_buf_line_count(source)
     if start_row >= eof_row then
       return nil
     end
 
+    local lines ---@type string[]
     if end_col == 0 then
       lines = a.nvim_buf_get_lines(source, start_row, end_row, true)
       end_col = -1
@@ -267,8 +277,13 @@ function M.get_node_text(node, source, opts)
   end
 end
 
+---@alias TSMatch table<integer,TSNode>
+
+---@alias TSPredicate fun(match: TSMatch, _, _, predicate: any[]): boolean
+
 -- Predicate handler receive the following arguments
 -- (match, pattern, bufnr, predicate)
+---@type table<string,TSPredicate>
 local predicate_handlers = {
   ['eq?'] = function(match, _, source, predicate)
     local node = match[predicate[2]]
@@ -277,13 +292,13 @@ local predicate_handlers = {
     end
     local node_text = M.get_node_text(node, source)
 
-    local str
+    local str ---@type string
     if type(predicate[3]) == 'string' then
       -- (#eq? @aa "foo")
       str = predicate[3]
     else
       -- (#eq? @aa @bb)
-      str = M.get_node_text(match[predicate[3]], source)
+      str = M.get_node_text(match[predicate[3]], source) --[[@as string]]
     end
 
     if node_text ~= str or str == nil then
@@ -299,7 +314,7 @@ local predicate_handlers = {
       return true
     end
     local regex = predicate[3]
-    return string.find(M.get_node_text(node, source), regex)
+    return string.find(M.get_node_text(node, source) --[[@as string]], regex) ~= nil
   end,
 
   ['match?'] = (function()
@@ -321,10 +336,12 @@ local predicate_handlers = {
     })
 
     return function(match, _, source, pred)
+      ---@cast match TSMatch
       local node = match[pred[2]]
       if not node then
         return true
       end
+      ---@diagnostic disable-next-line no-unknown
       local regex = compiled_vim_regexes[pred[3]]
       return regex:match_str(M.get_node_text(node, source))
     end
@@ -335,7 +352,7 @@ local predicate_handlers = {
     if not node then
       return true
     end
-    local node_text = M.get_node_text(node, source)
+    local node_text = M.get_node_text(node, source) --[[@as string]]
 
     for i = 3, #predicate do
       if string.find(node_text, predicate[i], 1, true) then
@@ -359,6 +376,7 @@ local predicate_handlers = {
     if not string_set then
       string_set = {}
       for i = 3, #predicate do
+        ---@diagnostic disable-next-line:no-unknown
         string_set[predicate[i]] = true
       end
       predicate['string_set'] = string_set
@@ -371,21 +389,39 @@ local predicate_handlers = {
 -- As we provide lua-match? also expose vim-match?
 predicate_handlers['vim-match?'] = predicate_handlers['match?']
 
+---@class TSMetadata
+---@field [integer] TSMetadata
+---@field [string] integer|string
+---@field range Range
+
+---@alias TSDirective fun(match: TSMatch, _, _, predicate: any[], metadata: TSMetadata)
+
+-- Predicate handler receive the following arguments
+-- (match, pattern, bufnr, predicate)
+
 -- Directives store metadata or perform side effects against a match.
 -- Directives should always end with a `!`.
 -- Directive handler receive the following arguments
 -- (match, pattern, bufnr, predicate, metadata)
+---@type table<string,TSDirective>
 local directive_handlers = {
   ['set!'] = function(_, _, _, pred, metadata)
     if #pred == 4 then
       -- (#set! @capture "key" "value")
+      ---@diagnostic disable-next-line:no-unknown
       local _, capture_id, key, value = unpack(pred)
+      ---@cast value integer|string
+      ---@cast capture_id integer
+      ---@cast key string
       if not metadata[capture_id] then
         metadata[capture_id] = {}
       end
       metadata[capture_id][key] = value
     else
+      ---@diagnostic disable-next-line:no-unknown
       local _, key, value = unpack(pred)
+      ---@cast value integer|string
+      ---@cast key string
       -- (#set! "key" "value")
       metadata[key] = value
     end
@@ -393,9 +429,11 @@ local directive_handlers = {
   -- Shifts the range of a node.
   -- Example: (#offset! @_node 0 1 0 -1)
   ['offset!'] = function(match, _, _, pred, metadata)
+    ---@cast pred integer[]
     local capture_id = pred[2]
     local offset_node = match[capture_id]
     local range = { offset_node:range() }
+    ---@cast range integer[] bug in sumneko
     local start_row_offset = pred[3] or 0
     local start_col_offset = pred[4] or 0
     local end_row_offset = pred[5] or 0
@@ -419,8 +457,9 @@ local directive_handlers = {
 --- Adds a new predicate to be used in queries
 ---
 ---@param name string Name of the predicate, without leading #
----@param handler function(match:table, pattern:string, bufnr:number, predicate:string[])
+---@param handler function(match:table<string,TSNode>, pattern:string, bufnr:number, predicate:string[])
 ---   - see |vim.treesitter.query.add_directive()| for argument meanings
+---@param force boolean
 function M.add_predicate(name, handler, force)
   if predicate_handlers[name] and not force then
     error(string.format('Overriding %s', name))
@@ -437,12 +476,13 @@ end
 --- metadata table `metadata[capture_id].key = value`
 ---
 ---@param name string Name of the directive, without leading #
----@param handler function(match:table, pattern:string, bufnr:number, predicate:string[], metadata:table)
+---@param handler function(match:table<string,TSNode>, pattern:string, bufnr:number, predicate:string[], metadata:table)
 ---   - match: see |treesitter-query|
 ---      - node-level data are accessible via `match[capture_id]`
 ---   - pattern: see |treesitter-query|
 ---   - predicate: list of strings containing the full directive being called, e.g.
 ---     `(node (#set! conceal "-"))` would get the predicate `{ "#set!", "conceal", "-" }`
+---@param force boolean
 function M.add_directive(name, handler, force)
   if directive_handlers[name] and not force then
     error(string.format('Overriding %s', name))
@@ -474,6 +514,9 @@ local function is_directive(name)
 end
 
 ---@private
+---@param match TSMatch
+---@param pattern string
+---@param source integer|string
 function Query:match_preds(match, pattern, source)
   local preds = self.info.patterns[pattern]
 
@@ -482,8 +525,9 @@ function Query:match_preds(match, pattern, source)
     -- continue on the other case. This way unknown predicates will not be considered,
     -- which allows some testing and easier user extensibility (#12173).
     -- Also, tree-sitter strips the leading # from predicates for us.
-    local pred_name
-    local is_not
+    local pred_name ---@type string
+
+    local is_not ---@type boolean
 
     -- Skip over directives... they will get processed after all the predicates.
     if not is_directive(pred[1]) then
@@ -513,6 +557,8 @@ function Query:match_preds(match, pattern, source)
 end
 
 ---@private
+---@param match TSMatch
+---@param metadata TSMetadata
 function Query:apply_directives(match, pattern, source, metadata)
   local preds = self.info.patterns[pattern]
 
@@ -534,6 +580,10 @@ end
 -- When the node's range is used, the stop is incremented by 1
 -- to make the search inclusive.
 ---@private
+---@param start integer
+---@param stop integer
+---@param node TSNode
+---@return integer, integer
 local function value_or_node_range(start, stop, node)
   if start == nil and stop == nil then
     local node_start, _, node_stop, _ = node:range()
@@ -565,14 +615,12 @@ end
 --- end
 --- </pre>
 ---
----@param node userdata |tsnode| under which the search will occur
----@param source (number|string) Source buffer or string to extract text from
+---@param node TSNode under which the search will occur
+---@param source (integer|string) Source buffer or string to extract text from
 ---@param start number Starting line for the search
 ---@param stop number Stopping line for the search (end-exclusive)
 ---
----@return number capture Matching capture id
----@return table capture_node Capture for {node}
----@return table metadata for the {capture}
+---@return (fun(): integer, TSNode, TSMetadata): capture id, capture node, metadata
 function Query:iter_captures(node, source, start, stop)
   if type(source) == 'number' and source == 0 then
     source = vim.api.nvim_get_current_buf()
@@ -622,14 +670,12 @@ end
 --- end
 --- </pre>
 ---
----@param node userdata |tsnode| under which the search will occur
----@param source (number|string) Source buffer or string to search
----@param start number Starting line for the search
----@param stop number Stopping line for the search (end-exclusive)
+---@param node TSNode under which the search will occur
+---@param source (integer|string) Source buffer or string to search
+---@param start integer Starting line for the search
+---@param stop integer Stopping line for the search (end-exclusive)
 ---
----@return number pattern id
----@return table match
----@return table metadata
+---@return (fun(): integer, table<integer,TSNode>, table): pattern id, match, metadata
 function Query:iter_matches(node, source, start, stop)
   if type(source) == 'number' and source == 0 then
     source = vim.api.nvim_get_current_buf()
@@ -638,6 +684,7 @@ function Query:iter_matches(node, source, start, stop)
   start, stop = value_or_node_range(start, stop, node)
 
   local raw_iter = node:_rawquery(self.query, false, start, stop)
+  ---@cast raw_iter fun(): string, any
   local function iter()
     local pattern, match = raw_iter()
     local metadata = {}