From d2ce9e21c86552166bc09fca3f8a89e72892e9b1 Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Sat, 7 Dec 2024 13:14:38 -0600
Subject: [PATCH 01/13] test: add treesitter long lines benchmark

---
 test/benchmark/decor_spec.lua      |  38 +++++++-
 test/benchmark/treesitter_spec.lua | 148 ++++++++++++++++++++++++++++-
 2 files changed, 183 insertions(+), 3 deletions(-)

diff --git a/test/benchmark/decor_spec.lua b/test/benchmark/decor_spec.lua
index 0994023c2d..2e04d44686 100644
--- a/test/benchmark/decor_spec.lua
+++ b/test/benchmark/decor_spec.lua
@@ -6,8 +6,7 @@ describe('decor perf', function()
   before_each(n.clear)
 
   it('can handle long lines', function()
-    local screen = Screen.new(100, 101)
-    screen:attach()
+    Screen.new(100, 101)
 
     local result = exec_lua [==[
       local ephemeral_pattern = {
@@ -99,4 +98,39 @@ describe('decor perf', function()
 
     print('\nTotal ' .. fmt(total) .. '\nDecoration provider: ' .. fmt(provider))
   end)
+
+  it('can handle long lines with treesitter highlighting', function()
+    Screen.new(100, 51)
+
+    local result = exec_lua [==[
+      local long_line = 'local a = { ' .. ('a = 5, '):rep(2000) .. '}'
+      vim.api.nvim_buf_set_lines(0, 0, 0, false, { long_line })
+      vim.api.nvim_win_set_cursor(0, { 1, 0 })
+      vim.treesitter.start(0, 'lua')
+
+      local total = {}
+      for i = 1, 50 do
+        local tic = vim.uv.hrtime()
+        vim.cmd'redraw!'
+        local toc = vim.uv.hrtime()
+        table.insert(total, toc - tic)
+      end
+
+      return { total }
+    ]==]
+
+    local total = unpack(result)
+    table.sort(total)
+
+    local ms = 1 / 1000000
+    local res = string.format(
+      'min, 25%%, median, 75%%, max:\n\t%0.1fms,\t%0.1fms,\t%0.1fms,\t%0.1fms,\t%0.1fms',
+      total[1] * ms,
+      total[1 + math.floor(#total * 0.25)] * ms,
+      total[1 + math.floor(#total * 0.5)] * ms,
+      total[1 + math.floor(#total * 0.75)] * ms,
+      total[#total] * ms
+    )
+    print('\nTotal ' .. res)
+  end)
 end)
diff --git a/test/benchmark/treesitter_spec.lua b/test/benchmark/treesitter_spec.lua
index b13b0fd9da..f83aa51eb0 100644
--- a/test/benchmark/treesitter_spec.lua
+++ b/test/benchmark/treesitter_spec.lua
@@ -1,10 +1,11 @@
 local n = require('test.functional.testnvim')()
+local Screen = require('test.functional.ui.screen')
 
 local clear = n.clear
 local exec_lua = n.exec_lua
 
 describe('treesitter perf', function()
-  setup(function()
+  before_each(function()
     clear()
   end)
 
@@ -47,4 +48,149 @@ describe('treesitter perf', function()
       return vim.uv.hrtime() - start
     ]]
   end)
+
+  local function test_long_line(pos, wrap, line, grid)
+    local screen = Screen.new(20, 11)
+
+    local result = exec_lua(
+      [==[
+        local pos, wrap, line = ...
+
+        vim.api.nvim_buf_set_lines(0, 0, 0, false, { line })
+        vim.api.nvim_win_set_cursor(0, pos)
+        vim.api.nvim_set_option_value('wrap', wrap, { win = 0 })
+
+        vim.treesitter.start(0, 'lua')
+
+        local total = {}
+        for i = 1, 100 do
+          local tic = vim.uv.hrtime()
+          vim.cmd'redraw!'
+          local toc = vim.uv.hrtime()
+          table.insert(total, toc - tic)
+        end
+
+        return { total }
+      ]==],
+      pos,
+      wrap,
+      line
+    )
+
+    screen:expect({ grid = grid or '' })
+
+    local total = unpack(result)
+    table.sort(total)
+
+    local ms = 1 / 1000000
+    local res = string.format(
+      'min, 25%%, median, 75%%, max:\n\t%0.2fms,\t%0.2fms,\t%0.2fms,\t%0.2fms,\t%0.2fms',
+      total[1] * ms,
+      total[1 + math.floor(#total * 0.25)] * ms,
+      total[1 + math.floor(#total * 0.5)] * ms,
+      total[1 + math.floor(#total * 0.75)] * ms,
+      total[#total] * ms
+    )
+    print('\nTotal ' .. res)
+  end
+
+  local long_line = 'local a = { ' .. ('a = 5, '):rep(500) .. '}'
+  it('can redraw the beginning of a long line with wrapping', function()
+    local grid = [[
+      {15:^local} {25:a} {15:=} {16:{} {25:a} {15:=} {26:5}{16:,} {25:a}|
+       {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} |
+      {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,}|
+       {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}|
+      {16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} |
+      {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=}|
+       {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} |
+      {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a}|
+       {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} |
+      {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,}|
+                          |
+    ]]
+    test_long_line({ 1, 0 }, true, long_line, grid)
+  end)
+
+  it('can redraw the middle of a long line with wrapping', function()
+    local grid = [[
+      {1:<<<}{26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} |
+      {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,}|
+       {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}|
+      {16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} |
+      {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=}|
+       {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} |
+      {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a}|
+       {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} |
+      {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,}|
+       {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a}^ {15:=} {26:5}|
+                          |
+    ]]
+    test_long_line({ 1, math.floor(#long_line / 2) }, true, long_line, grid)
+  end)
+
+  it('can redraw the end of a long line with wrapping', function()
+    local grid = [[
+      {1:<<<}{25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=}|
+       {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} |
+      {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a}|
+       {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} |
+      {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,}|
+       {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}|
+      {16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} |
+      {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=}|
+       {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {25:a} |
+      {15:=} {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {16:^}}       |
+                          |
+    ]]
+    test_long_line({ 1, #long_line - 1 }, true, long_line, grid)
+  end)
+
+  it('can redraw the beginning of a long line without wrapping', function()
+    local grid = [[
+      {15:^local} {25:a} {15:=} {16:{} {25:a} {15:=} {26:5}{16:,} {25:a}|
+                          |
+      {1:~                   }|*8
+                          |
+    ]]
+    test_long_line({ 1, 0 }, false, long_line, grid)
+  end)
+
+  it('can redraw the middle of a long line without wrapping', function()
+    local grid = [[
+      {16:,} {25:a} {15:=} {26:5}{16:,} {25:a}^ {15:=} {26:5}{16:,} {25:a} {15:=} |
+                          |
+      {1:~                   }|*8
+                          |
+    ]]
+    test_long_line({ 1, math.floor(#long_line / 2) }, false, long_line, grid)
+  end)
+
+  it('can redraw the end of a long line without wrapping', function()
+    local grid = [[
+      {26:5}{16:,} {25:a} {15:=} {26:5}{16:,} {16:^}}         |
+                          |
+      {1:~                   }|*8
+                          |
+    ]]
+    test_long_line({ 1, #long_line - 1 }, false, long_line, grid)
+  end)
+
+  local long_line_mb = 'local a = { ' .. ('À = 5, '):rep(500) .. '}'
+  it('can redraw the middle of a long line with multibyte characters', function()
+    local grid = [[
+      {1:<<<}{26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} |
+      {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,}|
+       {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}|
+      {16:,} {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} |
+      {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À} {15:=}|
+       {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À} |
+      {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À}|
+       {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} |
+      {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,}|
+       {25:À} {15:=} {26:5}{16:,} {25:À} {15:=} {26:5}{16:,} {25:À}^ {15:=} {26:5}|
+                          |
+    ]]
+    test_long_line({ 1, math.floor(#long_line_mb / 2) }, true, long_line_mb, grid)
+  end)
 end)

From 260b4bb3a58379afedaa0f1b9c75fff04476e54c Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Fri, 29 Nov 2024 22:47:33 -0600
Subject: [PATCH 02/13] perf: add on_range in treesitter highlighting

---
 runtime/lua/vim/treesitter/highlighter.lua | 42 +++++++++++------
 runtime/lua/vim/treesitter/query.lua       | 35 ++++++++++++++
 src/nvim/api/extmark.c                     |  1 +
 src/nvim/api/keysets_defs.h                |  2 +
 src/nvim/decoration_defs.h                 |  1 +
 src/nvim/decoration_provider.c             | 29 +++++++++++-
 src/nvim/drawline.c                        | 54 ++++++++++++++++++++--
 src/nvim/lua/executor.c                    |  3 ++
 src/nvim/lua/treesitter.c                  | 47 +++++++++++++++++++
 9 files changed, 193 insertions(+), 21 deletions(-)

diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index 8ce8652f7d..1cb61b7af4 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -55,6 +55,7 @@ end
 ---@class (private) vim.treesitter.highlighter.State
 ---@field tstree TSTree
 ---@field next_row integer
+---@field next_col integer
 ---@field iter vim.treesitter.highlighter.Iter?
 ---@field highlighter_query vim.treesitter.highlighter.Query
 
@@ -197,6 +198,7 @@ function TSHighlighter:prepare_highlight_states(srow, erow)
     table.insert(self._highlight_states, {
       tstree = tstree,
       next_row = 0,
+      next_col = 0,
       iter = nil,
       highlighter_query = highlighter_query,
     })
@@ -280,27 +282,34 @@ end
 ---@param buf integer
 ---@param line integer
 ---@param is_spell_nav boolean
-local function on_line_impl(self, buf, line, is_spell_nav)
+local function on_range_impl(self, buf, from_line, from_col, until_line, until_col, is_spell_nav)
   self:for_each_highlight_state(function(state)
     local root_node = state.tstree:root()
     local root_start_row, _, root_end_row, _ = root_node:range()
 
     -- Only consider trees that contain this line
-    if root_start_row > line or root_end_row < line then
+    if root_start_row > until_line or root_end_row < from_line then
       return
     end
 
-    if state.iter == nil or state.next_row < line then
+    if
+      state.iter == nil
+      or state.next_row < from_line
+      or (state.next_row == from_line and state.next_col < from_col)
+    then
       -- Mainly used to skip over folds
 
       -- TODO(lewis6991): Creating a new iterator loses the cached predicate results for query
       -- matches. Move this logic inside iter_captures() so we can maintain the cache.
-      state.iter =
-        state.highlighter_query:query():iter_captures(root_node, self.bufnr, line, root_end_row + 1)
+      state.iter = state.highlighter_query
+        :query()
+        :iter_captures2(root_node, self.bufnr, from_line, from_col, root_end_row + 1, 0)
     end
 
-    while line >= state.next_row do
-      local capture, node, metadata, match = state.iter(line)
+    while
+      until_line > state.next_row or (until_line == state.next_row and until_col > state.next_col)
+    do
+      local capture, node, metadata, match = state.iter(until_line, until_col)
 
       local range = { root_end_row + 1, 0, root_end_row + 1, 0 }
       if node then
@@ -326,7 +335,7 @@ local function on_line_impl(self, buf, line, is_spell_nav)
 
         local url = get_url(match, buf, capture, metadata)
 
-        if hl and end_row >= line and (not is_spell_nav or spell ~= nil) then
+        if hl and end_row >= from_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,
@@ -340,8 +349,13 @@ local function on_line_impl(self, buf, line, is_spell_nav)
         end
       end
 
-      if start_row > line then
+      if start_row > until_line or (start_row == until_line and start_col >= until_col) then
         state.next_row = start_row
+        state.next_col = start_col
+      end
+
+      if not capture then
+        break
       end
     end
   end)
@@ -351,13 +365,13 @@ end
 ---@param _win integer
 ---@param buf integer
 ---@param line integer
-function TSHighlighter._on_line(_, _win, buf, line, _)
+function TSHighlighter._on_range(_, _win, buf, sr, sc, er, ec, _)
   local self = TSHighlighter.active[buf]
   if not self then
     return
   end
 
-  on_line_impl(self, buf, line, false)
+  on_range_impl(self, buf, sr, sc, er, ec, false)
 end
 
 ---@private
@@ -375,9 +389,7 @@ function TSHighlighter._on_spell_nav(_, _, buf, srow, _, erow, _)
   local highlight_states = self._highlight_states
   self:prepare_highlight_states(srow, erow)
 
-  for row = srow, erow do
-    on_line_impl(self, buf, row, true)
-  end
+  on_range_impl(self, buf, srow, 0, erow, 0, true)
   self._highlight_states = highlight_states
 end
 
@@ -399,7 +411,7 @@ end
 
 api.nvim_set_decoration_provider(ns, {
   on_win = TSHighlighter._on_win,
-  on_line = TSHighlighter._on_line,
+  on_range = TSHighlighter._on_range,
   _on_spell_nav = TSHighlighter._on_spell_nav,
 })
 
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index dbe3d54c2f..93d5ddf920 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -915,6 +915,41 @@ function Query:iter_captures(node, source, start, stop)
   return iter
 end
 
+function Query:iter_captures2(node, source, sr, sc, er, ec)
+  if type(source) == 'number' and source == 0 then
+    source = api.nvim_get_current_buf()
+  end
+
+  local cursor =
+    vim._create_ts_querycursor2(node, self.query, sr, sc, er, ec, { match_limit = 256 })
+
+  local apply_directives = memoize(match_id_hash, self.apply_directives, true)
+  local match_preds = memoize(match_id_hash, self.match_preds, true)
+
+  local function iter(end_line, end_col)
+    local capture, captured_node, match = cursor:next_capture()
+
+    if not capture then
+      return
+    end
+
+    if not match_preds(self, match, source) then
+      local match_id = match:info()
+      cursor:remove_match(match_id)
+      local row, col = captured_node:range()
+      if row > end_line or (row == end_line and col > end_col) then
+        return nil, captured_node, nil, nil
+      end
+      return iter(end_line, end_col) -- tail call: try next match
+    end
+
+    local metadata = apply_directives(self, match, source)
+
+    return capture, captured_node, metadata, match
+  end
+  return iter
+end
+
 --- Iterates the matches of self on a given range.
 ---
 --- Iterate over all matches within a {node}. The arguments are the same as for
diff --git a/src/nvim/api/extmark.c b/src/nvim/api/extmark.c
index c94b8df9ea..97945c2d7e 100644
--- a/src/nvim/api/extmark.c
+++ b/src/nvim/api/extmark.c
@@ -1049,6 +1049,7 @@ void nvim_set_decoration_provider(Integer ns_id, Dict(set_decoration_provider) *
     { "on_buf", &opts->on_buf, &p->redraw_buf },
     { "on_win", &opts->on_win, &p->redraw_win },
     { "on_line", &opts->on_line, &p->redraw_line },
+    { "on_range", &opts->on_range, &p->redraw_range },
     { "on_end", &opts->on_end, &p->redraw_end },
     { "_on_hl_def", &opts->_on_hl_def, &p->hl_def },
     { "_on_spell_nav", &opts->_on_spell_nav, &p->spell_nav },
diff --git a/src/nvim/api/keysets_defs.h b/src/nvim/api/keysets_defs.h
index 48f5f7246c..f6c822b29b 100644
--- a/src/nvim/api/keysets_defs.h
+++ b/src/nvim/api/keysets_defs.h
@@ -17,6 +17,8 @@ typedef struct {
   LuaRefOf(("buf" _, Integer bufnr, Integer tick)) on_buf;
   LuaRefOf(("win" _, Integer winid, Integer bufnr, Integer toprow, Integer botrow)) on_win;
   LuaRefOf(("line" _, Integer winid, Integer bufnr, Integer row)) on_line;
+  LuaRefOf(("range" _, Integer winid, Integer bufnr, Integer sr, Integer sc, Integer er,
+            Integer ec)) on_range;
   LuaRefOf(("end" _, Integer tick)) on_end;
   LuaRefOf(("hl_def" _)) _on_hl_def;
   LuaRefOf(("spell_nav" _)) _on_spell_nav;
diff --git a/src/nvim/decoration_defs.h b/src/nvim/decoration_defs.h
index 58ba93a7ba..716fd34c4a 100644
--- a/src/nvim/decoration_defs.h
+++ b/src/nvim/decoration_defs.h
@@ -140,6 +140,7 @@ typedef struct {
   LuaRef redraw_buf;
   LuaRef redraw_win;
   LuaRef redraw_line;
+  LuaRef redraw_range;
   LuaRef redraw_end;
   LuaRef hl_def;
   LuaRef spell_nav;
diff --git a/src/nvim/decoration_provider.c b/src/nvim/decoration_provider.c
index e5d2658720..6560d28ee3 100644
--- a/src/nvim/decoration_provider.c
+++ b/src/nvim/decoration_provider.c
@@ -29,7 +29,7 @@ static kvec_t(DecorProvider) decor_providers = KV_INITIAL_VALUE;
 
 #define DECORATION_PROVIDER_INIT(ns_id) (DecorProvider) \
   { ns_id, kDecorProviderDisabled, LUA_NOREF, LUA_NOREF, \
-    LUA_NOREF, LUA_NOREF, LUA_NOREF, \
+    LUA_NOREF, LUA_NOREF, LUA_NOREF, LUA_NOREF, \
     LUA_NOREF, -1, false, false, 0 }
 
 static void decor_provider_error(DecorProvider *provider, const char *name, const char *msg)
@@ -178,6 +178,32 @@ void decor_providers_invoke_line(win_T *wp, int row, bool *has_decor)
   decor_state.running_decor_provider = false;
 }
 
+void decor_providers_invoke_range(win_T *wp, int sr, int sc, int er, int ec, bool *has_decor)
+{
+  decor_state.running_decor_provider = true;
+  for (size_t i = 0; i < kv_size(decor_providers); i++) {
+    DecorProvider *p = &kv_A(decor_providers, i);
+    if (p->state == kDecorProviderActive && p->redraw_range != LUA_NOREF) {
+      MAXSIZE_TEMP_ARRAY(args, 6);
+      ADD_C(args, WINDOW_OBJ(wp->handle));
+      ADD_C(args, BUFFER_OBJ(wp->w_buffer->handle));
+      ADD_C(args, INTEGER_OBJ(sr));
+      ADD_C(args, INTEGER_OBJ(sc));
+      ADD_C(args, INTEGER_OBJ(er));
+      ADD_C(args, INTEGER_OBJ(ec));
+      if (decor_provider_invoke((int)i, "range", p->redraw_range, args, true)) {
+        *has_decor = true;
+      } else {
+        // return 'false' or error: skip rest of this window
+        kv_A(decor_providers, i).state = kDecorProviderWinDisabled;
+      }
+
+      hl_check_ns();
+    }
+  }
+  decor_state.running_decor_provider = false;
+}
+
 /// For each provider invoke the 'buf' callback for a given buffer.
 ///
 /// @param      buf       Buffer
@@ -262,6 +288,7 @@ void decor_provider_clear(DecorProvider *p)
   NLUA_CLEAR_REF(p->redraw_buf);
   NLUA_CLEAR_REF(p->redraw_win);
   NLUA_CLEAR_REF(p->redraw_line);
+  NLUA_CLEAR_REF(p->redraw_range);
   NLUA_CLEAR_REF(p->redraw_end);
   NLUA_CLEAR_REF(p->spell_nav);
   p->state = kDecorProviderDisabled;
diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index b8f21d7454..731708a374 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -1027,6 +1027,9 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
   buf_T *buf = wp->w_buffer;
   const bool end_fill = (lnum == buf->b_ml.ml_line_count + 1);
 
+  int decor_provider_end_col;
+  bool check_decor_providers = false;
+
   if (col_rows == 0) {
     // To speed up the loop below, set extra_check when there is linebreak,
     // trailing white space and/or syntax processing to be done.
@@ -1052,11 +1055,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
     has_decor = decor_redraw_line(wp, lnum - 1, &decor_state);
 
     if (!end_fill) {
-      decor_providers_invoke_line(wp, lnum - 1, &has_decor);
-    }
-
-    if (has_decor) {
-      extra_check = true;
+      check_decor_providers = true;
     }
 
     // Check for columns to display for 'colorcolumn'.
@@ -1466,6 +1465,27 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
     }
   }
 
+  if (check_decor_providers) {
+    decor_provider_end_col = (int)(ptr - line);
+    int new_col = decor_provider_end_col + 500;
+
+    bool added_decor = false;
+    decor_providers_invoke_line(wp, lnum - 1, &added_decor);
+    has_decor |= added_decor;
+    decor_providers_invoke_range(wp,
+                                 lnum - 1,
+                                 decor_provider_end_col,
+                                 lnum - 1,
+                                 new_col,
+                                 &added_decor);
+    decor_provider_end_col = new_col;
+    has_decor |= added_decor;
+  }
+
+  if (has_decor) {
+    extra_check = true;
+  }
+
   // Correct highlighting for cursor that can't be disabled.
   // Avoids having to check this for each character.
   if (wlv.fromcol >= 0) {
@@ -1515,6 +1535,19 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
 
     bool did_decrement_ptr = false;
 
+    if (check_decor_providers && (int)(ptr - line) >= decor_provider_end_col) {
+      int const new_col = decor_provider_end_col + 500;
+
+      bool added_decor = false;
+      decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
+                                   lnum - 1, new_col, &added_decor);
+      decor_provider_end_col = new_col;
+      has_decor |= added_decor;
+    }
+    if (has_decor) {
+      extra_check = true;
+    }
+
     // Skip this quickly when working on the text.
     if (draw_cols) {
       if (cul_screenline) {
@@ -2551,6 +2584,17 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
 
     // At end of the text line.
     if (mb_schar == NUL) {
+      if (check_decor_providers) {
+        bool added_decor = false;
+        decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
+                                     lnum, 0, &added_decor);
+        decor_provider_end_col = INT_MAX;
+        has_decor |= added_decor;
+      }
+      if (has_decor) {
+        extra_check = true;
+      }
+
       // Highlight 'cursorcolumn' & 'colorcolumn' past end of the line.
 
       // check if line ends before left margin
diff --git a/src/nvim/lua/executor.c b/src/nvim/lua/executor.c
index c4fa8b0fff..4b65d4c268 100644
--- a/src/nvim/lua/executor.c
+++ b/src/nvim/lua/executor.c
@@ -1887,6 +1887,9 @@ static void nlua_add_treesitter(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
   lua_pushcfunction(lstate, tslua_push_querycursor);
   lua_setfield(lstate, -2, "_create_ts_querycursor");
 
+  lua_pushcfunction(lstate, tslua_push_querycursor2);
+  lua_setfield(lstate, -2, "_create_ts_querycursor2");
+
   lua_pushcfunction(lstate, tslua_add_language_from_object);
   lua_setfield(lstate, -2, "_ts_add_language_from_object");
 
diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index c80e7b7672..489cd1a0d6 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -1333,6 +1333,53 @@ int tslua_push_querycursor(lua_State *L)
   return 1;
 }
 
+int tslua_push_querycursor2(lua_State *L)
+{
+  TSNode node = node_check(L, 1);
+
+  TSQuery *query = query_check(L, 2);
+  TSQueryCursor *cursor = ts_query_cursor_new();
+  ts_query_cursor_exec(cursor, query, node);
+
+  uint32_t sr = (uint32_t)luaL_checkinteger(L, 3);
+  uint32_t sc = (uint32_t)luaL_checkinteger(L, 4);
+  uint32_t er = (uint32_t)luaL_checkinteger(L, 5);
+  uint32_t ec = (uint32_t)luaL_checkinteger(L, 6);
+
+  ts_query_cursor_set_point_range(cursor, (TSPoint){ sr, sc }, (TSPoint){ er, ec });
+
+  if (lua_gettop(L) >= 7 && !lua_isnil(L, 7)) {
+    luaL_argcheck(L, lua_istable(L, 7), 7, "table expected");
+    lua_pushnil(L);  // [dict, ..., nil]
+    while (lua_next(L, 7)) {
+      // [dict, ..., key, value]
+      if (lua_type(L, -2) == LUA_TSTRING) {
+        char *k = (char *)lua_tostring(L, -2);
+        if (strequal("max_start_depth", k)) {
+          uint32_t max_start_depth = (uint32_t)lua_tointeger(L, -1);
+          ts_query_cursor_set_max_start_depth(cursor, max_start_depth);
+        } else if (strequal("match_limit", k)) {
+          uint32_t match_limit = (uint32_t)lua_tointeger(L, -1);
+          ts_query_cursor_set_match_limit(cursor, match_limit);
+        }
+      }
+      // pop the value; lua_next will pop the key.
+      lua_pop(L, 1);  // [dict, ..., key]
+    }
+  }
+
+  TSQueryCursor **ud = lua_newuserdata(L, sizeof(*ud));  // [node, query, ..., udata]
+  *ud = cursor;
+  lua_getfield(L, LUA_REGISTRYINDEX, TS_META_QUERYCURSOR);  // [node, query, ..., udata, meta]
+  lua_setmetatable(L, -2);  // [node, query, ..., udata]
+
+  // Copy the fenv which contains the nodes tree.
+  lua_getfenv(L, 1);  // [udata, reftable]
+  lua_setfenv(L, -2);  // [udata]
+
+  return 1;
+}
+
 static int querycursor_remove_match(lua_State *L)
 {
   TSQueryCursor *cursor = querycursor_check(L, 1);

From d43690d82a49737cc977be54ea9528987f007f98 Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Sat, 30 Nov 2024 19:06:36 -0600
Subject: [PATCH 03/13] fix: check columns in root bounds

---
 runtime/lua/vim/treesitter/highlighter.lua | 142 ++++++++++++---------
 src/nvim/drawline.c                        |   4 +-
 2 files changed, 87 insertions(+), 59 deletions(-)

diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index 1cb61b7af4..4524f29d07 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -278,100 +278,128 @@ local function get_spell(capture_name)
   return nil, 0
 end
 
+local function cmp_eq(rowa, cola, rowb, colb)
+  return rowa == rowb and cola == colb
+end
+
+local function cmp_lt(rowa, cola, rowb, colb)
+  return rowa < rowb or (rowa == rowb and cola < colb)
+end
+
+local function cmp_lte(rowa, cola, rowb, colb)
+  return rowa < rowb or (rowa == rowb and cola <= colb)
+end
+
 ---@param self vim.treesitter.highlighter
 ---@param buf integer
----@param line integer
+---@param range_br integer
+---@param range_bc integer
+---@param range_er integer
+---@param range_ec integer
 ---@param is_spell_nav boolean
-local function on_range_impl(self, buf, from_line, from_col, until_line, until_col, is_spell_nav)
+local function on_range_impl(self, buf, range_br, range_bc, range_er, range_ec, is_spell_nav)
   self:for_each_highlight_state(function(state)
     local root_node = state.tstree:root()
-    local root_start_row, _, root_end_row, _ = root_node:range()
+    local root_br, root_bc, root_er, root_ec = root_node:range()
 
-    -- Only consider trees that contain this line
-    if root_start_row > until_line or root_end_row < from_line then
+    local root_intersects = false
+
+    local range_empty = cmp_eq(range_br, range_bc, range_er, range_ec)
+    local root_empty = cmp_eq(root_br, root_bc, root_er, root_ec)
+    if not range_empty and not root_empty then
+      root_intersects = cmp_lt(range_br, range_bc, root_er, root_ec)
+        and cmp_lt(root_br, root_bc, range_er, range_ec)
+    else
+      root_intersects = cmp_lte(range_br, range_bc, root_er, root_ec)
+        and cmp_lte(root_br, root_bc, range_er, range_ec)
+    end
+
+    -- Only consider trees within the visible range
+    if not root_intersects then
       return
     end
 
-    if
-      state.iter == nil
-      or state.next_row < from_line
-      or (state.next_row == from_line and state.next_col < from_col)
-    then
+    local next_row = state.next_row
+    local next_col = state.next_col
+
+    if state.iter == nil or cmp_lt(next_row, next_col, range_br, range_bc) then
       -- Mainly used to skip over folds
 
       -- TODO(lewis6991): Creating a new iterator loses the cached predicate results for query
       -- matches. Move this logic inside iter_captures() so we can maintain the cache.
       state.iter = state.highlighter_query
         :query()
-        :iter_captures2(root_node, self.bufnr, from_line, from_col, root_end_row + 1, 0)
+        :iter_captures2(root_node, self.bufnr, range_br, range_bc, root_er + 1, 0)
     end
 
-    while
-      until_line > state.next_row or (until_line == state.next_row and until_col > state.next_col)
-    do
-      local capture, node, metadata, match = state.iter(until_line, until_col)
-
-      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])
+    while cmp_lt(next_row, next_col, range_er, range_ec) do
+      local capture, node, metadata, match = state.iter(range_er, range_ec)
+      if not node then
+        next_row = math.huge
+        next_col = math.huge
+        break
       end
+
+      local range = vim.treesitter.get_range(node, buf, metadata and metadata[capture])
       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)
-
-        local capture_name = state.highlighter_query:query().captures[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 = (
-          tonumber(metadata.priority or metadata[capture] and metadata[capture].priority)
-          or vim.hl.priorities.treesitter
-        ) + spell_pri_offset
-
-        -- 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
-
-        local url = get_url(match, buf, capture, metadata)
-
-        if hl and end_row >= from_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 > until_line or (start_row == until_line and start_col >= until_col) then
-        state.next_row = start_row
-        state.next_col = start_col
+      if cmp_lt(next_row, next_col, start_row, start_col) then
+        next_row = start_row
+        next_col = start_col
       end
 
       if not capture then
         break
       end
+
+      local hl = state.highlighter_query:get_hl_from_capture(capture)
+      local capture_name = state.highlighter_query:query().captures[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 = (
+        tonumber(metadata.priority or metadata[capture] and metadata[capture].priority)
+        or vim.hl.priorities.treesitter
+      ) + spell_pri_offset
+
+      -- 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
+
+      local url = get_url(match, buf, capture, metadata)
+
+      if 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,
+          ephemeral = true,
+          priority = priority,
+          hl_group = hl,
+          conceal = conceal,
+          spell = spell,
+          url = url,
+        })
+      end
     end
+
+    state.next_row = next_row
+    state.next_col = next_col
   end)
 end
 
 ---@private
 ---@param _win integer
 ---@param buf integer
----@param line integer
-function TSHighlighter._on_range(_, _win, buf, sr, sc, er, ec, _)
+---@param br integer
+---@param bc integer
+---@param er integer
+---@param ec integer
+function TSHighlighter._on_range(_, _win, buf, br, bc, er, ec, _)
   local self = TSHighlighter.active[buf]
   if not self then
     return
   end
 
-  on_range_impl(self, buf, sr, sc, er, ec, false)
+  on_range_impl(self, buf, br, bc, er, ec, false)
 end
 
 ---@private
diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index 731708a374..07bd8265ae 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -2588,8 +2588,8 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
         bool added_decor = false;
         decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
                                      lnum, 0, &added_decor);
-        decor_provider_end_col = INT_MAX;
-        has_decor |= added_decor;
+        // decor_provider_end_col = INT_MAX;
+        // has_decor |= added_decor;
       }
       if (has_decor) {
         extra_check = true;

From 112a5cf67985749a1102be288ae00c7f17a9a603 Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Sat, 30 Nov 2024 20:46:48 -0600
Subject: [PATCH 04/13] feat: approximate decor provider range

---
 runtime/lua/vim/treesitter/highlighter.lua |  3 +-
 src/nvim/drawline.c                        | 58 ++++++++++++----------
 2 files changed, 34 insertions(+), 27 deletions(-)

diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index 4524f29d07..da3b1dd017 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -302,7 +302,8 @@ local function on_range_impl(self, buf, range_br, range_bc, range_er, range_ec,
     local root_node = state.tstree:root()
     local root_br, root_bc, root_er, root_ec = root_node:range()
 
-    local root_intersects = false
+    --- @type boolean
+    local root_intersects
 
     local range_empty = cmp_eq(range_br, range_bc, range_er, range_ec)
     local root_empty = cmp_eq(root_br, root_bc, root_er, root_ec)
diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index 07bd8265ae..33bfbfd167 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -1466,20 +1466,30 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
   }
 
   if (check_decor_providers) {
+    int line_len = ml_get_buf_len(buf, lnum);
     decor_provider_end_col = (int)(ptr - line);
-    int new_col = decor_provider_end_col + 500;
+    int rem_bytes = line_len - decor_provider_end_col;
 
+    // Approximate the number of bytes that will be drawn.
+    // Assume we're dealing with 1-cell ascii and ignore
+    // the effects of 'linebreak', 'breakindent', etc.
+    int rem_vcols;
+    if (wp->w_p_wrap) {
+      int width = wp->w_width_inner - win_col_off(wp);
+      int width2 = width + win_col_off2(wp);
+      rem_vcols = width + (endrow - startrow) * width2;
+    } else {
+      rem_vcols = wp->w_width_inner - win_col_off(wp);
+    }
+
+    int const new_col = decor_provider_end_col + MIN(rem_bytes, rem_vcols + 1);
     bool added_decor = false;
     decor_providers_invoke_line(wp, lnum - 1, &added_decor);
     has_decor |= added_decor;
-    decor_providers_invoke_range(wp,
-                                 lnum - 1,
-                                 decor_provider_end_col,
-                                 lnum - 1,
-                                 new_col,
-                                 &added_decor);
-    decor_provider_end_col = new_col;
+    decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
+                                 lnum - 1, new_col, &added_decor);
     has_decor |= added_decor;
+    decor_provider_end_col = new_col;
   }
 
   if (has_decor) {
@@ -1536,13 +1546,20 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
     bool did_decrement_ptr = false;
 
     if (check_decor_providers && (int)(ptr - line) >= decor_provider_end_col) {
-      int const new_col = decor_provider_end_col + 500;
-
-      bool added_decor = false;
-      decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
-                                   lnum - 1, new_col, &added_decor);
-      decor_provider_end_col = new_col;
-      has_decor |= added_decor;
+      if (*ptr == NUL) {
+        bool added_decor = false;
+        decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
+                                     lnum, 0, &added_decor);
+        has_decor |= added_decor;
+        decor_provider_end_col = INT_MAX;
+      } else {
+        int const new_col = decor_provider_end_col + 100;
+        bool added_decor = false;
+        decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
+                                     lnum - 1, new_col, &added_decor);
+        has_decor |= added_decor;
+        decor_provider_end_col = new_col;
+      }
     }
     if (has_decor) {
       extra_check = true;
@@ -2584,17 +2601,6 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
 
     // At end of the text line.
     if (mb_schar == NUL) {
-      if (check_decor_providers) {
-        bool added_decor = false;
-        decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
-                                     lnum, 0, &added_decor);
-        // decor_provider_end_col = INT_MAX;
-        // has_decor |= added_decor;
-      }
-      if (has_decor) {
-        extra_check = true;
-      }
-
       // Highlight 'cursorcolumn' & 'colorcolumn' past end of the line.
 
       // check if line ends before left margin

From f94e98556e378b15ad976ab08b3132aa5e89b2e0 Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Sat, 30 Nov 2024 21:35:35 -0600
Subject: [PATCH 05/13] fix: invalidate line after calling decor providers

---
 src/nvim/drawline.c | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index 33bfbfd167..3ab9fc193d 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -1466,9 +1466,8 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
   }
 
   if (check_decor_providers) {
-    int line_len = ml_get_buf_len(buf, lnum);
-    decor_provider_end_col = (int)(ptr - line);
-    int rem_bytes = line_len - decor_provider_end_col;
+    int const line_len = ml_get_buf_len(buf, lnum);
+    int const col = (int)(ptr - line);
 
     // Approximate the number of bytes that will be drawn.
     // Assume we're dealing with 1-cell ascii and ignore
@@ -1482,14 +1481,16 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
       rem_vcols = wp->w_width_inner - win_col_off(wp);
     }
 
-    int const new_col = decor_provider_end_col + MIN(rem_bytes, rem_vcols + 1);
+    int const new_col = col + MIN(line_len - col, rem_vcols + 1);
     bool added_decor = false;
     decor_providers_invoke_line(wp, lnum - 1, &added_decor);
     has_decor |= added_decor;
-    decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
-                                 lnum - 1, new_col, &added_decor);
+    decor_providers_invoke_range(wp, lnum - 1, col, lnum - 1, new_col, &added_decor);
     has_decor |= added_decor;
     decor_provider_end_col = new_col;
+
+    line = ml_get_buf(wp->w_buffer, lnum);
+    ptr = line + col;
   }
 
   if (has_decor) {
@@ -1546,6 +1547,8 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
     bool did_decrement_ptr = false;
 
     if (check_decor_providers && (int)(ptr - line) >= decor_provider_end_col) {
+      int const col = (int)(ptr - line);
+
       if (*ptr == NUL) {
         bool added_decor = false;
         decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
@@ -1560,6 +1563,9 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
         has_decor |= added_decor;
         decor_provider_end_col = new_col;
       }
+
+      line = ml_get_buf(wp->w_buffer, lnum);
+      ptr = line + col;
     }
     if (has_decor) {
       extra_check = true;

From b730a4c6250d6c8cb7ae702f4ec1827d69e71142 Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Sat, 7 Dec 2024 23:45:28 -0600
Subject: [PATCH 06/13] fix: make range estimation more optimal

---
 src/nvim/drawline.c | 43 +++++++++++++++++++++++++++++--------------
 1 file changed, 29 insertions(+), 14 deletions(-)

diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index 3ab9fc193d..fc42b6fc3f 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -1466,7 +1466,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
   }
 
   if (check_decor_providers) {
-    int const line_len = ml_get_buf_len(buf, lnum);
+    int const line_len = ml_get_buf_len(wp->w_buffer, lnum);
     int const col = (int)(ptr - line);
 
     // Approximate the number of bytes that will be drawn.
@@ -1476,18 +1476,30 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
     if (wp->w_p_wrap) {
       int width = wp->w_width_inner - win_col_off(wp);
       int width2 = width + win_col_off2(wp);
-      rem_vcols = width + (endrow - startrow) * width2;
+
+      int first_row_width = start_col == 0 ? width : width2;
+      rem_vcols = first_row_width + (endrow - startrow - 1) * width2;
     } else {
       rem_vcols = wp->w_width_inner - win_col_off(wp);
     }
+    int expected_end = col + rem_vcols + 1;
 
-    int const new_col = col + MIN(line_len - col, rem_vcols + 1);
+    // Call it here since we need to invalidate the line pointer anyway.
     bool added_decor = false;
     decor_providers_invoke_line(wp, lnum - 1, &added_decor);
     has_decor |= added_decor;
-    decor_providers_invoke_range(wp, lnum - 1, col, lnum - 1, new_col, &added_decor);
-    has_decor |= added_decor;
-    decor_provider_end_col = new_col;
+
+    if (line_len > expected_end) {
+      decor_providers_invoke_range(wp, lnum - 1, col, lnum - 1, 
+                                   expected_end, &added_decor);
+      has_decor |= added_decor;
+      decor_provider_end_col = expected_end;
+    }
+    else {
+      decor_providers_invoke_range(wp, lnum - 1, col, lnum, 0, &added_decor);
+      has_decor |= added_decor;
+      decor_provider_end_col = INT_MAX;
+    }
 
     line = ml_get_buf(wp->w_buffer, lnum);
     ptr = line + col;
@@ -1547,21 +1559,24 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
     bool did_decrement_ptr = false;
 
     if (check_decor_providers && (int)(ptr - line) >= decor_provider_end_col) {
+      int const line_len = ml_get_buf_len(wp->w_buffer, lnum);
       int const col = (int)(ptr - line);
 
-      if (*ptr == NUL) {
+      int const col_increment = 100;
+      if (line_len - decor_provider_end_col > col_increment) {
+        int const new_col = decor_provider_end_col + col_increment;
+        bool added_decor = false;
+        decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col, 
+                                     lnum - 1, new_col, &added_decor);
+        has_decor |= added_decor;
+        decor_provider_end_col = new_col;
+      }
+      else {
         bool added_decor = false;
         decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
                                      lnum, 0, &added_decor);
         has_decor |= added_decor;
         decor_provider_end_col = INT_MAX;
-      } else {
-        int const new_col = decor_provider_end_col + 100;
-        bool added_decor = false;
-        decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
-                                     lnum - 1, new_col, &added_decor);
-        has_decor |= added_decor;
-        decor_provider_end_col = new_col;
       }
 
       line = ml_get_buf(wp->w_buffer, lnum);

From fb58bc8da5e7a83eb8270485227c98f4b8e50c37 Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Sun, 8 Dec 2024 22:12:46 -0600
Subject: [PATCH 07/13] refactor: improve API

---
 runtime/lua/vim/treesitter/_meta/misc.lua  |   6 +-
 runtime/lua/vim/treesitter/highlighter.lua |  27 ++---
 runtime/lua/vim/treesitter/query.lua       |  58 ++++-------
 src/nvim/api/keysets_defs.h                |   4 +-
 src/nvim/drawline.c                        |  10 +-
 src/nvim/lua/executor.c                    |   3 -
 src/nvim/lua/treesitter.c                  | 116 +++++++++------------
 7 files changed, 87 insertions(+), 137 deletions(-)

diff --git a/runtime/lua/vim/treesitter/_meta/misc.lua b/runtime/lua/vim/treesitter/_meta/misc.lua
index 33701ef254..db6b9fd70d 100644
--- a/runtime/lua/vim/treesitter/_meta/misc.lua
+++ b/runtime/lua/vim/treesitter/_meta/misc.lua
@@ -71,8 +71,6 @@ function TSQueryCursor:next_match() end
 
 --- @param node TSNode
 --- @param query TSQuery
---- @param start integer?
---- @param stop integer?
---- @param opts? { max_start_depth?: integer, match_limit?: integer}
+--- @param opts? { row_begin?: integer, col_begin?: integer, row_end?: integer, col_end?: integer, max_start_depth?: integer, match_limit?: integer }
 --- @return TSQueryCursor
-function vim._create_ts_querycursor(node, query, start, stop, opts) end
+function vim._create_ts_querycursor(node, query, opts) end
diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index da3b1dd017..b6ce0112ef 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -1,10 +1,13 @@
 local api = vim.api
 local query = vim.treesitter.query
 local Range = require('vim.treesitter._range')
+local cmp_eq = Range.cmp_pos.eq
+local cmp_lt = Range.cmp_pos.lt
+local cmp_lte = Range.cmp_pos.le
 
 local ns = api.nvim_create_namespace('treesitter/highlighter')
 
----@alias vim.treesitter.highlighter.Iter fun(end_line: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch
+---@alias vim.treesitter.highlighter.Iter fun(end_line: integer|nil, end_col: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch
 
 ---@class (private) vim.treesitter.highlighter.Query
 ---@field private _query vim.treesitter.Query?
@@ -278,18 +281,6 @@ local function get_spell(capture_name)
   return nil, 0
 end
 
-local function cmp_eq(rowa, cola, rowb, colb)
-  return rowa == rowb and cola == colb
-end
-
-local function cmp_lt(rowa, cola, rowb, colb)
-  return rowa < rowb or (rowa == rowb and cola < colb)
-end
-
-local function cmp_lte(rowa, cola, rowb, colb)
-  return rowa < rowb or (rowa == rowb and cola <= colb)
-end
-
 ---@param self vim.treesitter.highlighter
 ---@param buf integer
 ---@param range_br integer
@@ -328,9 +319,13 @@ local function on_range_impl(self, buf, range_br, range_bc, range_er, range_ec,
 
       -- TODO(lewis6991): Creating a new iterator loses the cached predicate results for query
       -- matches. Move this logic inside iter_captures() so we can maintain the cache.
-      state.iter = state.highlighter_query
-        :query()
-        :iter_captures2(root_node, self.bufnr, range_br, range_bc, root_er + 1, 0)
+      state.iter = state.highlighter_query:query():iter_captures(
+        root_node,
+        self.bufnr,
+        range_br,
+        root_er,
+        { col_begin = range_bc, col_end = root_ec }
+      )
     end
 
     while cmp_lt(next_row, next_col, range_er, range_ec) do
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 93d5ddf920..31d88fcb7f 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -1,6 +1,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 M = {}
 
@@ -875,53 +876,28 @@ end
 ---@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 opts? { col_begin?: integer, col_end?: integer } Optional parameters.
 ---
----@return (fun(end_line: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch):
+---@return (fun(end_line: integer|nil, end_col: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch):
 ---        capture id, capture node, metadata, match
 ---
 ---@note Captures are only returned if the query pattern of a specific capture contained predicates.
-function Query:iter_captures(node, source, start, stop)
+function Query:iter_captures(node, source, start, stop, opts)
   if type(source) == 'number' and source == 0 then
     source = api.nvim_get_current_buf()
   end
 
+  opts = opts or {}
   start, stop = value_or_node_range(start, stop, node)
 
-  local cursor = vim._create_ts_querycursor(node, self.query, start, stop, { match_limit = 256 })
-
-  local apply_directives = memoize(match_id_hash, self.apply_directives, true)
-  local match_preds = memoize(match_id_hash, self.match_preds, true)
-
-  local function iter(end_line)
-    local capture, captured_node, match = cursor:next_capture()
-
-    if not capture then
-      return
-    end
-
-    if not match_preds(self, match, source) then
-      local match_id = match:info()
-      cursor:remove_match(match_id)
-      if end_line and captured_node:range() > end_line then
-        return nil, captured_node, nil, nil
-      end
-      return iter(end_line) -- tail call: try next match
-    end
-
-    local metadata = apply_directives(self, match, source)
-
-    return capture, captured_node, metadata, match
-  end
-  return iter
-end
-
-function Query:iter_captures2(node, source, sr, sc, er, ec)
-  if type(source) == 'number' and source == 0 then
-    source = api.nvim_get_current_buf()
-  end
-
-  local cursor =
-    vim._create_ts_querycursor2(node, self.query, sr, sc, er, ec, { match_limit = 256 })
+  local cursor_opts = {
+    row_begin = start,
+    col_begin = opts.col_begin,
+    row_end = stop,
+    col_end = opts.col_end,
+    match_limit = 256,
+  }
+  local cursor = vim._create_ts_querycursor(node, self.query, cursor_opts)
 
   local apply_directives = memoize(match_id_hash, self.apply_directives, true)
   local match_preds = memoize(match_id_hash, self.match_preds, true)
@@ -937,7 +913,7 @@ function Query:iter_captures2(node, source, sr, sc, er, ec)
       local match_id = match:info()
       cursor:remove_match(match_id)
       local row, col = captured_node:range()
-      if row > end_line or (row == end_line and col > end_col) then
+      if end_line and cmp_ge(row, col, end_line, end_col or 0) then
         return nil, captured_node, nil, nil
       end
       return iter(end_line, end_col) -- tail call: try next match
@@ -996,8 +972,12 @@ function Query:iter_matches(node, source, start, stop, opts)
   end
 
   start, stop = value_or_node_range(start, stop, node)
+  opts.row_begin = start
+  opts.col_begin = 0
+  opts.row_end = stop
+  opts.col_end = 0
 
-  local cursor = vim._create_ts_querycursor(node, self.query, start, stop, opts)
+  local cursor = vim._create_ts_querycursor(node, self.query, opts)
 
   local function iter()
     local match = cursor:next_match()
diff --git a/src/nvim/api/keysets_defs.h b/src/nvim/api/keysets_defs.h
index f6c822b29b..317bb2e20f 100644
--- a/src/nvim/api/keysets_defs.h
+++ b/src/nvim/api/keysets_defs.h
@@ -17,8 +17,8 @@ typedef struct {
   LuaRefOf(("buf" _, Integer bufnr, Integer tick)) on_buf;
   LuaRefOf(("win" _, Integer winid, Integer bufnr, Integer toprow, Integer botrow)) on_win;
   LuaRefOf(("line" _, Integer winid, Integer bufnr, Integer row)) on_line;
-  LuaRefOf(("range" _, Integer winid, Integer bufnr, Integer sr, Integer sc, Integer er,
-            Integer ec)) on_range;
+  LuaRefOf(("range" _, Integer winid, Integer bufnr, Integer row_begin, Integer col_begin,
+            Integer row_end, Integer col_end)) on_range;
   LuaRefOf(("end" _, Integer tick)) on_end;
   LuaRefOf(("hl_def" _)) _on_hl_def;
   LuaRefOf(("spell_nav" _)) _on_spell_nav;
diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index fc42b6fc3f..1cc20cba15 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -1490,12 +1490,11 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
     has_decor |= added_decor;
 
     if (line_len > expected_end) {
-      decor_providers_invoke_range(wp, lnum - 1, col, lnum - 1, 
+      decor_providers_invoke_range(wp, lnum - 1, col, lnum - 1,
                                    expected_end, &added_decor);
       has_decor |= added_decor;
       decor_provider_end_col = expected_end;
-    }
-    else {
+    } else {
       decor_providers_invoke_range(wp, lnum - 1, col, lnum, 0, &added_decor);
       has_decor |= added_decor;
       decor_provider_end_col = INT_MAX;
@@ -1566,12 +1565,11 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
       if (line_len - decor_provider_end_col > col_increment) {
         int const new_col = decor_provider_end_col + col_increment;
         bool added_decor = false;
-        decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col, 
+        decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
                                      lnum - 1, new_col, &added_decor);
         has_decor |= added_decor;
         decor_provider_end_col = new_col;
-      }
-      else {
+      } else {
         bool added_decor = false;
         decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
                                      lnum, 0, &added_decor);
diff --git a/src/nvim/lua/executor.c b/src/nvim/lua/executor.c
index 4b65d4c268..c4fa8b0fff 100644
--- a/src/nvim/lua/executor.c
+++ b/src/nvim/lua/executor.c
@@ -1887,9 +1887,6 @@ static void nlua_add_treesitter(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
   lua_pushcfunction(lstate, tslua_push_querycursor);
   lua_setfield(lstate, -2, "_create_ts_querycursor");
 
-  lua_pushcfunction(lstate, tslua_push_querycursor2);
-  lua_setfield(lstate, -2, "_create_ts_querycursor2");
-
   lua_pushcfunction(lstate, tslua_add_language_from_object);
   lua_setfield(lstate, -2, "_ts_add_language_from_object");
 
diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index 489cd1a0d6..14eb825485 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -1287,87 +1287,69 @@ static struct luaL_Reg querycursor_meta[] = {
   { NULL, NULL }
 };
 
+static inline bool set_uint32(lua_State *L, int idx, char const *name, uint32_t *res)
+  FUNC_ATTR_ALWAYS_INLINE FUNC_ATTR_NONNULL_ALL
+{
+  lua_getfield(L, idx, name);
+
+  if (!lua_isnil(L, -1)) {
+    int64_t value = lua_tointeger(L, -1);
+    if (value < 0U) {
+      *res = 0U;
+      return true;
+    } else if (value > UINT32_MAX) {
+      *res = UINT32_MAX;
+      return true;
+    } else {
+      *res = (uint32_t)value;
+      return true;
+    }
+  }
+  lua_pop(L, 1);
+
+  return false;
+}
+
 int tslua_push_querycursor(lua_State *L)
 {
   TSNode node = node_check(L, 1);
 
   TSQuery *query = query_check(L, 2);
   TSQueryCursor *cursor = ts_query_cursor_new();
-  ts_query_cursor_exec(cursor, query, node);
 
-  if (lua_gettop(L) >= 3) {
-    uint32_t start = (uint32_t)luaL_checkinteger(L, 3);
-    uint32_t end = lua_gettop(L) >= 4 ? (uint32_t)luaL_checkinteger(L, 4) : MAXLNUM;
-    ts_query_cursor_set_point_range(cursor, (TSPoint){ start, 0 }, (TSPoint){ end, 0 });
-  }
+  if (lua_gettop(L) >= 3 && !lua_isnil(L, 3)) {
+    luaL_argcheck(L, lua_istable(L, 3), 3, "table expected");
 
-  if (lua_gettop(L) >= 5 && !lua_isnil(L, 5)) {
-    luaL_argcheck(L, lua_istable(L, 5), 5, "table expected");
-    lua_pushnil(L);  // [dict, ..., nil]
-    while (lua_next(L, 5)) {
-      // [dict, ..., key, value]
-      if (lua_type(L, -2) == LUA_TSTRING) {
-        char *k = (char *)lua_tostring(L, -2);
-        if (strequal("max_start_depth", k)) {
-          uint32_t max_start_depth = (uint32_t)lua_tointeger(L, -1);
-          ts_query_cursor_set_max_start_depth(cursor, max_start_depth);
-        } else if (strequal("match_limit", k)) {
-          uint32_t match_limit = (uint32_t)lua_tointeger(L, -1);
-          ts_query_cursor_set_match_limit(cursor, match_limit);
-        }
+    TSPoint begin = (TSPoint){ 0, 0 };
+    TSPoint end = (TSPoint){ UINT32_MAX, UINT32_MAX };
+
+    if (set_uint32(L, 3, "row_begin", &begin.row)) {
+      if (!set_uint32(L, 3, "col_begin", &begin.column)) {
+        begin.column = 0;
       }
-      // pop the value; lua_next will pop the key.
-      lua_pop(L, 1);  // [dict, ..., key]
+    }
+
+    if (set_uint32(L, 3, "row_end", &end.row)) {
+      if (!set_uint32(L, 3, "col_end", &end.column)) {
+        end.column = 0;
+      }
+    }
+
+    ts_query_cursor_set_point_range(cursor, begin, end);
+
+    uint32_t max_start_depth;
+    if (set_uint32(L, 3, "max_start_depth", &max_start_depth)) {
+      ts_query_cursor_set_max_start_depth(cursor, max_start_depth);
+    }
+
+    uint32_t match_limit;
+    if (set_uint32(L, 3, "match_limit", &match_limit)) {
+      ts_query_cursor_set_match_limit(cursor, match_limit);
     }
   }
 
-  TSQueryCursor **ud = lua_newuserdata(L, sizeof(*ud));  // [node, query, ..., udata]
-  *ud = cursor;
-  lua_getfield(L, LUA_REGISTRYINDEX, TS_META_QUERYCURSOR);  // [node, query, ..., udata, meta]
-  lua_setmetatable(L, -2);  // [node, query, ..., udata]
-
-  // Copy the fenv which contains the nodes tree.
-  lua_getfenv(L, 1);  // [udata, reftable]
-  lua_setfenv(L, -2);  // [udata]
-
-  return 1;
-}
-
-int tslua_push_querycursor2(lua_State *L)
-{
-  TSNode node = node_check(L, 1);
-
-  TSQuery *query = query_check(L, 2);
-  TSQueryCursor *cursor = ts_query_cursor_new();
   ts_query_cursor_exec(cursor, query, node);
 
-  uint32_t sr = (uint32_t)luaL_checkinteger(L, 3);
-  uint32_t sc = (uint32_t)luaL_checkinteger(L, 4);
-  uint32_t er = (uint32_t)luaL_checkinteger(L, 5);
-  uint32_t ec = (uint32_t)luaL_checkinteger(L, 6);
-
-  ts_query_cursor_set_point_range(cursor, (TSPoint){ sr, sc }, (TSPoint){ er, ec });
-
-  if (lua_gettop(L) >= 7 && !lua_isnil(L, 7)) {
-    luaL_argcheck(L, lua_istable(L, 7), 7, "table expected");
-    lua_pushnil(L);  // [dict, ..., nil]
-    while (lua_next(L, 7)) {
-      // [dict, ..., key, value]
-      if (lua_type(L, -2) == LUA_TSTRING) {
-        char *k = (char *)lua_tostring(L, -2);
-        if (strequal("max_start_depth", k)) {
-          uint32_t max_start_depth = (uint32_t)lua_tointeger(L, -1);
-          ts_query_cursor_set_max_start_depth(cursor, max_start_depth);
-        } else if (strequal("match_limit", k)) {
-          uint32_t match_limit = (uint32_t)lua_tointeger(L, -1);
-          ts_query_cursor_set_match_limit(cursor, match_limit);
-        }
-      }
-      // pop the value; lua_next will pop the key.
-      lua_pop(L, 1);  // [dict, ..., key]
-    }
-  }
-
   TSQueryCursor **ud = lua_newuserdata(L, sizeof(*ud));  // [node, query, ..., udata]
   *ud = cursor;
   lua_getfield(L, LUA_REGISTRYINDEX, TS_META_QUERYCURSOR);  // [node, query, ..., udata, meta]

From 1e42a0e748a037539a97d8c92e4fcaa7348bc191 Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Mon, 9 Dec 2024 00:02:30 -0600
Subject: [PATCH 08/13] docs: generate docs

---
 runtime/doc/treesitter.txt            | 6 ++++--
 runtime/lua/vim/_meta/api_keysets.lua | 1 +
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index a0860c60a6..3df84ecc37 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -1348,7 +1348,7 @@ parse({lang}, {query})                          *vim.treesitter.query.parse()*
       • |vim.treesitter.query.get()|
 
                                                        *Query:iter_captures()*
-Query:iter_captures({node}, {source}, {start}, {stop})
+Query:iter_captures({node}, {source}, {start}, {stop}, {opts})
     Iterate over all captures from all matches inside {node}
 
     {source} is needed if the query contains predicates; then the caller must
@@ -1382,9 +1382,11 @@ Query:iter_captures({node}, {source}, {start}, {stop})
                   `node:start()`.
       • {stop}    (`integer?`) Stopping line for the search (end-exclusive).
                   Defaults to `node:end_()`.
+      • {opts}    (`{ col_begin?: integer, col_end?: integer }?`) Optional
+                  parameters.
 
     Return: ~
-        (`fun(end_line: integer?): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch`)
+        (`fun(end_line: integer?, end_col: integer?): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch`)
         capture id, capture node, metadata, match
 
                                                         *Query:iter_matches()*
diff --git a/runtime/lua/vim/_meta/api_keysets.lua b/runtime/lua/vim/_meta/api_keysets.lua
index e11dddb2d3..f478cd3511 100644
--- a/runtime/lua/vim/_meta/api_keysets.lua
+++ b/runtime/lua/vim/_meta/api_keysets.lua
@@ -231,6 +231,7 @@ error('Cannot require a meta file')
 --- @field on_buf? fun(_: "buf", bufnr: integer, tick: integer)
 --- @field on_win? fun(_: "win", winid: integer, bufnr: integer, toprow: integer, botrow: integer)
 --- @field on_line? fun(_: "line", winid: integer, bufnr: integer, row: integer)
+--- @field on_range? fun(_: "range", winid: integer, bufnr: integer, row_begin: integer, col_begin: integer, row_end: integer, col_end: integer)
 --- @field on_end? fun(_: "end", tick: integer)
 --- @field _on_hl_def? fun(_: "hl_def")
 --- @field _on_spell_nav? fun(_: "spell_nav")

From 7d2d0a55491dd69d71b217987835576784dc7b45 Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Mon, 9 Dec 2024 00:03:08 -0600
Subject: [PATCH 09/13] test: test iter_captures columns

---
 test/functional/treesitter/query_spec.lua | 41 +++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/test/functional/treesitter/query_spec.lua b/test/functional/treesitter/query_spec.lua
index 634f8af83d..a549df64af 100644
--- a/test/functional/treesitter/query_spec.lua
+++ b/test/functional/treesitter/query_spec.lua
@@ -717,6 +717,47 @@ void ui_refresh(void)
     eq({ { 1, 10, 1, 13 } }, ret)
   end)
 
+  it('iter_captures supports columns', function()
+    local txt = [[
+int aaa = 1, bbb = 2;
+int foo = 1, bar = 2;
+int baz = 3, qux = 4;
+int ccc = 1, ddd = 2;
+    ]]
+
+    local function test(opts)
+      local parser = vim.treesitter.get_string_parser(txt, 'c')
+
+      local nodes = {}
+      local query = vim.treesitter.query.parse('c', '((identifier) @foo)')
+      local root = assert(parser:parse()[1]:root())
+      local iter = query:iter_captures(root, txt, 1, 2, opts)
+
+      while true do
+        local capture, node = iter()
+        if not capture then
+          break
+        end
+        table.insert(nodes, { node:range() })
+      end
+
+      return nodes
+    end
+
+    local ret
+    ret = exec_lua(test, { col_begin = 7, col_end = 13 })
+    eq({ { 1, 13, 1, 16 }, { 2, 4, 2, 7 } }, ret)
+
+    ret = exec_lua(test, { col_begin = 7 })
+    eq({ { 1, 13, 1, 16 } }, ret)
+
+    ret = exec_lua(test, { col_end = 13 })
+    eq({ { 1, 4, 1, 7 }, { 1, 13, 1, 16 }, { 2, 4, 2, 7 } }, ret)
+
+    ret = exec_lua(test, {})
+    eq({ { 1, 4, 1, 7 }, { 1, 13, 1, 16 } }, ret)
+  end)
+
   it('fails to load queries', function()
     local function test(exp, cquery)
       eq(exp, pcall_err(exec_lua, "vim.treesitter.query.parse('c', ...)", cquery))

From e007e977d2b55efb0c9374d52f56014a3bd9a50e Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Fri, 13 Dec 2024 16:42:25 -0600
Subject: [PATCH 10/13] docs: update news

---
 runtime/doc/news.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 5e70d75cfd..47603acdb5 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -290,6 +290,7 @@ TREESITTER
 • |treesitter-directive-trim!| can trim all whitespace (not just empty lines)
   from both sides of a node.
 • |vim.treesitter.get_captures_at_pos()| now returns the `id` of each capture
+• |Query:iter_captures()| supports specifying starting and ending columns.
 
 TUI
 

From be6b7cdc5c2848d38e890d10289ab1728cc77af7 Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Fri, 13 Dec 2024 17:11:39 -0600
Subject: [PATCH 11/13] feat: on_range bounds are always on the first byte

---
 src/nvim/drawline.c | 56 +++++++++++++++++++++------------------------
 1 file changed, 26 insertions(+), 30 deletions(-)

diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index 1cc20cba15..4e7cb7578d 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -1466,7 +1466,6 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
   }
 
   if (check_decor_providers) {
-    int const line_len = ml_get_buf_len(wp->w_buffer, lnum);
     int const col = (int)(ptr - line);
 
     // Approximate the number of bytes that will be drawn.
@@ -1482,23 +1481,13 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
     } else {
       rem_vcols = wp->w_width_inner - win_col_off(wp);
     }
-    int expected_end = col + rem_vcols + 1;
 
     // Call it here since we need to invalidate the line pointer anyway.
     bool added_decor = false;
     decor_providers_invoke_line(wp, lnum - 1, &added_decor);
     has_decor |= added_decor;
 
-    if (line_len > expected_end) {
-      decor_providers_invoke_range(wp, lnum - 1, col, lnum - 1,
-                                   expected_end, &added_decor);
-      has_decor |= added_decor;
-      decor_provider_end_col = expected_end;
-    } else {
-      decor_providers_invoke_range(wp, lnum - 1, col, lnum, 0, &added_decor);
-      has_decor |= added_decor;
-      decor_provider_end_col = INT_MAX;
-    }
+    decor_provider_end_col = invoke_range_next(wp, lnum, col, rem_vcols + 1, &has_decor);
 
     line = ml_get_buf(wp->w_buffer, lnum);
     ptr = line + col;
@@ -1558,25 +1547,8 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
     bool did_decrement_ptr = false;
 
     if (check_decor_providers && (int)(ptr - line) >= decor_provider_end_col) {
-      int const line_len = ml_get_buf_len(wp->w_buffer, lnum);
       int const col = (int)(ptr - line);
-
-      int const col_increment = 100;
-      if (line_len - decor_provider_end_col > col_increment) {
-        int const new_col = decor_provider_end_col + col_increment;
-        bool added_decor = false;
-        decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
-                                     lnum - 1, new_col, &added_decor);
-        has_decor |= added_decor;
-        decor_provider_end_col = new_col;
-      } else {
-        bool added_decor = false;
-        decor_providers_invoke_range(wp, lnum - 1, decor_provider_end_col,
-                                     lnum, 0, &added_decor);
-        has_decor |= added_decor;
-        decor_provider_end_col = INT_MAX;
-      }
-
+      decor_provider_end_col = invoke_range_next(wp, lnum, col, 100, &has_decor);
       line = ml_get_buf(wp->w_buffer, lnum);
       ptr = line + col;
     }
@@ -3058,3 +3030,27 @@ static void wlv_put_linebuf(win_T *wp, const winlinevars_T *wlv, int endcol, boo
   grid_adjust(&grid, &row, &coloff);
   grid_put_linebuf(grid, row, coloff, startcol, endcol, clear_width, bg_attr, wlv->vcol - 1, flags);
 }
+
+/// @return New begin column, or INT_MAX.
+static int invoke_range_next(win_T *wp, int lnum, colnr_T begin_col, colnr_T col_off,
+                             bool *has_decor)
+{
+  char const *const line = ml_get_buf(wp->w_buffer, lnum);
+  int const line_len = ml_get_buf_len(wp->w_buffer, lnum);
+  col_off = MAX(col_off, 1);
+
+  bool added_decor = false;
+  colnr_T new_col;
+  if (col_off <= line_len - begin_col) {
+    int end_col = begin_col + col_off;
+    end_col += mb_off_next(line, line + end_col);
+    decor_providers_invoke_range(wp, lnum - 1, begin_col, lnum - 1, end_col, &added_decor);
+    new_col = end_col;
+  } else {
+    decor_providers_invoke_range(wp, lnum - 1, begin_col, lnum, 0, &added_decor);
+    new_col = INT_MAX;
+  }
+
+  *has_decor |= added_decor;
+  return new_col;
+}

From a61c1eee21e29118b393f8fac8cf1edb9b80f5f6 Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Fri, 13 Dec 2024 19:31:15 -0600
Subject: [PATCH 12/13] docs: add docs for on_range()

---
 runtime/doc/api.txt           | 23 ++++++++++++++++-------
 runtime/doc/news.txt          |  1 +
 runtime/lua/vim/_meta/api.lua | 15 ++++++++++++---
 src/nvim/api/extmark.c        | 15 ++++++++++++---
 4 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/runtime/doc/api.txt b/runtime/doc/api.txt
index cb3b2a3f77..ab05dc17b4 100644
--- a/runtime/doc/api.txt
+++ b/runtime/doc/api.txt
@@ -2825,11 +2825,12 @@ nvim_set_decoration_provider({ns_id}, {opts})
     Note: this function should not be called often. Rather, the callbacks
     themselves can be used to throttle unneeded callbacks. the `on_start`
     callback can return `false` to disable the provider until the next redraw.
-    Similarly, return `false` in `on_win` will skip the `on_line` calls for
-    that window (but any extmarks set in `on_win` will still be used). A
-    plugin managing multiple sources of decoration should ideally only set one
-    provider, and merge the sources internally. You can use multiple `ns_id`
-    for the extmarks set/modified inside the callback anyway.
+    Similarly, return `false` in `on_win` will skip the `on_line` and
+    `on_range` calls for that window (but any extmarks set in `on_win` will
+    still be used). A plugin managing multiple sources of decoration should
+    ideally only set one provider, and merge the sources internally. You can
+    use multiple `ns_id` for the extmarks set/modified inside the callback
+    anyway.
 
     Note: doing anything other than setting extmarks is considered
     experimental. Doing things like changing options are not explicitly
@@ -2837,8 +2838,8 @@ nvim_set_decoration_provider({ns_id}, {opts})
     consumption). Doing `vim.rpcnotify` should be OK, but `vim.rpcrequest` is
     quite dubious for the moment.
 
-    Note: It is not allowed to remove or update extmarks in `on_line`
-    callbacks.
+    Note: It is not allowed to remove or update extmarks in `on_line` and
+    `on_range` callbacks.
 
     Attributes: ~
         Lua |vim.api| only
@@ -2859,6 +2860,14 @@ nvim_set_decoration_provider({ns_id}, {opts})
                  • on_line: called for each buffer line being redrawn. (The
                    interaction with fold lines is subject to change) >
                     ["line", winid, bufnr, row]
+<
+                 • on_range: called for each buffer range being redrawn. Range
+                   is end-exclusive and may span multiple lines. Range bounds
+                   point to the first byte of a character. An end position of
+                   the form (lnum, 0), including (number of lines, 0), is
+                   valid and indicates that EOL of the preceding line is
+                   included. >
+                    ["range", winid, bufnr, begin_row, begin_col, end_row, end_col]
 <
                  • on_end: called at the end of a redraw cycle >
                     ["end", tick]
diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 47603acdb5..9898160775 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -172,6 +172,7 @@ API
 
 • |nvim__ns_set()| can set properties for a namespace
 • |vim.json.encode()| has an option to enable forward slash escaping
+• Decoration provider has `on_range()` callback.
 
 DEFAULTS
 
diff --git a/runtime/lua/vim/_meta/api.lua b/runtime/lua/vim/_meta/api.lua
index b2385197bd..1a9d8921d1 100644
--- a/runtime/lua/vim/_meta/api.lua
+++ b/runtime/lua/vim/_meta/api.lua
@@ -2116,8 +2116,8 @@ function vim.api.nvim_set_current_win(window) end
 --- Note: this function should not be called often. Rather, the callbacks
 --- themselves can be used to throttle unneeded callbacks. the `on_start`
 --- callback can return `false` to disable the provider until the next redraw.
---- Similarly, return `false` in `on_win` will skip the `on_line` calls
---- for that window (but any extmarks set in `on_win` will still be used).
+--- Similarly, return `false` in `on_win` will skip the `on_line` and `on_range`
+--- calls for that window (but any extmarks set in `on_win` will still be used).
 --- A plugin managing multiple sources of decoration should ideally only set
 --- one provider, and merge the sources internally. You can use multiple `ns_id`
 --- for the extmarks set/modified inside the callback anyway.
@@ -2128,7 +2128,8 @@ function vim.api.nvim_set_current_win(window) end
 --- Doing `vim.rpcnotify` should be OK, but `vim.rpcrequest` is quite dubious
 --- for the moment.
 ---
---- Note: It is not allowed to remove or update extmarks in `on_line` callbacks.
+--- Note: It is not allowed to remove or update extmarks
+--- in `on_line` and `on_range` callbacks.
 ---
 --- @param ns_id integer Namespace id from `nvim_create_namespace()`
 --- @param opts vim.api.keyset.set_decoration_provider Table of callbacks:
@@ -2150,6 +2151,14 @@ function vim.api.nvim_set_current_win(window) end
 ---   ```
 ---     ["line", winid, bufnr, row]
 ---   ```
+--- - on_range: called for each buffer range being redrawn.
+---   Range is end-exclusive and may span multiple lines. Range
+---   bounds point to the first byte of a character. An end position
+---   of the form (lnum, 0), including (number of lines, 0), is valid
+---   and indicates that EOL of the preceding line is included.
+---   ```
+---     ["range", winid, bufnr, begin_row, begin_col, end_row, end_col]
+---   ```
 --- - on_end: called at the end of a redraw cycle
 ---   ```
 ---     ["end", tick]
diff --git a/src/nvim/api/extmark.c b/src/nvim/api/extmark.c
index 97945c2d7e..0c7d32c0e5 100644
--- a/src/nvim/api/extmark.c
+++ b/src/nvim/api/extmark.c
@@ -992,8 +992,8 @@ void nvim_buf_clear_namespace(Buffer buffer, Integer ns_id, Integer line_start,
 /// Note: this function should not be called often. Rather, the callbacks
 /// themselves can be used to throttle unneeded callbacks. the `on_start`
 /// callback can return `false` to disable the provider until the next redraw.
-/// Similarly, return `false` in `on_win` will skip the `on_line` calls
-/// for that window (but any extmarks set in `on_win` will still be used).
+/// Similarly, return `false` in `on_win` will skip the `on_line` and `on_range`
+/// calls for that window (but any extmarks set in `on_win` will still be used).
 /// A plugin managing multiple sources of decoration should ideally only set
 /// one provider, and merge the sources internally. You can use multiple `ns_id`
 /// for the extmarks set/modified inside the callback anyway.
@@ -1004,7 +1004,8 @@ void nvim_buf_clear_namespace(Buffer buffer, Integer ns_id, Integer line_start,
 /// Doing `vim.rpcnotify` should be OK, but `vim.rpcrequest` is quite dubious
 /// for the moment.
 ///
-/// Note: It is not allowed to remove or update extmarks in `on_line` callbacks.
+/// Note: It is not allowed to remove or update extmarks
+/// in `on_line` and `on_range` callbacks.
 ///
 /// @param ns_id  Namespace id from |nvim_create_namespace()|
 /// @param opts  Table of callbacks:
@@ -1026,6 +1027,14 @@ void nvim_buf_clear_namespace(Buffer buffer, Integer ns_id, Integer line_start,
 ///               ```
 ///                 ["line", winid, bufnr, row]
 ///               ```
+///             - on_range: called for each buffer range being redrawn.
+///               Range is end-exclusive and may span multiple lines. Range
+///               bounds point to the first byte of a character. An end position
+///               of the form (lnum, 0), including (number of lines, 0), is valid
+///               and indicates that EOL of the preceding line is included.
+///               ```
+///                 ["range", winid, bufnr, begin_row, begin_col, end_row, end_col]
+///               ```
 ///             - on_end: called at the end of a redraw cycle
 ///               ```
 ///                 ["end", tick]

From fab5cbad0042e46ed035550c66981990f2627018 Mon Sep 17 00:00:00 2001
From: vanaigr <vanaigranov@gmail.com>
Date: Sat, 14 Dec 2024 22:37:15 -0600
Subject: [PATCH 13/13] test: add on_range test

---
 test/functional/ui/decorations_spec.lua | 139 ++++++++++++++++++++----
 1 file changed, 117 insertions(+), 22 deletions(-)

diff --git a/test/functional/ui/decorations_spec.lua b/test/functional/ui/decorations_spec.lua
index da0fb9849a..1e702b1954 100644
--- a/test/functional/ui/decorations_spec.lua
+++ b/test/functional/ui/decorations_spec.lua
@@ -15,32 +15,36 @@ local eq = t.eq
 local assert_alive = n.assert_alive
 local pcall_err = t.pcall_err
 
+local function setup_screen(screen)
+  screen:set_default_attr_ids {
+    [1] = {bold=true, foreground=Screen.colors.Blue};
+    [2] = {foreground = Screen.colors.Grey100, background = Screen.colors.Red};
+    [3] = {foreground = Screen.colors.Brown};
+    [4] = {foreground = Screen.colors.Blue1};
+    [5] = {foreground = Screen.colors.Magenta};
+    [6] = {bold = true, foreground = Screen.colors.Brown};
+    [7] = {background = Screen.colors.Gray90};
+    [8] = {bold = true, reverse = true};
+    [9] = {reverse = true};
+    [10] = {italic = true, background = Screen.colors.Magenta};
+    [11] = {foreground = Screen.colors.Red, background = tonumber('0x005028')};
+    [12] = {foreground = tonumber('0x990000')};
+    [13] = {background = Screen.colors.LightBlue};
+    [14] = {background = Screen.colors.WebGray, foreground = Screen.colors.DarkBlue};
+    [15] = {special = Screen.colors.Blue, undercurl = true},
+    [16] = {special = Screen.colors.Red, undercurl = true},
+    [17] = {foreground = Screen.colors.Red},
+    [18] = {bold = true, foreground = Screen.colors.SeaGreen};
+    [19] = {bold = true};
+  }
+end
+
 describe('decorations providers', function()
   local screen
   before_each(function()
     clear()
     screen = Screen.new(40, 8)
-    screen:set_default_attr_ids {
-      [1] = {bold=true, foreground=Screen.colors.Blue};
-      [2] = {foreground = Screen.colors.Grey100, background = Screen.colors.Red};
-      [3] = {foreground = Screen.colors.Brown};
-      [4] = {foreground = Screen.colors.Blue1};
-      [5] = {foreground = Screen.colors.Magenta};
-      [6] = {bold = true, foreground = Screen.colors.Brown};
-      [7] = {background = Screen.colors.Gray90};
-      [8] = {bold = true, reverse = true};
-      [9] = {reverse = true};
-      [10] = {italic = true, background = Screen.colors.Magenta};
-      [11] = {foreground = Screen.colors.Red, background = tonumber('0x005028')};
-      [12] = {foreground = tonumber('0x990000')};
-      [13] = {background = Screen.colors.LightBlue};
-      [14] = {background = Screen.colors.WebGray, foreground = Screen.colors.DarkBlue};
-      [15] = {special = Screen.colors.Blue, undercurl = true},
-      [16] = {special = Screen.colors.Red, undercurl = true},
-      [17] = {foreground = Screen.colors.Red},
-      [18] = {bold = true, foreground = Screen.colors.SeaGreen};
-      [19] = {bold = true};
-    }
+    setup_screen(screen)
   end)
 
   local mulholland = [[
@@ -64,7 +68,7 @@ describe('decorations providers', function()
     ]]) .. [[
       api.nvim_set_decoration_provider(_G.ns1, {
         on_start = on_do; on_buf = on_do;
-        on_win = on_do; on_line = on_do;
+        on_win = on_do; on_line = on_do; on_range = on_do;
         on_end = on_do; _on_spell_nav = on_do;
       })
       return _G.ns1
@@ -107,12 +111,19 @@ describe('decorations providers', function()
       { "start", 4 };
       { "win", 1000, 1, 0, 6 };
       { "line", 1000, 1, 0 };
+      { "range", 1000, 1, 0, 0, 1, 0 };
       { "line", 1000, 1, 1 };
+      { "range", 1000, 1, 1, 0, 2, 0 };
       { "line", 1000, 1, 2 };
+      { "range", 1000, 1, 2, 0, 3, 0 };
       { "line", 1000, 1, 3 };
+      { "range", 1000, 1, 3, 0, 4, 0 };
       { "line", 1000, 1, 4 };
+      { "range", 1000, 1, 4, 0, 5, 0 };
       { "line", 1000, 1, 5 };
+      { "range", 1000, 1, 5, 0, 6, 0 };
       { "line", 1000, 1, 6 };
+      { "range", 1000, 1, 6, 0, 7, 0 };
       { "end", 4 };
     }
 
@@ -132,6 +143,7 @@ describe('decorations providers', function()
       { "buf", 1, 5 };
       { "win", 1000, 1, 0, 6 };
       { "line", 1000, 1, 6 };
+      { "range", 1000, 1, 6, 0, 7, 0 };
       { "end", 5 };
     }
   end)
@@ -199,9 +211,13 @@ describe('decorations providers', function()
       { "start", 5 };
       { "win", 1000, 1, 0, 3 };
       { "line", 1000, 1, 0 };
+      { "range", 1000, 1, 0, 0, 1, 0 };
       { "line", 1000, 1, 1 };
+      { "range", 1000, 1, 1, 0, 2, 0 };
       { "line", 1000, 1, 2 };
+      { "range", 1000, 1, 2, 0, 3, 0 };
       { "line", 1000, 1, 3 };
+      { "range", 1000, 1, 3, 0, 4, 0 };
       { "end", 5 };
     }
 
@@ -694,6 +710,85 @@ describe('decorations providers', function()
     ]])
   end)
 
+  it('on_range is invoked on all visible characters', function()
+    clear()
+    screen = Screen.new(20, 4)
+    setup_screen(screen)
+
+    local function record()
+      exec_lua [[
+        p_min = { math.huge, math.huge }
+        p_max = { -math.huge, -math.huge }
+        function pos_gt(a, b)
+          return a[1] > b[1] or (a[1] == b[1] and a[2] > b[2])
+        end
+        function pos_lt(a, b)
+          return a[1] < b[1] or (a[1] == b[1] and a[2] < b[2])
+        end
+      ]]
+      setup_provider [[
+        local function on_do(kind, _, bufnr, br, bc, er, ec)
+          if kind == 'range' then
+            local b = { br, bc }
+            local e = { er, ec }
+            if pos_gt(p_min, b) then
+              p_min = b
+            end
+            if pos_lt(p_max, e) then
+              p_max = e
+            end
+          end
+        end
+      ]]
+    end
+    local function check(min, max)
+      local function pos_gt(a, b)
+        return a[1] > b[1] or (a[1] == b[1] and a[2] > b[2])
+      end
+      local function pos_lt(a, b)
+        return a[1] < b[1] or (a[1] == b[1] and a[2] < b[2])
+      end
+      local p_min = exec_lua('return p_min')
+      assert(not pos_gt(p_min, min),
+        "minimum position " .. vim.inspect(p_min)
+        .. " should be before the first char")
+      local p_max = exec_lua('return p_max')
+      assert(not pos_lt(p_max, max),
+        "maximum position " .. vim.inspect(p_max)
+        .. " should be on or after the last char")
+    end
+
+    -- Multiple lines.
+    exec_lua([[
+      local lines = { ('a'):rep(40), ('b'):rep(40), ('c'):rep(40) }
+      vim.api.nvim_buf_set_lines(0, 0, -1, true, lines)
+      vim.api.nvim_win_set_cursor(0, { 2, 0 })
+    ]])
+    record()
+    screen:expect([[
+      ^bbbbbbbbbbbbbbbbbbbb|
+      bbbbbbbbbbbbbbbbbbbb|
+      ccccccccccccccccc{1:@@@}|
+                          |
+    ]])
+    check({ 1, 0 }, { 2, 21 })
+
+    -- One long line.
+    exec_lua([[
+      local lines = { ('a'):rep(100) }
+      vim.api.nvim_buf_set_lines(0, 0, -1, true, lines)
+      vim.api.nvim_win_set_cursor(0, { 1, 70 })
+    ]])
+    record()
+    screen:expect([[
+      {1:<<<}aaaaaaaaaaaaaaaaa|
+      aaaaaaaaaaaaaaaaaaaa|
+      aaaaaaaaaa^aaaaaaaaaa|
+                          |
+    ]])
+    check({ 0, 20 }, { 0, 81 })
+  end)
+
   it('errors gracefully', function()
     insert(mulholland)