feat(emmy): add emmylua annotations

8 files changed, 257 insertions, 121 deletions

diff --git a/lua/nvim-treesitter/configs.lua b/lua/nvim-treesitter/configs.lua
index 87bb8d1a1..12fbb4cd1 100644
--- a/lua/nvim-treesitter/configs.lua
+++ b/lua/nvim-treesitter/configs.lua
@@ -8,6 +8,16 @@ local caching = require "nvim-treesitter.caching"
 
 local M = {}
 
+---@class TSConfig
+---@field modules {[string]:TSModule}
+---@field sync_install boolean
+---@field ensure_installed string[]
+---@field ignore_install string[]
+---@field auto_install boolean
+---@field update_strategy string
+---@field parser_install_dir string|nil
+
+---@type TSConfig
 local config = {
   modules = {},
   sync_install = false,
@@ -21,6 +31,17 @@ local config = {
 local queued_modules_defs = {}
 -- Whether we've initialized the plugin yet.
 local is_initialized = false
+
+---@class TSModule
+---@field module_path string
+---@field enable boolean|string[]|function(string): boolean
+---@field disable boolean|string[]|function(string): boolean
+---@field is_supported function(string): boolean
+---@field attach function(string)
+---@field detach function(string)
+---@field enabled_buffers table<integer, boolean>
+
+---@type {[string]: TSModule}
 local builtin_modules = {
   highlight = {
     module_path = "nvim-treesitter.highlight",
@@ -54,7 +75,9 @@ local builtin_modules = {
 
 local attached_buffers_by_module = caching.create_buffer_cache()
 
--- Resolves a module by requiring the `module_path` or using the module definition.
+---Resolves a module by requiring the `module_path` or using the module definition.
+---@param mod_name string
+---@return TSModule|nil
 local function resolve_module(mod_name)
   local config_mod = M.get_module(mod_name)
 
@@ -69,10 +92,10 @@ local function resolve_module(mod_name)
   end
 end
 
--- Enables and attaches the module to a buffer for lang.
--- @param mod path to module
--- @param bufnr buffer number, defaults to current buffer
--- @param lang language, defaults to current language
+---Enables and attaches the module to a buffer for lang.
+---@param mod string path to module
+---@param bufnr integer|nil buffer number, defaults to current buffer
+---@param lang string|nil language, defaults to current language
 local function enable_module(mod, bufnr, lang)
   local module = M.get_module(mod)
   if not module then
@@ -93,9 +116,9 @@ local function enable_module(mod, bufnr, lang)
   M.attach_module(mod, bufnr, lang)
 end
 
--- Enables autocomands for the module.
--- After the module is loaded `loaded` will be set to true for the module.
--- @param mod path to module
+---Enables autocomands for the module.
+---After the module is loaded `loaded` will be set to true for the module.
+---@param mod string path to module
 local function enable_mod_conf_autocmd(mod)
   local config_mod = M.get_module(mod)
   if not config_mod or config_mod.loaded then
@@ -113,9 +136,9 @@ local function enable_mod_conf_autocmd(mod)
   config_mod.loaded = true
 end
 
--- Enables the module globally and for all current buffers.
--- After enabled, `enable` will be set to true for the module.
--- @param mod path to module
+---Enables the module globally and for all current buffers.
+---After enabled, `enable` will be set to true for the module.
+---@param mod string path to module
 local function enable_all(mod)
   local config_mod = M.get_module(mod)
   if not config_mod then
@@ -131,9 +154,9 @@ local function enable_all(mod)
   end
 end
 
--- Disables and detaches the module for a buffer.
--- @param mod path to module
--- @param bufnr buffer number, defaults to current buffer
+---Disables and detaches the module for a buffer.
+---@param mod string path to module
+---@param bufnr integer buffer number, defaults to current buffer
 local function disable_module(mod, bufnr)
   local module = M.get_module(mod)
   if not module then
@@ -147,9 +170,9 @@ local function disable_module(mod, bufnr)
   M.detach_module(mod, bufnr)
 end
 
--- Disables autocomands for the module.
--- After the module is unloaded `loaded` will be set to false for the module.
--- @param mod path to module
+---Disables autocomands for the module.
+---After the module is unloaded `loaded` will be set to false for the module.
+---@param mod string path to module
 local function disable_mod_conf_autocmd(mod)
   local config_mod = M.get_module(mod)
   if not config_mod or not config_mod.loaded then
@@ -159,9 +182,9 @@ local function disable_mod_conf_autocmd(mod)
   config_mod.loaded = false
 end
 
--- Disables the module globally and for all current buffers.
--- After disabled, `enable` will be set to false for the module.
--- @param mod path to module
+---Disables the module globally and for all current buffers.
+---After disabled, `enable` will be set to false for the module.
+---@param mod string path to module
 local function disable_all(mod)
   local config_mod = M.get_module(mod)
   if not config_mod then
@@ -177,10 +200,10 @@ local function disable_all(mod)
   end
 end
 
--- Toggles a module for a buffer
--- @param mod path to module
--- @param bufnr buffer number, defaults to current buffer
--- @param lang language, defaults to current language
+---Toggles a module for a buffer
+---@param mod string path to module
+---@param bufnr integer buffer number, defaults to current buffer
+---@param lang string language, defaults to current language
 local function toggle_module(mod, bufnr, lang)
   bufnr = bufnr or api.nvim_get_current_buf()
   lang = lang or parsers.get_buf_lang(bufnr)
@@ -207,10 +230,10 @@ local function toggle_all(mod)
   end
 end
 
--- Recurses through all modules including submodules
--- @param accumulator function called for each module
--- @param root root configuration table to start at
--- @param path prefix path
+---Recurses through all modules including submodules
+---@param accumulator function called for each module
+---@param root {[string]: TSModule} root configuration table to start at
+---@param path string|nil prefix path
 local function recurse_modules(accumulator, root, path)
   root = root or config.modules
 
@@ -225,9 +248,9 @@ local function recurse_modules(accumulator, root, path)
   end
 end
 
--- Shows current configuration of all nvim-treesitter modules
--- @param process_function function used as the `process` parameter
---        for vim.inspect (https://github.com/kikito/inspect.lua#optionsprocess)
+---Shows current configuration of all nvim-treesitter modules
+---@param process_function function used as the `process` parameter
+---       for vim.inspect (https://github.com/kikito/inspect.lua#optionsprocess)
 local function config_info(process_function)
   process_function = process_function
     or function(item, path)
@@ -242,6 +265,8 @@ local function config_info(process_function)
   print(vim.inspect(config, { process = process_function }))
 end
 
+---@param query_group string
+---@param lang string
 function M.edit_query_file(query_group, lang)
   lang = lang or parsers.get_buf_lang()
   local files = ts_query.get_query_files(lang, query_group, true)
@@ -259,6 +284,8 @@ function M.edit_query_file(query_group, lang)
   end
 end
 
+---@param query_group string
+---@param lang string
 function M.edit_query_file_user_after(query_group, lang)
   lang = lang or parsers.get_buf_lang()
   local folder = utils.join_path(vim.fn.stdpath "config", "after", "queries", lang)
@@ -340,9 +367,9 @@ M.commands = {
   },
 }
 
--- @param mod: module (string)
--- @param lang: the language of the buffer (string)
--- @param bufnr: the bufnr (number)
+---@param mod string module
+---@param lang string the language of the buffer
+---@param bufnr integer the bufnr
 function M.is_enabled(mod, lang, bufnr)
   if not parsers.has_parser(lang) then
     return false
@@ -376,8 +403,8 @@ function M.is_enabled(mod, lang, bufnr)
   return true
 end
 
--- Setup call for users to override module configurations.
--- @param user_data module overrides
+---Setup call for users to override module configurations.
+---@param user_data TSConfig module overrides
 function M.setup(user_data)
   config.modules = vim.tbl_deep_extend("force", config.modules, user_data)
   config.ignore_install = user_data.ignore_install or {}
@@ -411,32 +438,33 @@ function M.setup(user_data)
   end, config.modules)
 end
 
--- Defines a table of modules that can be attached/detached to buffers
--- based on language support. A module consist of the following properties:
--- * @enable Whether the modules is enabled. Can be true or false.
--- * @disable A list of languages to disable the module for. Only relevant if enable is true.
--- * @keymaps A list of user mappings for a given module if relevant.
--- * @is_supported A function which, given a ft, will return true if the ft works on the module.
--- * @module_path A string path to a module file using `require`. The exported module must contain
---                an `attach` and `detach` function. This path is not required if `attach` and `detach`
---                functions are provided directly on the module definition.
--- * @attach An attach function that is called for each buffer that the module is enabled for. This is required
---           if a `module_path` is not specified.
--- * @detach A detach function that is called for each buffer that the module is enabled for. This is required
---           if a `module_path` is not specified.
--- Modules are not setup until `init` is invoked by the plugin. This allows modules to be defined in any order
--- and can be loaded lazily.
--- @example
--- require"nvim-treesitter".define_modules {
---   my_cool_module = {
---     attach = function()
---       do_some_cool_setup()
---     end,
---     detach = function()
---       do_some_cool_teardown()
---     end
---   }
--- }
+---Defines a table of modules that can be attached/detached to buffers
+---based on language support. A module consist of the following properties:
+---* @enable Whether the modules is enabled. Can be true or false.
+---* @disable A list of languages to disable the module for. Only relevant if enable is true.
+---* @keymaps A list of user mappings for a given module if relevant.
+---* @is_supported A function which, given a ft, will return true if the ft works on the module.
+---* @module_path A string path to a module file using `require`. The exported module must contain
+---               an `attach` and `detach` function. This path is not required if `attach` and `detach`
+---               functions are provided directly on the module definition.
+---* @attach An attach function that is called for each buffer that the module is enabled for. This is required
+---          if a `module_path` is not specified.
+---* @detach A detach function that is called for each buffer that the module is enabled for. This is required
+---          if a `module_path` is not specified.
+---Modules are not setup until `init` is invoked by the plugin. This allows modules to be defined in any order
+---and can be loaded lazily.
+---@example
+---require"nvim-treesitter".define_modules {
+---  my_cool_module = {
+---    attach = function()
+---      do_some_cool_setup()
+---    end,
+---    detach = function()
+---      do_some_cool_teardown()
+---    end
+---  }
+---}
+---@param mod_defs TSModule[]
 function M.define_modules(mod_defs)
   if not is_initialized then
     table.insert(queued_modules_defs, mod_defs)
@@ -463,10 +491,10 @@ function M.define_modules(mod_defs)
   end
 end
 
--- Attaches a module to a buffer
--- @param mod_name the module name
--- @param bufnr the bufnr
--- @param lang the language of the buffer
+---Attaches a module to a buffer
+---@param mod_name string the module name
+---@param bufnr integer the bufnr
+---@param lang string the language of the buffer
 function M.attach_module(mod_name, bufnr, lang)
   bufnr = bufnr or api.nvim_get_current_buf()
   lang = lang or parsers.get_buf_lang(bufnr)
@@ -478,9 +506,9 @@ function M.attach_module(mod_name, bufnr, lang)
   end
 end
 
--- Detaches a module to a buffer
--- @param mod_name the module name
--- @param bufnr the bufnr
+---Detaches a module to a buffer
+---@param mod_name string the module name
+---@param bufnr integer the bufnr
 function M.detach_module(mod_name, bufnr)
   local resolved_mod = resolve_module(mod_name)
   bufnr = bufnr or api.nvim_get_current_buf()
@@ -491,17 +519,17 @@ function M.detach_module(mod_name, bufnr)
   end
 end
 
--- Same as attach_module, but if the module is already attached, detach it first.
--- @param mod_name the module name
--- @param bufnr the bufnr
--- @param lang the language of the buffer
+---Same as attach_module, but if the module is already attached, detach it first.
+---@param mod_name string the module name
+---@param bufnr integer the bufnr
+---@param lang string the language of the buffer
 function M.reattach_module(mod_name, bufnr, lang)
   M.detach_module(mod_name, bufnr)
   M.attach_module(mod_name, bufnr, lang)
 end
 
--- Gets available modules
--- @param root root table to find modules
+---Gets available modules
+---@param root {[string]:TSModule} table to find modules
 function M.available_modules(root)
   local modules = {}
 
@@ -512,25 +540,26 @@ function M.available_modules(root)
   return modules
 end
 
--- Gets a module config by path
--- @param mod_path path to the module
--- @returns the module or nil
+---Gets a module config by path
+---@param mod_path string path to the module
+---@return TSModule|nil the module or nil
 function M.get_module(mod_path)
   local mod = utils.get_at_path(config.modules, mod_path)
 
   return M.is_module(mod) and mod or nil
 end
 
--- Determines whether the provided table is a module.
--- A module should contain an attach and detach function.
--- @param mod the module table
+---Determines whether the provided table is a module.
+---A module should contain an attach and detach function.
+---@param mod table the module table
+---@return boolean
 function M.is_module(mod)
   return type(mod) == "table"
     and ((type(mod.attach) == "function" and type(mod.detach) == "function") or type(mod.module_path) == "string")
 end
 
--- Initializes built-in modules and any queued modules
--- registered by plugins or the user.
+---Initializes built-in modules and any queued modules
+---registered by plugins or the user.
 function M.init()
   is_initialized = true
   M.define_modules(builtin_modules)
@@ -540,11 +569,13 @@ function M.init()
   end
 end
 
--- If parser_install_dir is not nil is used or created.
--- If parser_install_dir is nil try the package dir of the nvim-treesitter
--- plugin first, followed by the "site" dir from "runtimepath". "site" dir will
--- be created if it doesn't exist. Using only the package dir won't work when
--- the plugin is installed with Nix, since the "/nix/store" is read-only.
+---If parser_install_dir is not nil is used or created.
+---If parser_install_dir is nil try the package dir of the nvim-treesitter
+---plugin first, followed by the "site" dir from "runtimepath". "site" dir will
+---be created if it doesn't exist. Using only the package dir won't work when
+---the plugin is installed with Nix, since the "/nix/store" is read-only.
+---@param folder_name string
+---@return string|nil, string|nil
 function M.get_parser_install_dir(folder_name)
   folder_name = folder_name or "parser"
 
diff --git a/lua/nvim-treesitter/fold.lua b/lua/nvim-treesitter/fold.lua
index 1f417414f..50272622f 100644
--- a/lua/nvim-treesitter/fold.lua
+++ b/lua/nvim-treesitter/fold.lua
@@ -102,6 +102,8 @@ local folds_levels = tsutils.memoize_by_buf_tick(function(bufnr)
   return levels
 end)
 
+---@param lnum integer
+---@return string
 function M.get_fold_indic(lnum)
   if not parsers.has_parser() or not lnum then
     return "0"
diff --git a/lua/nvim-treesitter/highlight.lua b/lua/nvim-treesitter/highlight.lua
index 502914fe4..5f7dead61 100644
--- a/lua/nvim-treesitter/highlight.lua
+++ b/lua/nvim-treesitter/highlight.lua
@@ -6,6 +6,9 @@ local configs = require "nvim-treesitter.configs"
 
 local M = {}
 
+---@param config table
+---@param lang string
+---@return boolean
 local function should_enable_vim_regex(config, lang)
   local additional_hl = config.additional_vim_regex_highlighting
   local is_table = type(additional_hl) == "table"
@@ -13,29 +16,36 @@ local function should_enable_vim_regex(config, lang)
   return additional_hl and (not is_table or vim.tbl_contains(additional_hl, lang))
 end
 
+---@param bufnr integer
 local function enable_syntax(bufnr)
   api.nvim_buf_set_option(bufnr, "syntax", "ON")
 end
 
+---@param bufnr integer
 function M.stop(bufnr)
   if ts.highlighter.active[bufnr] then
     ts.highlighter.active[bufnr]:destroy()
   end
 end
 
+---@param bufnr integer
+---@param lang string
 function M.start(bufnr, lang)
   local parser = parsers.get_parser(bufnr, lang)
   ts.highlighter.new(parser, {})
 end
 
+---@param bufnr integer
+---@param lang string
 function M.attach(bufnr, lang)
   local config = configs.get_module "highlight"
   M.start(bufnr, lang)
-  if should_enable_vim_regex(config, lang) then
+  if config and should_enable_vim_regex(config, lang) then
     enable_syntax(bufnr)
   end
 end
 
+---@param bufnr integer
 function M.detach(bufnr)
   M.stop(bufnr)
   enable_syntax(bufnr)
diff --git a/lua/nvim-treesitter/install.lua b/lua/nvim-treesitter/install.lua
index 95a8a9f92..48810b1e2 100644
--- a/lua/nvim-treesitter/install.lua
+++ b/lua/nvim-treesitter/install.lua
@@ -42,11 +42,14 @@ local function get_job_status()
     .. "]"
 end
 
+---@param lang string
+---@param validate boolean|nil
+---@return InstallInfo
 local function get_parser_install_info(lang, validate)
   local parser_config = parsers.get_parser_configs()[lang]
 
   if not parser_config then
-    return error('Parser not available for language "' .. lang .. '"')
+    error('Parser not available for language "' .. lang .. '"')
   end
 
   local install_info = parser_config.install_info
@@ -79,6 +82,8 @@ local function get_revision(lang)
   return (lockfile[lang] and lockfile[lang].revision)
 end
 
+---@param lang string
+---@return string|nil
 local function get_installed_revision(lang)
   local lang_file = utils.join_path(configs.get_parser_info_dir(), lang .. ".revision")
   if vim.fn.filereadable(lang_file) == 1 then
@@ -86,23 +91,30 @@ local function get_installed_revision(lang)
   end
 end
 
+---@param lang string
+---@return boolean
 local function is_installed(lang)
   return #api.nvim_get_runtime_file("parser/" .. lang .. ".so", false) > 0
 end
 
+---@param lang string
+---@return boolean
 local function needs_update(lang)
   local revision = get_revision(lang)
   return not revision or revision ~= get_installed_revision(lang)
 end
 
+---@return table
 local function outdated_parsers()
   return vim.tbl_filter(function(lang)
     return needs_update(lang)
   end, info.installed_parsers())
 end
 
+---@param handle userdata
+---@param is_stderr boolean
 local function onread(handle, is_stderr)
-  return function(err, data)
+  return function(_, data)
     if data then
       if is_stderr then
         complete_error_output[handle] = (complete_error_output[handle] or "") .. data
@@ -147,6 +159,7 @@ function M.iter_cmd(cmd_list, i, lang, success_message)
     local stdout = luv.new_pipe(false)
     local stderr = luv.new_pipe(false)
     attr.opts.stdio = { nil, stdout, stderr }
+    ---@type userdata
     handle = luv.spawn(
       attr.cmd,
       attr.opts,
@@ -225,6 +238,12 @@ local function iter_cmd_sync(cmd_list)
   return true
 end
 
+---@param cache_folder string
+---@param install_folder string
+---@param lang string
+---@param repo InstallInfo
+---@param with_sync boolean
+---@param generate_from_grammar boolean
 local function run_install(cache_folder, install_folder, lang, repo, with_sync, generate_from_grammar)
   parsers.reset_cache()
 
@@ -361,6 +380,12 @@ local function run_install(cache_folder, install_folder, lang, repo, with_sync,
   end
 end
 
+---@param lang string
+---@param ask_reinstall boolean
+---@param cache_folder string
+---@param install_folder string
+---@param with_sync boolean
+---@param generate_from_grammar boolean
 local function install_lang(lang, ask_reinstall, cache_folder, install_folder, with_sync, generate_from_grammar)
   if is_installed(lang) and ask_reinstall ~= "force" then
     if not ask_reinstall then
@@ -389,6 +414,7 @@ local function install_lang(lang, ask_reinstall, cache_folder, install_folder, w
   run_install(cache_folder, install_folder, lang, install_info, with_sync, generate_from_grammar)
 end
 
+---@return function
 local function install(options)
   options = options or {}
   local with_sync = options.with_sync
@@ -405,11 +431,14 @@ local function install(options)
     if err then
       return api.nvim_err_writeln(err)
     end
+    assert(cache_folder)
 
-    local install_folder, err = configs.get_parser_install_dir()
+    local install_folder
+    install_folder, err = configs.get_parser_install_dir()
     if err then
       return api.nvim_err_writeln(err)
     end
+    assert(install_folder)
 
     local languages
     local ask
diff --git a/lua/nvim-treesitter/parsers.lua b/lua/nvim-treesitter/parsers.lua
index ee79c09d6..8fe4225e0 100644
--- a/lua/nvim-treesitter/parsers.lua
+++ b/lua/nvim-treesitter/parsers.lua
@@ -20,6 +20,22 @@ local filetype_to_parsername = {
   cs = "c_sharp",
 }
 
+---@class InstallInfo
+---@field url string
+---@field branch string|nil
+---@field revision string|nil
+---@field files string[]
+---@field generate_requires_npm boolean|nil
+---@field requires_generate_from_grammar boolean|nil
+---@field location string|nil
+
+---@class ParserInfo
+---@field install_info InstallInfo
+---@field filetype string
+---@field maintainers string[]
+---@field experimental boolean|nil
+
+---@type ParserInfo[]
 local list = setmetatable({}, {
   __newindex = function(table, parsername, parserconfig)
     rawset(
diff --git a/lua/nvim-treesitter/query.lua b/lua/nvim-treesitter/query.lua
index aeaa683be..bc80d51e6 100644
--- a/lua/nvim-treesitter/query.lua
+++ b/lua/nvim-treesitter/query.lua
@@ -11,8 +11,10 @@ local EMPTY_ITER = function() end
 
 M.built_in_query_groups = { "highlights", "locals", "folds", "indents", "injections" }
 
--- Creates a function that checks whether a given query exists
--- for a specific language.
+--- Creates a function that checks whether a given query exists
+--- for a specific language.
+---@param query string
+---@return function(string): boolean
 local function get_query_guard(query)
   return function(lang)
     return M.has_query_files(lang, query)
@@ -23,6 +25,7 @@ for _, query in ipairs(M.built_in_query_groups) do
   M["has_" .. query] = get_query_guard(query)
 end
 
+---@return string[]
 function M.available_query_groups()
   local query_files = api.nvim_get_runtime_file("queries/*/*.scm", true)
   local groups = {}
@@ -57,11 +60,19 @@ do
   end
 end
 
+---@param lang string
+---@param query_name string
+---@return string[]
 local function runtime_queries(lang, query_name)
   return api.nvim_get_runtime_file(string.format("queries/%s/%s.scm", lang, query_name), true) or {}
 end
 
+---@type table<string, table<string, boolean>>
 local query_files_cache = {}
+
+---@param lang string
+---@param query_name string
+---@return boolean
 function M.has_query_files(lang, query_name)
   if not query_files_cache[lang] then
     query_files_cache[lang] = {}
@@ -86,6 +97,8 @@ do
   local cache = setmetatable({}, mt)
 
   --- Same as `vim.treesitter.query` except will return cached values
+  ---@param lang string
+  ---@param query_name string
   function M.get_query(lang, query_name)
     if cache[lang][query_name] == nil then
       cache[lang][query_name] = tsq.get_query(lang, query_name)
@@ -98,6 +111,8 @@ do
   --- If lang and query_name is both present, will reload for only the lang and query_name.
   --- If only lang is present, will reload all query_names for that lang
   --- If none are present, will reload everything
+  ---@param lang string
+  ---@param query_name string
   function M.invalidate_query_cache(lang, query_name)
     if lang and query_name then
       cache[lang][query_name] = nil
@@ -106,14 +121,14 @@ do
       end
     elseif lang and not query_name then
       query_files_cache[lang] = nil
-      for query_name, _ in pairs(cache[lang]) do
-        M.invalidate_query_cache(lang, query_name)
+      for query_name0, _ in pairs(cache[lang]) do
+        M.invalidate_query_cache(lang, query_name0)
       end
     elseif not lang and not query_name then
       query_files_cache = {}
-      for lang, _ in pairs(cache) do
-        for query_name, _ in pairs(cache[lang]) do
-          M.invalidate_query_cache(lang, query_name)
+      for lang0, _ in pairs(cache) do
+        for query_name0, _ in pairs(cache[lang0]) do
+          M.invalidate_query_cache(lang0, query_name0)
         end
       end
     else
@@ -123,11 +138,23 @@ do
 end
 
 --- This function is meant for an autocommand and not to be used. Only use if file is a query file.
+---@param fname string
 function M.invalidate_query_file(fname)
   local fnamemodify = vim.fn.fnamemodify
   M.invalidate_query_cache(fnamemodify(fname, ":p:h:t"), fnamemodify(fname, ":t:r"))
 end
 
+---@class QueryInfo
+---@field root LanguageTree
+---@field source integer
+---@field start integer
+---@field stop integer
+
+---@param bufnr integer
+---@param query_name string
+---@param root LanguageTree
+---@param root_lang string|nil
+---@return Query|nil, QueryInfo|nil
 local function prepare_query(bufnr, query_name, root, root_lang)
   local buf_lang = parsers.get_buf_lang(bufnr)
 
@@ -181,6 +208,10 @@ local function prepare_query(bufnr, query_name, root, root_lang)
     }
 end
 
+---@param query Query
+---@param bufnr integer
+---@param start_row integer
+---@param end_row integer
 function M.iter_prepared_matches(query, qnode, bufnr, start_row, end_row)
   -- A function that splits  a string on '.'
   local function split(string)
@@ -249,15 +280,16 @@ function M.iter_prepared_matches(query, qnode, bufnr, start_row, end_row)
 end
 
 --- Return all nodes corresponding to a specific capture path (like @definition.var, @reference.type)
--- Works like M.get_references or M.get_scopes except you can choose the capture
--- Can also be a nested capture like @definition.function to get all nodes defining a function.
---
--- @param bufnr the buffer
--- @param captures a single string or a list of strings
--- @param query_group the name of query group (highlights or injections for example)
--- @param root (optional) node from where to start the search
--- @param lang (optional) the language from where to get the captures.
---             Root nodes can have several languages.
+---Works like M.get_references or M.get_scopes except you can choose the capture
+---Can also be a nested capture like @definition.function to get all nodes defining a function.
+---
+---@param bufnr integer the buffer
+---@param captures string|string[]
+---@param query_group string the name of query group (highlights or injections for example)
+---@param root LanguageTree|nil node from where to start the search
+---@param lang string|nil the language from where to get the captures.
+---            Root nodes can have several languages.
+---@return table|nil
 function M.get_capture_matches(bufnr, captures, query_group, root, lang)
   if type(captures) == "string" then
     captures = { captures }
@@ -289,6 +321,7 @@ function M.iter_captures(bufnr, query_name, root, lang)
   if not query then
     return EMPTY_ITER
   end
+  assert(params)
 
   local iter = query:iter_captures(params.root, params.source, params.start, params.stop)
 
@@ -336,16 +369,17 @@ function M.find_best_match(bufnr, capture_string, query_group, filter_predicate,
   return best
 end
 
--- Iterates matches from a query file.
--- @param bufnr the buffer
--- @param query_group the query file to use
--- @param root the root node
--- @param root the root node lang, if known
+---Iterates matches from a query file.
+---@param bufnr integer the buffer
+---@param query_group string the query file to use
+---@param root LanguageTree the root node
+---@param root_lang string|nil the root node lang, if known
 function M.iter_group_results(bufnr, query_group, root, root_lang)
   local query, params = prepare_query(bufnr, query_group, root, root_lang)
   if not query then
     return EMPTY_ITER
   end
+  assert(params)
 
   return M.iter_prepared_matches(query, params.root, params.source, params.start, params.stop)
 end
@@ -360,18 +394,25 @@ function M.collect_group_results(bufnr, query_group, root, lang)
   return matches
 end
 
+---@alias CaptureResFn function(string, LanguageTree, LanguageTree): string, string
+
 --- Same as get_capture_matches except this will recursively get matches for every language in the tree.
--- @param bufnr The bufnr
--- @param capture_or_fn The capture to get. If a function is provided then that
---                      function will be used to resolve both the capture and query argument.
---                      The function can return `nil` to ignore that tree.
--- @param query_type The query to get the capture from. This is ignore if a function is provided
---                   for the captuer argument.
+---@param bufnr integer The bufnr
+---@param capture_or_fn string|CaptureResFn The capture to get. If a function is provided then that
+---                     function will be used to resolve both the capture and query argument.
+---                     The function can return `nil` to ignore that tree.
+---@param query_type string The query to get the capture from. This is ignore if a function is provided
+---                  for the captuer argument.
 function M.get_capture_matches_recursively(bufnr, capture_or_fn, query_type)
-  local type_fn = type(capture_or_fn) == "function" and capture_or_fn
-    or function()
+  ---@type CaptureResFn
+  local type_fn
+  if type(capture_or_fn) == "function" then
+    type_fn = capture_or_fn
+  else
+    type_fn = function(_, _, _)
       return capture_or_fn, query_type
     end
+  end
   local parser = parsers.get_parser(bufnr)
   local matches = {}
 
diff --git a/lua/nvim-treesitter/shell_command_selectors.lua b/lua/nvim-treesitter/shell_command_selectors.lua
index 522ac727a..59bbb0e41 100644
--- a/lua/nvim-treesitter/shell_command_selectors.lua
+++ b/lua/nvim-treesitter/shell_command_selectors.lua
@@ -49,6 +49,7 @@ function M.select_rm_file_cmd(file, info_msg)
   end
 end
 
+---@return string|nil
 function M.select_executable(executables)
   return vim.tbl_filter(function(c)
     return c ~= vim.NIL and fn.executable(c) == 1
diff --git a/lua/nvim-treesitter/utils.lua b/lua/nvim-treesitter/utils.lua
index 81ce70173..156ba9c31 100644
--- a/lua/nvim-treesitter/utils.lua
+++ b/lua/nvim-treesitter/utils.lua
@@ -11,6 +11,7 @@ function M.notify(msg, log_level, opts)
 end
 
 -- Returns the system specific path seperator.
+---@return string
 function M.get_path_sep()
   return fn.has "win32" == 1 and "\\" or "/"
 end
@@ -82,6 +83,10 @@ function M.setup_commands(mod, commands)
   end
 end
 
+---@param dir string
+---@param create_err string
+---@param writeable_err string
+---@return string|nil, string|nil
 function M.create_or_reuse_writable_dir(dir, create_err, writeable_err)
   create_err = create_err or M.join_space("Could not create dir '", dir, "': ")
   writeable_err = writeable_err or M.join_space("Invalid rights, '", dir, "' should be read/write")
@@ -201,6 +206,7 @@ function M.to_func(a)
   return type(a) == "function" and a or M.constant(a)
 end
 
+---@return string|nil
 function M.ts_cli_version()
   if fn.executable "tree-sitter" == 1 then
     local handle = io.popen "tree-sitter  -V"