fix(treesitter): revert to using iter_captures in highlighter

Fixes #27895

2 files changed, 36 insertions, 51 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
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 67b8c596b8..30cd00c617 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -811,12 +811,13 @@ end
 --- as the {node}, i.e., to get syntax highlight matches in the current
 --- viewport). When omitted, the {start} and {stop} row values are used from the given node.
 ---
---- The iterator returns three values: a numeric id identifying the capture,
---- the captured node, and metadata from any directives processing the match.
+--- The iterator returns four values: a numeric id identifying the capture,
+--- the captured node, metadata from any directives processing the match,
+--- and the match itself.
 --- The following example shows how to get captures by name:
 ---
 --- ```lua
---- for id, node, metadata in query:iter_captures(tree:root(), bufnr, first, last) do
+--- for id, node, metadata, match in query:iter_captures(tree:root(), bufnr, first, last) do
 ---   local name = query.captures[id] -- name of the capture in the query
 ---   -- typically useful info about the node:
 ---   local type = node:type() -- type of the captured node
@@ -830,8 +831,8 @@ end
 ---@param start? integer Starting line for the search. Defaults to `node:start()`.
 ---@param stop? integer Stopping line for the search (end-exclusive). Defaults to `node:end_()`.
 ---
----@return (fun(end_line: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata):
----        capture id, capture node, metadata
+---@return (fun(end_line: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, table<integer, TSNode>):
+---        capture id, capture node, metadata, match
 function Query:iter_captures(node, source, start, stop)
   if type(source) == 'number' and source == 0 then
     source = api.nvim_get_current_buf()
@@ -856,7 +857,7 @@ function Query:iter_captures(node, source, start, stop)
 
       self:apply_directives(match, match.pattern, source, metadata)
     end
-    return capture, captured_node, metadata
+    return capture, captured_node, metadata, match
   end
   return iter
 end