fix(treesitter): revert to using iter_captures in highlighter

Fixes #27895

1 files changed, 29 insertions, 45 deletions

diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index 7bc6e5c019..1e6f128461 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -4,7 +4,7 @@ local Range = require('vim.treesitter._range')
 
 local ns = api.nvim_create_namespace('treesitter/highlighter')
 
----@alias vim.treesitter.highlighter.Iter fun(): integer, table<integer, TSNode[]>, vim.treesitter.query.TSMetadata
+---@alias vim.treesitter.highlighter.Iter fun(end_line: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, table<integer, TSNode[]>
 
 ---@class (private) vim.treesitter.highlighter.Query
 ---@field private _query vim.treesitter.Query?
@@ -57,7 +57,6 @@ 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
@@ -193,20 +192,12 @@ 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)
@@ -296,9 +287,6 @@ end
 ---@param is_spell_nav boolean
 local function on_line_impl(self, buf, line, is_spell_nav)
   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()
 
@@ -308,23 +296,25 @@ local function on_line_impl(self, buf, line, is_spell_nav)
     end
 
     if state.iter == nil or state.next_row < line then
-      state.iter = state.highlighter_query
-        :query()
-        :iter_matches(root_node, self.bufnr, line, root_end_row + 1, { all = true })
+      state.iter =
+        state.highlighter_query:query():iter_captures(root_node, self.bufnr, line, root_end_row + 1)
     end
 
     while line >= state.next_row do
-      local pattern, match, metadata = state.iter()
+      local capture, node, metadata, match = state.iter(line)
 
-      if not match then
-        state.next_row = root_end_row + 1
+      local range = { root_end_row + 1, 0, root_end_row + 1, 0 }
+      if node then
+        range = vim.treesitter.get_range(node, buf, metadata and metadata[capture])
       end
+      local start_row, start_col, end_row, end_col = Range.unpack4(range)
+
+      if capture then
+        local hl = state.highlighter_query:get_hl_from_capture(capture)
 
-      for capture, nodes in pairs(match or {}) do
         local capture_name = state.highlighter_query:query().captures[capture]
-        local spell, spell_pri_offset = get_spell(capture_name)
 
-        local hl = state.highlighter_query:get_hl_from_capture(capture)
+        local spell, spell_pri_offset = get_spell(capture_name)
 
         -- The "priority" attribute can be set at the pattern level or on a particular capture
         local priority = (
@@ -332,34 +322,28 @@ local function on_line_impl(self, buf, line, is_spell_nav)
           or vim.highlight.priorities.treesitter
         ) + spell_pri_offset
 
-        local url = get_url(match, buf, capture, metadata)
-
         -- The "conceal" attribute can be set at the pattern level or on a particular capture
         local conceal = metadata.conceal or metadata[capture] and metadata[capture].conceal
 
-        for _, node in ipairs(nodes) do
-          local range = vim.treesitter.get_range(node, buf, metadata[capture])
-          local start_row, start_col, end_row, end_col = Range.unpack4(range)
-
-          if hl and end_row >= line and (not is_spell_nav or spell ~= nil) then
-            api.nvim_buf_set_extmark(buf, ns, start_row, start_col, {
-              end_line = end_row,
-              end_col = end_col,
-              hl_group = hl,
-              ephemeral = true,
-              priority = priority,
-              _subpriority = pattern_offset + pattern,
-              conceal = conceal,
-              spell = spell,
-              url = url,
-            })
-          end
-
-          if start_row > line then
-            state.next_row = start_row
-          end
+        local url = get_url(match, buf, capture, metadata)
+
+        if hl and end_row >= line and (not is_spell_nav or spell ~= nil) then
+          api.nvim_buf_set_extmark(buf, ns, start_row, start_col, {
+            end_line = end_row,
+            end_col = end_col,
+            hl_group = hl,
+            ephemeral = true,
+            priority = priority,
+            conceal = conceal,
+            spell = spell,
+            url = url,
+          })
         end
       end
+
+      if start_row > line then
+        state.next_row = start_row
+      end
     end
   end)
 end