perf(treesitter): do not scan past given line for predicate match

Problem
---
If a highlighter query returns a significant number of predicate
non-matches, the highlighter will scan well past the end of the window.

Solution
---
In the iterator returned from `iter_captures`, accept an optional
parameter `end_line`. If no parameter provided, the behavior is
unchanged, hence this is a non-invasive tweak.

Fixes: #25113 nvim-treesitter/nvim-treesitter#5057

3 files changed, 42 insertions, 35 deletions

diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index 64b4ca7ca2..ee34c45cce 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -976,8 +976,8 @@ Query:iter_captures({node}, {source}, {start}, {stop})
       • {stop}    (integer) Stopping line for the search (end-exclusive)
 
     Return: ~
-        (fun(): integer, TSNode, TSMetadata): capture id, capture node,
-        metadata
+        (fun(end_line: integer|nil): integer, TSNode, TSMetadata): capture id,
+        capture node, metadata
 
                                                         *Query:iter_matches()*
 Query:iter_matches({node}, {source}, {start}, {stop}, {opts})
diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index 8d4d6a9337..496193c6ed 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -2,7 +2,7 @@ local api = vim.api
 local query = vim.treesitter.query
 local Range = require('vim.treesitter._range')
 
----@alias TSHlIter fun(): integer, TSNode, TSMetadata
+---@alias TSHlIter fun(end_line: integer|nil): integer, TSNode, TSMetadata
 
 ---@class TSHighlightState
 ---@field next_row integer
@@ -241,40 +241,43 @@ local function on_line_impl(self, buf, line, is_spell_nav)
     end
 
     while line >= state.next_row do
-      local capture, node, metadata = state.iter()
+      local capture, node, metadata = state.iter(line)
 
-      if capture == nil then
-        break
+      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 range = vim.treesitter.get_range(node, buf, metadata[capture])
       local start_row, start_col, end_row, end_col = Range.unpack4(range)
-      local hl = highlighter_query.hl_cache[capture]
-
-      local capture_name = highlighter_query:query().captures[capture]
-      local spell = nil ---@type boolean?
-      if capture_name == 'spell' then
-        spell = true
-      elseif capture_name == 'nospell' then
-        spell = false
-      end
 
-      -- Give nospell a higher priority so it always overrides spell captures.
-      local spell_pri_offset = capture_name == 'nospell' and 1 or 0
-
-      if hl and end_row >= line and (not is_spell_nav or spell ~= nil) then
-        local priority = (tonumber(metadata.priority) or vim.highlight.priorities.treesitter)
-          + spell_pri_offset
-        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 = metadata.conceal,
-          spell = spell,
-        })
+      if capture then
+        local hl = highlighter_query.hl_cache[capture]
+
+        local capture_name = highlighter_query:query().captures[capture]
+        local spell = nil ---@type boolean?
+        if capture_name == 'spell' then
+          spell = true
+        elseif capture_name == 'nospell' then
+          spell = false
+        end
+
+        -- Give nospell a higher priority so it always overrides spell captures.
+        local spell_pri_offset = capture_name == 'nospell' and 1 or 0
+
+        if hl and end_row >= line and (not is_spell_nav or spell ~= nil) then
+          local priority = (tonumber(metadata.priority) or vim.highlight.priorities.treesitter)
+            + spell_pri_offset
+          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 = metadata.conceal,
+            spell = spell,
+          })
+        end
       end
+
       if start_row > line then
         state.next_row = start_row
       end
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index d7973cc48f..6d9b214d4a 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -708,7 +708,8 @@ end
 ---@param start integer Starting line for the search
 ---@param stop integer Stopping line for the search (end-exclusive)
 ---
----@return (fun(): integer, TSNode, TSMetadata): capture id, capture node, metadata
+---@return (fun(end_line: integer|nil): integer, TSNode, TSMetadata):
+---        capture id, capture node, metadata
 function Query:iter_captures(node, source, start, stop)
   if type(source) == 'number' and source == 0 then
     source = api.nvim_get_current_buf()
@@ -717,7 +718,7 @@ function Query:iter_captures(node, source, start, stop)
   start, stop = value_or_node_range(start, stop, node)
 
   local raw_iter = node:_rawquery(self.query, true, start, stop)
-  local function iter()
+  local function iter(end_line)
     local capture, captured_node, match = raw_iter()
     local metadata = {}
 
@@ -725,7 +726,10 @@ function Query:iter_captures(node, source, start, stop)
       local active = self:match_preds(match, match.pattern, source)
       match.active = active
       if not active then
-        return iter() -- tail call: try next match
+        if end_line and captured_node:range() > end_line then
+          return nil, captured_node, nil
+        end
+        return iter(end_line) -- tail call: try next match
       end
 
       self:apply_directives(match, match.pattern, source, metadata)