fix(treesitter): highlight injections properly

`on_line_impl` doesn't highlight single lines, so using pattern indexes
to offset priority doesn't work.

2 files changed, 22 insertions, 20 deletions

diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index cbab5e990e..6175977b49 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -57,6 +57,7 @@ end
 ---@field next_row integer
 ---@field iter vim.treesitter.highlighter.Iter?
 ---@field highlighter_query vim.treesitter.highlighter.Query
+---@field level integer Injection level
 
 ---@nodoc
 ---@class vim.treesitter.highlighter
@@ -192,12 +193,20 @@ function TSHighlighter:prepare_highlight_states(srow, erow)
       return
     end
 
+    local level = 0
+    local t = tree
+    while t do
+      t = t:parent()
+      level = level + 1
+    end
+
     -- _highlight_states should be a list so that the highlights are added in the same order as
     -- for_each_tree traversal. This ensures that parents' highlight don't override children's.
     table.insert(self._highlight_states, {
       tstree = tstree,
       next_row = 0,
       iter = nil,
+      level = level,
       highlighter_query = highlighter_query,
     })
   end)
@@ -248,14 +257,10 @@ end
 ---@param line integer
 ---@param is_spell_nav boolean
 local function on_line_impl(self, buf, line, is_spell_nav)
-  -- Track the maximum pattern index encountered in each tree. For subsequent
-  -- trees, the subpriority passed to nvim_buf_set_extmark is offset by the
-  -- largest pattern index from the prior tree. This ensures that extmarks
-  -- from subsequent trees always appear "on top of" extmarks from previous
-  -- trees (e.g. injections should always appear over base highlights).
-  local pattern_offset = 0
-
   self:for_each_highlight_state(function(state)
+    -- Use the injection level to offset the subpriority passed to nvim_buf_set_extmark
+    -- so injections always appear over base highlights.
+    local pattern_offset = state.level * 1000
     local root_node = state.tstree:root()
     local root_start_row, _, root_end_row, _ = root_node:range()
 
@@ -270,14 +275,9 @@ local function on_line_impl(self, buf, line, is_spell_nav)
         :iter_matches(root_node, self.bufnr, line, root_end_row + 1, { all = true })
     end
 
-    local max_pattern_index = 0
     while line >= state.next_row do
       local pattern, match, metadata = state.iter()
 
-      if pattern and pattern > max_pattern_index then
-        max_pattern_index = pattern
-      end
-
       if not match then
         state.next_row = root_end_row + 1
       end
@@ -343,8 +343,6 @@ local function on_line_impl(self, buf, line, is_spell_nav)
         end
       end
     end
-
-    pattern_offset = pattern_offset + max_pattern_index
   end)
 end
 
diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index 62714d3f1b..ec933f5194 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -81,7 +81,7 @@ local TSCallbackNames = {
 ---List of regions this tree should manage and parse. If nil then regions are
 ---taken from _trees. This is mostly a short-lived cache for included_regions()
 ---@field private _lang string Language name
----@field private _parent_lang? string Parent language name
+---@field private _parent? vim.treesitter.LanguageTree Parent LanguageTree
 ---@field private _source (integer|string) Buffer or string to parse
 ---@field private _trees table<integer, TSTree> Reference to parsed tree (one for each language).
 ---Each key is the index of region, which is synced with _regions and _valid.
@@ -106,9 +106,8 @@ LanguageTree.__index = LanguageTree
 ---@param source (integer|string) Buffer or text string to parse
 ---@param lang string Root language of this tree
 ---@param opts vim.treesitter.LanguageTree.new.Opts?
----@param parent_lang? string Parent language name of this tree
 ---@return vim.treesitter.LanguageTree parser object
-function LanguageTree.new(source, lang, opts, parent_lang)
+function LanguageTree.new(source, lang, opts)
   language.add(lang)
   opts = opts or {}
 
@@ -122,7 +121,6 @@ function LanguageTree.new(source, lang, opts, parent_lang)
   local self = {
     _source = source,
     _lang = lang,
-    _parent_lang = parent_lang,
     _children = {},
     _trees = {},
     _opts = opts,
@@ -505,19 +503,25 @@ function LanguageTree:add_child(lang)
     self:remove_child(lang)
   end
 
-  local child = LanguageTree.new(self._source, lang, self._opts, self:lang())
+  local child = LanguageTree.new(self._source, lang, self._opts)
 
   -- Inherit recursive callbacks
   for nm, cb in pairs(self._callbacks_rec) do
     vim.list_extend(child._callbacks_rec[nm], cb)
   end
 
+  child._parent = self
   self._children[lang] = child
   self:_do_callback('child_added', self._children[lang])
 
   return self._children[lang]
 end
 
+--- @package
+function LanguageTree:parent()
+  return self._parent
+end
+
 --- Removes a child language from this |LanguageTree|.
 ---
 ---@private
@@ -792,7 +796,7 @@ function LanguageTree:_get_injection(match, metadata)
   local combined = metadata['injection.combined'] ~= nil
   local injection_lang = metadata['injection.language'] --[[@as string?]]
   local lang = metadata['injection.self'] ~= nil and self:lang()
-    or metadata['injection.parent'] ~= nil and self._parent_lang
+    or metadata['injection.parent'] ~= nil and self._parent
     or (injection_lang and resolve_lang(injection_lang))
   local include_children = metadata['injection.include-children'] ~= nil