Merge remote-tracking branch 'upstream/master' into floattitlefloattitle

5 files changed, 266 insertions, 226 deletions

diff --git a/runtime/lua/vim/treesitter/health.lua b/runtime/lua/vim/treesitter/health.lua
index 3bd59ca282..4995c80a02 100644
--- a/runtime/lua/vim/treesitter/health.lua
+++ b/runtime/lua/vim/treesitter/health.lua
@@ -3,7 +3,7 @@ local ts = vim.treesitter
 
 --- Lists the parsers currently installed
 ---
----@return A list of parsers
+---@return string[] list of parser files
 function M.list_parsers()
   return vim.api.nvim_get_runtime_file('parser/*', true)
 end
diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index e27a5fa9c3..83a26aff13 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -2,6 +2,7 @@ local a = vim.api
 local query = require('vim.treesitter.query')
 
 -- support reload for quick experimentation
+---@class TSHighlighter
 local TSHighlighter = rawget(vim.treesitter, 'TSHighlighter') or {}
 TSHighlighter.__index = TSHighlighter
 
@@ -12,105 +13,18 @@ TSHighlighterQuery.__index = TSHighlighterQuery
 
 local ns = a.nvim_create_namespace('treesitter/highlighter')
 
-local _default_highlights = {}
-local _link_default_highlight_once = function(from, to)
-  if not _default_highlights[from] then
-    _default_highlights[from] = true
-    a.nvim_set_hl(0, from, { link = to, default = true })
-  end
-
-  return from
-end
-
--- If @definition.special does not exist use @definition instead
-local subcapture_fallback = {
-  __index = function(self, capture)
-    local rtn
-    local shortened = capture
-    while not rtn and shortened do
-      shortened = shortened:match('(.*)%.')
-      rtn = shortened and rawget(self, shortened)
-    end
-    rawset(self, capture, rtn or '__notfound')
-    return rtn
-  end,
-}
-
-TSHighlighter.hl_map = setmetatable({
-  ['error'] = 'Error',
-  ['text.underline'] = 'Underlined',
-  ['todo'] = 'Todo',
-  ['debug'] = 'Debug',
-
-  -- Miscs
-  ['comment'] = 'Comment',
-  ['punctuation.delimiter'] = 'Delimiter',
-  ['punctuation.bracket'] = 'Delimiter',
-  ['punctuation.special'] = 'Delimiter',
-
-  -- Constants
-  ['constant'] = 'Constant',
-  ['constant.builtin'] = 'Special',
-  ['constant.macro'] = 'Define',
-  ['define'] = 'Define',
-  ['macro'] = 'Macro',
-  ['string'] = 'String',
-  ['string.regex'] = 'String',
-  ['string.escape'] = 'SpecialChar',
-  ['character'] = 'Character',
-  ['character.special'] = 'SpecialChar',
-  ['number'] = 'Number',
-  ['boolean'] = 'Boolean',
-  ['float'] = 'Float',
-
-  -- Functions
-  ['function'] = 'Function',
-  ['function.special'] = 'Function',
-  ['function.builtin'] = 'Special',
-  ['function.macro'] = 'Macro',
-  ['parameter'] = 'Identifier',
-  ['method'] = 'Function',
-  ['field'] = 'Identifier',
-  ['property'] = 'Identifier',
-  ['constructor'] = 'Special',
-
-  -- Keywords
-  ['conditional'] = 'Conditional',
-  ['repeat'] = 'Repeat',
-  ['label'] = 'Label',
-  ['operator'] = 'Operator',
-  ['keyword'] = 'Keyword',
-  ['exception'] = 'Exception',
-
-  ['type'] = 'Type',
-  ['type.builtin'] = 'Type',
-  ['type.qualifier'] = 'Type',
-  ['type.definition'] = 'Typedef',
-  ['storageclass'] = 'StorageClass',
-  ['structure'] = 'Structure',
-  ['include'] = 'Include',
-  ['preproc'] = 'PreProc',
-}, subcapture_fallback)
-
----@private
-local function is_highlight_name(capture_name)
-  local firstc = string.sub(capture_name, 1, 1)
-  return firstc ~= string.lower(firstc)
-end
-
 ---@private
 function TSHighlighterQuery.new(lang, query_string)
   local self = setmetatable({}, { __index = TSHighlighterQuery })
 
   self.hl_cache = setmetatable({}, {
     __index = function(table, capture)
-      local hl, is_vim_highlight = self:_get_hl_from_capture(capture)
-      if not is_vim_highlight then
-        hl = _link_default_highlight_once(lang .. hl, hl)
+      local name = self._query.captures[capture]
+      local id = 0
+      if not vim.startswith(name, '_') then
+        id = a.nvim_get_hl_id_by_name('@' .. name .. '.' .. lang)
       end
 
-      local id = a.nvim_get_hl_id_by_name(hl)
-
       rawset(table, capture, id)
       return id
     end,
@@ -130,25 +44,12 @@ function TSHighlighterQuery:query()
   return self._query
 end
 
----@private
---- Get the hl from capture.
---- Returns a tuple { highlight_name: string, is_builtin: bool }
-function TSHighlighterQuery:_get_hl_from_capture(capture)
-  local name = self._query.captures[capture]
-
-  if is_highlight_name(name) then
-    -- From "Normal.left" only keep "Normal"
-    return vim.split(name, '.', true)[1], true
-  else
-    return TSHighlighter.hl_map[name] or 0, false
-  end
-end
-
 --- Creates a new highlighter using @param tree
 ---
----@param tree The language tree to use for highlighting
----@param opts Table used to configure the highlighter
----           - queries: Table to overwrite queries used by the highlighter
+---@param tree LanguageTree |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
 function TSHighlighter.new(tree, opts)
   local self = setmetatable({}, TSHighlighter)
 
@@ -187,7 +88,7 @@ function TSHighlighter.new(tree, opts)
     end
   end
 
-  a.nvim_buf_set_option(self.bufnr, 'syntax', '')
+  vim.bo[self.bufnr].syntax = ''
 
   TSHighlighter.active[self.bufnr] = self
 
@@ -196,9 +97,13 @@ function TSHighlighter.new(tree, opts)
   -- syntax FileType autocmds. Later on we should integrate with the
   -- `:syntax` and `set syntax=...` machinery properly.
   if vim.g.syntax_on ~= 1 then
-    vim.api.nvim_command('runtime! syntax/synload.vim')
+    vim.cmd.runtime({ 'syntax/synload.vim', bang = true })
   end
 
+  a.nvim_buf_call(self.bufnr, function()
+    vim.opt_local.spelloptions:append('noplainbuffer')
+  end)
+
   self.tree:parse()
 
   return self
@@ -246,8 +151,10 @@ function TSHighlighter:on_changedtree(changes)
 end
 
 --- Gets the query used for @param lang
----
----@param lang A language used by the highlighter.
+--
+---@private
+---@param lang string Language used by the highlighter.
+---@return Query
 function TSHighlighter:get_query(lang)
   if not self._queries[lang] then
     self._queries[lang] = TSHighlighterQuery.new(lang)
@@ -257,7 +164,7 @@ function TSHighlighter:get_query(lang)
 end
 
 ---@private
-local function on_line_impl(self, buf, line)
+local function on_line_impl(self, buf, line, spell)
   self.tree:for_each_tree(function(tstree, tree)
     if not tstree then
       return
@@ -294,7 +201,9 @@ local function on_line_impl(self, buf, line)
       local start_row, start_col, end_row, end_col = node:range()
       local hl = highlighter_query.hl_cache[capture]
 
-      if hl and end_row >= line then
+      local is_spell = highlighter_query:query().captures[capture] == 'spell'
+
+      if hl and end_row >= line and (not spell or is_spell) then
         a.nvim_buf_set_extmark(buf, ns, start_row, start_col, {
           end_line = end_row,
           end_col = end_col,
@@ -302,6 +211,7 @@ local function on_line_impl(self, buf, line)
           ephemeral = true,
           priority = tonumber(metadata.priority) or 100, -- Low but leaves room below
           conceal = metadata.conceal,
+          spell = is_spell,
         })
       end
       if start_row > line then
@@ -318,7 +228,21 @@ function TSHighlighter._on_line(_, _win, buf, line, _)
     return
   end
 
-  on_line_impl(self, buf, line)
+  on_line_impl(self, buf, line, false)
+end
+
+---@private
+function TSHighlighter._on_spell_nav(_, _, buf, srow, _, erow, _)
+  local self = TSHighlighter.active[buf]
+  if not self then
+    return
+  end
+
+  self:reset_highlight_state()
+
+  for row = srow, erow do
+    on_line_impl(self, buf, row, true)
+  end
 end
 
 ---@private
@@ -345,6 +269,7 @@ a.nvim_set_decoration_provider(ns, {
   on_buf = TSHighlighter._on_buf,
   on_win = TSHighlighter._on_win,
   on_line = TSHighlighter._on_line,
+  _on_spell_nav = TSHighlighter._on_spell_nav,
 })
 
 return TSHighlighter
diff --git a/runtime/lua/vim/treesitter/language.lua b/runtime/lua/vim/treesitter/language.lua
index dfb6f5be84..c92d63b8c4 100644
--- a/runtime/lua/vim/treesitter/language.lua
+++ b/runtime/lua/vim/treesitter/language.lua
@@ -2,14 +2,16 @@ local a = vim.api
 
 local M = {}
 
---- Asserts that the provided language is installed, and optionally provide a path for the parser
+--- Asserts that a parser for the language {lang} is installed.
 ---
---- Parsers are searched in the `parser` runtime directory.
+--- Parsers are searched in the `parser` runtime directory, or the provided {path}
 ---
----@param lang The language the parser should parse
----@param path Optional path the parser is located at
----@param silent Don't throw an error if language not found
-function M.require_language(lang, path, silent)
+---@param lang string Language the parser should parse
+---@param path (string|nil) Optional path the parser is located at
+---@param silent (boolean|nil) Don't throw an error if language not found
+---@param symbol_name (string|nil) Internal symbol name for the language to load
+---@return boolean If the specified language is installed
+function M.require_language(lang, path, silent, symbol_name)
   if vim._ts_has_language(lang) then
     return true
   end
@@ -21,7 +23,6 @@ function M.require_language(lang, path, silent)
         return false
       end
 
-      -- TODO(bfredl): help tag?
       error("no parser for '" .. lang .. "' language, see :help treesitter-parsers")
     end
     path = paths[1]
@@ -29,10 +30,10 @@ function M.require_language(lang, path, silent)
 
   if silent then
     return pcall(function()
-      vim._ts_add_language(path, lang)
+      vim._ts_add_language(path, lang, symbol_name)
     end)
   else
-    vim._ts_add_language(path, lang)
+    vim._ts_add_language(path, lang, symbol_name)
   end
 
   return true
@@ -42,7 +43,8 @@ end
 ---
 --- Inspecting provides some useful information on the language like node names, ...
 ---
----@param lang The language.
+---@param lang string Language
+---@return table
 function M.inspect_language(lang)
   M.require_language(lang)
   return vim._ts_inspect_language(lang)
diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index 4d3b0631a2..e9d70c4204 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -2,19 +2,35 @@ 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
+
 local LanguageTree = {}
 LanguageTree.__index = LanguageTree
 
---- Represents a single treesitter parser for a language.
---- The language can contain child languages with in its range,
---- hence the tree.
+--- A |LanguageTree| holds the treesitter parser for a given language {lang} used
+--- to parse a buffer. As the buffer may contain injected languages, the LanguageTree
+--- needs to store parsers for these child languages as well (which in turn may contain
+--- child languages themselves, hence the name).
 ---
----@param source Can be a bufnr or a string of text to parse
----@param lang The language this tree represents
----@param opts Options table
----@param opts.injections A table of language to injection query strings.
----                      This is useful for overriding the built-in runtime file
----                      searching for the injection language query per language.
+---@param source (number|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
 function LanguageTree.new(source, lang, opts)
   language.require_language(lang)
   opts = opts or {}
@@ -94,6 +110,9 @@ end
 --- for the language this tree represents.
 --- 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
 function LanguageTree:parse()
   if self._valid then
     return self._trees
@@ -167,10 +186,10 @@ function LanguageTree:parse()
   return self._trees, changes
 end
 
---- Invokes the callback for each LanguageTree and it's children recursively
+--- Invokes the callback for each |LanguageTree| and its children recursively
 ---
----@param fn The function to invoke. This is invoked with arguments (tree: LanguageTree, lang: string)
----@param include_self Whether to include the invoking tree in the results.
+---@param fn function(tree: LanguageTree, lang: string)
+---@param include_self boolean 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)
@@ -181,12 +200,11 @@ function LanguageTree:for_each_child(fn, include_self)
   end
 end
 
---- Invokes the callback for each treesitter trees recursively.
+--- Invokes the callback for each |LanguageTree| recursively.
 ---
---- Note, this includes the invoking language tree's trees as well.
+--- Note: This includes the invoking tree's child trees as well.
 ---
----@param fn The callback to invoke. The callback is invoked with arguments
----         (tree: TSTree, languageTree: LanguageTree)
+---@param fn function(tree: TSTree, languageTree: LanguageTree)
 function LanguageTree:for_each_tree(fn)
   for _, tree in ipairs(self._trees) do
     fn(tree, self)
@@ -197,11 +215,13 @@ function LanguageTree:for_each_tree(fn)
   end
 end
 
---- Adds a child language to this tree.
+--- Adds a child language to this |LanguageTree|.
 ---
 --- If the language already exists as a child, it will first be removed.
 ---
----@param lang The language to add.
+---@private
+---@param lang string Language to add.
+---@return LanguageTree Injected |LanguageTree|
 function LanguageTree:add_child(lang)
   if self._children[lang] then
     self:remove_child(lang)
@@ -215,9 +235,10 @@ function LanguageTree:add_child(lang)
   return self._children[lang]
 end
 
---- Removes a child language from this tree.
+--- Removes a child language from this |LanguageTree|.
 ---
----@param lang The language to remove.
+---@private
+---@param lang string Language to remove.
 function LanguageTree:remove_child(lang)
   local child = self._children[lang]
 
@@ -229,12 +250,11 @@ function LanguageTree:remove_child(lang)
   end
 end
 
---- Destroys this language tree and all its children.
+--- Destroys this |LanguageTree| and all its children.
 ---
 --- Any cleanup logic should be performed here.
 ---
---- Note:
---- This DOES NOT remove this tree from a parent. Instead,
+--- Note: This DOES NOT remove this tree from a parent. Instead,
 --- `remove_child` must be called on the parent to remove it.
 function LanguageTree:destroy()
   -- Cleanup here
@@ -243,23 +263,24 @@ function LanguageTree:destroy()
   end
 end
 
---- Sets the included regions that should be parsed by this parser.
+--- Sets the included regions that should be parsed by this |LanguageTree|.
 --- A region is a set of nodes and/or ranges that will be parsed in the same context.
 ---
---- For example, `{ { node1 }, { node2} }` is two separate regions.
---- This will be parsed by the parser in two different contexts... thus resulting
+--- For example, `{ { node1 }, { node2} }` contains two separate regions.
+--- They will be parsed by the parser in two different contexts, thus resulting
 --- in two separate trees.
 ---
---- `{ { node1, node2 } }` is a single region consisting of two nodes.
---- This will be parsed by the parser in a single context... thus resulting
+--- On the other hand, `{ { node1, node2 } }` is a single region consisting of
+--- two nodes. This will be parsed by the parser in a single context, thus resulting
 --- in a single tree.
 ---
 --- This allows for embedded languages to be parsed together across different
 --- nodes, which is useful for templating languages like ERB and EJS.
 ---
---- Note, this call invalidates the tree and requires it to be parsed again.
+--- Note: This call invalidates the tree and requires it to be parsed again.
 ---
----@param regions (table) list of regions this tree should manage and parse.
+---@private
+---@param regions table 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
@@ -288,7 +309,7 @@ function LanguageTree:set_included_regions(regions)
   -- Trees are no longer valid now that we have changed regions.
   -- TODO(vigoux,steelsojka): Look into doing this smarter so we can use some of the
   --                          old trees for incremental parsing. Currently, this only
-  --                          effects injected languages.
+  --                          affects injected languages.
   self._trees = {}
   self:invalidate()
 end
@@ -299,7 +320,7 @@ function LanguageTree:included_regions()
 end
 
 ---@private
-local function get_node_range(node, id, metadata)
+local function get_range_from_metadata(node, id, metadata)
   if metadata[id] and metadata[id].range then
     return metadata[id].range
   end
@@ -362,7 +383,7 @@ function LanguageTree:_get_injections()
         elseif name == 'combined' then
           combined = true
         elseif name == 'content' and #ranges == 0 then
-          table.insert(ranges, get_node_range(node, id, metadata))
+          table.insert(ranges, get_range_from_metadata(node, id, metadata))
           -- Ignore any tags that start with "_"
           -- Allows for other tags to be used in matches
         elseif string.sub(name, 1, 1) ~= '_' then
@@ -371,7 +392,7 @@ function LanguageTree:_get_injections()
           end
 
           if #ranges == 0 then
-            table.insert(ranges, get_node_range(node, id, metadata))
+            table.insert(ranges, get_range_from_metadata(node, id, metadata))
           end
         end
       end
@@ -493,8 +514,8 @@ function LanguageTree:_on_detach(...)
   self:_do_callback('detach', ...)
 end
 
---- Registers callbacks for the parser.
----@param cbs table An |nvim_buf_attach()|-like table argument with the following keys :
+--- Registers callbacks for the |LanguageTree|.
+---@param cbs table An |nvim_buf_attach()|-like table argument with the following handlers:
 ---           - `on_bytes` : see |nvim_buf_attach()|, but this will be called _after_ the parsers callback.
 ---           - `on_changedtree` : a callback that will be called every time the tree has syntactical changes.
 ---              It will only be passed one argument, which is a table of the ranges (as node ranges) that
@@ -536,9 +557,10 @@ local function tree_contains(tree, range)
   return start_fits and end_fits
 end
 
---- Determines whether {range} is contained in this language tree
+--- Determines whether {range} is contained in the |LanguageTree|.
 ---
----@param range A range, that is a `{ start_line, start_col, end_line, end_col }` table.
+---@param range table `{ start_line, start_col, end_line, end_col }`
+---@return boolean
 function LanguageTree:contains(range)
   for _, tree in pairs(self._trees) do
     if tree_contains(tree, range) then
@@ -549,9 +571,50 @@ function LanguageTree:contains(range)
   return false
 end
 
---- Gets the appropriate language that contains {range}
+--- Gets the tree that contains {range}.
+---
+---@param range table `{ 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|
+function LanguageTree:tree_for_range(range, opts)
+  opts = opts or {}
+  local ignore = vim.F.if_nil(opts.ignore_injections, true)
+
+  if not ignore then
+    for _, child in pairs(self._children) do
+      for _, tree in pairs(child:trees()) do
+        if tree_contains(tree, range) then
+          return tree
+        end
+      end
+    end
+  end
+
+  for _, tree in pairs(self._trees) do
+    if tree_contains(tree, range) then
+      return tree
+    end
+  end
+
+  return nil
+end
+
+--- Gets the smallest named node that contains {range}.
+---
+---@param range table `{ 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|
+function LanguageTree:named_node_for_range(range, opts)
+  local tree = self:tree_for_range(range, opts)
+  return tree:root():named_descendant_for_range(unpack(range))
+end
+
+--- Gets the appropriate language that contains {range}.
 ---
----@param range A text range, see |LanguageTree:contains|
+---@param range table `{ 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
     if child:contains(range) then
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 103e85abfd..7ca7384a88 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -3,6 +3,11 @@ 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 query userdata Parsed query
 local Query = {}
 Query.__index = Query
 
@@ -34,11 +39,24 @@ local function safe_read(filename, read_quantifier)
   return content
 end
 
+---@private
+--- Adds {ilang} to {base_langs}, only if {ilang} is different than {lang}
+---
+---@return boolean true If lang == ilang
+local function add_included_lang(base_langs, lang, ilang)
+  if lang == ilang then
+    return true
+  end
+  table.insert(base_langs, ilang)
+  return false
+end
+
 --- Gets the list of files used to make up a query
 ---
----@param lang The language
----@param query_name The name of the query to load
----@param is_included Internal parameter, most of the time left as `nil`
+---@param lang string Language to get query for
+---@param query_name string Name of the query to load (e.g., "highlights")
+---@param is_included (boolean|nil) Internal parameter, most of the time left as `nil`
+---@return string[] query_files List of files to load for given query and language
 function M.get_query_files(lang, query_name, is_included)
   local query_path = string.format('queries/%s/%s.scm', lang, query_name)
   local lang_files = dedupe_files(a.nvim_get_runtime_file(query_path, true))
@@ -47,6 +65,9 @@ function M.get_query_files(lang, query_name, is_included)
     return {}
   end
 
+  local base_query = nil
+  local extensions = {}
+
   local base_langs = {}
 
   -- Now get the base languages by looking at the first line of every file
@@ -55,27 +76,53 @@ function M.get_query_files(lang, query_name, is_included)
   --
   -- {language} ::= {lang} | ({lang})
   local MODELINE_FORMAT = '^;+%s*inherits%s*:?%s*([a-z_,()]+)%s*$'
+  local EXTENDS_FORMAT = '^;+%s*extends%s*$'
 
-  for _, file in ipairs(lang_files) do
-    local modeline = safe_read(file, '*l')
+  for _, filename in ipairs(lang_files) do
+    local file, err = io.open(filename, 'r')
+    if not file then
+      error(err)
+    end
 
-    if modeline then
-      local langlist = modeline:match(MODELINE_FORMAT)
+    local extension = false
 
+    for modeline in
+      function()
+        return file:read('*l')
+      end
+    do
+      if not vim.startswith(modeline, ';') then
+        break
+      end
+
+      local langlist = modeline:match(MODELINE_FORMAT)
       if langlist then
         for _, incllang in ipairs(vim.split(langlist, ',', true)) do
           local is_optional = incllang:match('%(.*%)')
 
           if is_optional then
             if not is_included then
-              table.insert(base_langs, incllang:sub(2, #incllang - 1))
+              if add_included_lang(base_langs, lang, incllang:sub(2, #incllang - 1)) then
+                extension = true
+              end
             end
           else
-            table.insert(base_langs, incllang)
+            if add_included_lang(base_langs, lang, incllang) then
+              extension = true
+            end
           end
         end
+      elseif modeline:match(EXTENDS_FORMAT) then
+        extension = true
       end
     end
+
+    if extension then
+      table.insert(extensions, filename)
+    elseif base_query == nil then
+      base_query = filename
+    end
+    io.close(file)
   end
 
   local query_files = {}
@@ -83,7 +130,8 @@ function M.get_query_files(lang, query_name, is_included)
     local base_files = M.get_query_files(base_lang, query_name, true)
     vim.list_extend(query_files, base_files)
   end
-  vim.list_extend(query_files, lang_files)
+  vim.list_extend(query_files, { base_query })
+  vim.list_extend(query_files, extensions)
 
   return query_files
 end
@@ -109,24 +157,24 @@ local explicit_queries = setmetatable({}, {
   end,
 })
 
---- Sets the runtime query {query_name} for {lang}
+--- Sets the runtime query named {query_name} for {lang}
 ---
 --- This allows users to override any runtime files and/or configuration
 --- set by plugins.
 ---
----@param lang string: The language to use for the query
----@param query_name string: The name of the query (i.e. "highlights")
----@param text string: The query text (unparsed).
+---@param lang string Language to use for the query
+---@param query_name string Name of the query (e.g., "highlights")
+---@param text string Query text (unparsed).
 function M.set_query(lang, query_name, text)
   explicit_queries[lang][query_name] = M.parse_query(lang, text)
 end
 
 --- Returns the runtime query {query_name} for {lang}.
 ---
----@param lang The language to use for the query
----@param query_name The name of the query (i.e. "highlights")
+---@param lang string Language to use for the query
+---@param query_name string Name of the query (e.g. "highlights")
 ---
----@return The corresponding query, parsed.
+---@return Query Parsed query
 function M.get_query(lang, query_name)
   if explicit_queries[lang][query_name] then
     return explicit_queries[lang][query_name]
@@ -140,12 +188,9 @@ function M.get_query(lang, query_name)
   end
 end
 
-local query_cache = setmetatable({}, {
-  __index = function(tbl, key)
-    rawset(tbl, key, {})
-    return rawget(tbl, key)
-  end,
-})
+local query_cache = vim.defaulttable(function()
+  return setmetatable({}, { __mode = 'v' })
+end)
 
 --- Parse {query} as a string. (If the query is in a file, the caller
 --- should read the contents into a string before calling).
@@ -160,10 +205,10 @@ local query_cache = setmetatable({}, {
 ---   -` info.captures` also points to `captures`.
 ---   - `info.patterns` contains information about predicates.
 ---
----@param lang string The language
----@param query string A string containing the query (s-expr syntax)
+---@param lang string Language to use for the query
+---@param query string Query in s-expr syntax
 ---
----@returns The query
+---@return Query Parsed query
 function M.parse_query(lang, query)
   language.require_language(lang)
   local cached = query_cache[lang][query]
@@ -181,9 +226,15 @@ end
 
 --- Gets the text corresponding to a given node
 ---
----@param node the node
----@param source The buffer or string from which the node is extracted
-function M.get_node_text(node, source)
+---@param node userdata |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)
+function M.get_node_text(node, source, opts)
+  opts = opts or {}
+  local concat = vim.F.if_nil(opts.concat, true)
+
   local start_row, start_col, start_byte = node:start()
   local end_row, end_col, end_byte = node:end_()
 
@@ -210,7 +261,7 @@ function M.get_node_text(node, source)
       end
     end
 
-    return table.concat(lines, '\n')
+    return concat and table.concat(lines, '\n') or lines
   elseif type(source) == 'string' then
     return source:sub(start_byte + 1, end_byte)
   end
@@ -367,9 +418,8 @@ local directive_handlers = {
 
 --- Adds a new predicate to be used in queries
 ---
----@param name the name of the predicate, without leading #
----@param handler the handler function to be used
----      signature will be (match, pattern, bufnr, predicate)
+---@param name string Name of the predicate, without leading #
+---@param handler function(match:string, pattern:string, bufnr:number, predicate:function)
 function M.add_predicate(name, handler, force)
   if predicate_handlers[name] and not force then
     error(string.format('Overriding %s', name))
@@ -385,9 +435,8 @@ end
 --- can set node level data by using the capture id on the
 --- metadata table `metadata[capture_id].key = value`
 ---
----@param name the name of the directive, without leading #
----@param handler the handler function to be used
----      signature will be (match, pattern, bufnr, predicate, metadata)
+---@param name string Name of the directive, without leading #
+---@param handler function(match:string, pattern:string, bufnr:number, predicate:function, metadata:table)
 function M.add_directive(name, handler, force)
   if directive_handlers[name] and not force then
     error(string.format('Overriding %s', name))
@@ -397,12 +446,13 @@ function M.add_directive(name, handler, force)
 end
 
 --- Lists the currently available directives to use in queries.
----@return The list of supported directives.
+---@return string[] List of supported directives.
 function M.list_directives()
   return vim.tbl_keys(directive_handlers)
 end
 
----@return The list of supported predicates.
+--- Lists the currently available predicates to use in queries.
+---@return string[] List of supported predicates.
 function M.list_predicates()
   return vim.tbl_keys(predicate_handlers)
 end
@@ -489,17 +539,16 @@ end
 
 --- Iterate over all captures from all matches inside {node}
 ---
---- {source} is needed if the query contains predicates, then the caller
+--- {source} is needed if the query contains predicates; then the caller
 --- must ensure to use a freshly parsed tree consistent with the current
 --- text of the buffer (if relevant). {start_row} and {end_row} can be used to limit
 --- matches inside a row range (this is typically used with root node
---- as the node, i e to get syntax highlight matches in the current
---- viewport). When omitted the start and end row values are used from the given node.
+--- as the {node}, i.e., to get syntax highlight matches in the current
+--- viewport). When omitted, the {start} and {end} row values are used from the given node.
 ---
---- The iterator returns three values, a numeric id identifying the capture,
+--- The iterator returns three values: a numeric id identifying the capture,
 --- the captured node, and metadata from any directives processing the match.
 --- The following example shows how to get captures by name:
----
 --- <pre>
 --- for id, node, metadata in query:iter_captures(tree:root(), bufnr, first, last) do
 ---   local name = query.captures[id] -- name of the capture in the query
@@ -510,13 +559,14 @@ end
 --- end
 --- </pre>
 ---
----@param node The node under which the search will occur
----@param source The source buffer or string to extract text from
----@param start The starting line of the search
----@param stop The stopping line of the search (end-exclusive)
+---@param node userdata |tsnode| under which the search will occur
+---@param source (number|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)
 ---
----@returns The matching capture id
----@returns The captured node
+---@return number capture Matching capture id
+---@return table capture_node Capture for {node}
+---@return table metadata for the {capture}
 function Query:iter_captures(node, source, start, stop)
   if type(source) == 'number' and source == 0 then
     source = vim.api.nvim_get_current_buf()
@@ -546,14 +596,13 @@ end
 
 --- Iterates the matches of self on a given range.
 ---
---- Iterate over all matches within a node. The arguments are the same as
---- for |query:iter_captures()| but the iterated values are different:
+--- Iterate over all matches within a {node}. The arguments are the same as
+--- for |Query:iter_captures()| but the iterated values are different:
 --- an (1-based) index of the pattern in the query, a table mapping
 --- capture indices to nodes, and metadata from any directives processing the match.
---- If the query has more than one pattern the capture table might be sparse,
+--- If the query has more than one pattern, the capture table might be sparse
 --- and e.g. `pairs()` method should be used over `ipairs`.
---- Here an example iterating over all captures in every match:
----
+--- Here is an example iterating over all captures in every match:
 --- <pre>
 --- for pattern, match, metadata in cquery:iter_matches(tree:root(), bufnr, first, last) do
 ---   for id, node in pairs(match) do
@@ -567,13 +616,14 @@ end
 --- end
 --- </pre>
 ---
----@param node The node under which the search will occur
----@param source The source buffer or string to search
----@param start The starting line of the search
----@param stop The stopping line of the search (end-exclusive)
+---@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)
 ---
----@returns The matching pattern id
----@returns The matching match
+---@return number pattern id
+---@return table match
+---@return table metadata
 function Query:iter_matches(node, source, start, stop)
   if type(source) == 'number' and source == 0 then
     source = vim.api.nvim_get_current_buf()