From 5edbabdbec0ac3fba33be8afc008845130158583 Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Sun, 3 Aug 2025 21:39:55 -0500
Subject: perf: add on_range in treesitter highlighting

---
 runtime/lua/vim/treesitter/query.lua | 48 +++++++++++++++++++++++++++++-------
 1 file changed, 39 insertions(+), 9 deletions(-)

(limited to 'runtime/lua/vim/treesitter/query.lua')

diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index b93051b895..cdb29fdb61 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -4,6 +4,7 @@
 local api = vim.api
 local language = require('vim.treesitter.language')
 local memoize = vim.func._memoize
+local cmp_ge = require('vim.treesitter._range').cmp_pos.ge
 
 local MODELINE_FORMAT = '^;+%s*inherits%s*:?%s*([a-z_,()]+)%s*$'
 local EXTENDS_FORMAT = '^;+%s*extends%s*$'
@@ -951,18 +952,20 @@ end
 ---
 ---@param node TSNode under which the search will occur
 ---@param source (integer|string) Source buffer or string to extract text from
----@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_()`.
+---@param start_row? integer Starting line for the search. Defaults to `node:start()`.
+---@param end_row? integer Stopping line for the search (end-inclusive, unless `stop_col` is provided). Defaults to `node:end_()`.
 ---@param opts? table Optional keyword arguments:
 ---   - max_start_depth (integer) if non-zero, sets the maximum start depth
 ---     for each match. This is used to prevent traversing too deep into a tree.
 ---   - match_limit (integer) Set the maximum number of in-progress matches (Default: 256).
+---   - start_col (integer) Starting column for the search.
+---   - end_col (integer) Stopping column for the search (end-exclusive).
 ---
----@return (fun(end_line: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch, TSTree):
+---@return (fun(end_line: integer|nil, end_col: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch, TSTree):
 ---        capture id, capture node, metadata, match, tree
 ---
 ---@note Captures are only returned if the query pattern of a specific capture contained predicates.
-function Query:iter_captures(node, source, start, stop, opts)
+function Query:iter_captures(node, source, start_row, end_row, opts)
   opts = opts or {}
   opts.match_limit = opts.match_limit or 256
 
@@ -970,17 +973,24 @@ function Query:iter_captures(node, source, start, stop, opts)
     source = api.nvim_get_current_buf()
   end
 
-  start, stop = value_or_node_range(start, stop, node)
+  start_row, end_row = value_or_node_range(start_row, end_row, node)
 
   local tree = node:tree()
-  local cursor = vim._create_ts_querycursor(node, self.query, start, stop, opts)
+  local cursor = vim._create_ts_querycursor(node, self.query, {
+    start_row = start_row,
+    start_col = opts.start_col or 0,
+    end_row = end_row,
+    end_col = opts.end_col or 0,
+    max_start_depth = opts.max_start_depth,
+    match_limit = opts.match_limit or 256,
+  })
 
   -- For faster checks that a match is not in the cache.
   local highest_cached_match_id = -1
   ---@type table<integer, vim.treesitter.query.TSMetadata>
   local match_cache = {}
 
-  local function iter(end_line)
+  local function iter(end_line, end_col)
     local capture, captured_node, match = cursor:next_capture()
 
     if not capture then
@@ -1005,9 +1015,22 @@ function Query:iter_captures(node, source, start, stop, opts)
         local predicates = processed_pattern.predicates
         if not self:_match_predicates(predicates, pattern_i, captures, source) then
           cursor:remove_match(match_id)
-          if end_line and captured_node:range() > end_line then
+
+          local row, col = captured_node:range()
+
+          local outside = false
+          if end_line then
+            if end_col then
+              outside = cmp_ge(row, col, end_line, end_col)
+            else
+              outside = row > end_line
+            end
+          end
+
+          if outside then
             return nil, captured_node, nil, nil
           end
+
           return iter(end_line) -- tail call: try next match
         end
 
@@ -1072,7 +1095,14 @@ function Query:iter_matches(node, source, start, stop, opts)
   start, stop = value_or_node_range(start, stop, node)
 
   local tree = node:tree()
-  local cursor = vim._create_ts_querycursor(node, self.query, start, stop, opts)
+  local cursor = vim._create_ts_querycursor(node, self.query, {
+    start_row = start,
+    start_col = 0,
+    end_row = stop,
+    end_col = 0,
+    max_start_depth = opts.max_start_depth,
+    match_limit = opts.match_limit or 256,
+  })
 
   local function iter()
     local match = cursor:next_match()
-- 
cgit