Merge pull request #31400 from vanaigr/decor-provider-range

feat(decor): add range-based highlighting

5 files changed, 204 insertions, 98 deletions

diff --git a/runtime/lua/vim/_meta/api.lua b/runtime/lua/vim/_meta/api.lua
index b7132e5c40..cb552554cb 100644
--- a/runtime/lua/vim/_meta/api.lua
+++ b/runtime/lua/vim/_meta/api.lua
@@ -2128,7 +2128,7 @@ 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
+--- 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`
@@ -2140,7 +2140,7 @@ 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` or `on_range` callbacks.
 ---
 --- @param ns_id integer Namespace id from `nvim_create_namespace()`
 --- @param opts vim.api.keyset.set_decoration_provider Table of callbacks:
@@ -2162,6 +2162,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/runtime/lua/vim/_meta/api_keysets.lua b/runtime/lua/vim/_meta/api_keysets.lua
index 452a68eb50..a3c6b52409 100644
--- a/runtime/lua/vim/_meta/api_keysets.lua
+++ b/runtime/lua/vim/_meta/api_keysets.lua
@@ -382,6 +382,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): boolean?
 --- @field on_line? fun(_: "line", winid: integer, bufnr: integer, row: integer): boolean?
+--- @field on_range? fun(_: "range", winid: integer, bufnr: integer, start_row: integer, start_col: integer, end_row: integer, end_col: integer): boolean?
 --- @field on_end? fun(_: "end", tick: integer)
 --- @field _on_hl_def? fun(_: "hl_def")
 --- @field _on_spell_nav? fun(_: "spell_nav")
diff --git a/runtime/lua/vim/treesitter/_meta/misc.lua b/runtime/lua/vim/treesitter/_meta/misc.lua
index 971a1ceac2..8ed9b91ab5 100644
--- a/runtime/lua/vim/treesitter/_meta/misc.lua
+++ b/runtime/lua/vim/treesitter/_meta/misc.lua
@@ -80,8 +80,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? { start_row: integer, start_col: integer, end_row: integer, end_col: 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 9757a5f852..da13f7fe47 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -1,10 +1,11 @@
 local api = vim.api
 local query = vim.treesitter.query
 local Range = require('vim.treesitter._range')
+local cmp_lt = Range.cmp_pos.lt
 
 local ns = api.nvim_create_namespace('nvim.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, TSTree
 
 ---@class (private) vim.treesitter.highlighter.Query
 ---@field private _query vim.treesitter.Query?
@@ -57,6 +58,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
 ---@field prev_marks MarkInfo[]
@@ -233,6 +235,7 @@ function TSHighlighter:prepare_highlight_states(win, srow, erow)
     table.insert(self._highlight_states[win], {
       tstree = tstree,
       next_row = 0,
+      next_col = 0,
       iter = nil,
       highlighter_query = hl_query,
       prev_marks = {},
@@ -331,27 +334,36 @@ end
 ---Queues the remainder if the mark continues after the line.
 ---@param m MarkInfo
 ---@param buf integer
----@param line integer
+---@param range_start_row integer
+---@param range_start_col integer
+---@param range_end_row integer
+---@param range_end_col integer
 ---@param next_marks MarkInfo[]
-local function add_mark(m, buf, line, next_marks)
+local function add_mark(
+  m,
+  buf,
+  range_start_row,
+  range_start_col,
+  range_end_row,
+  range_end_col,
+  next_marks
+)
   local cur_start_l = m.start_line
   local cur_start_c = m.start_col
-  if cur_start_l < line then
-    cur_start_l = line
-    cur_start_c = 0
+  if cmp_lt(cur_start_l, cur_start_c, range_start_row, range_start_col) then
+    cur_start_l = range_start_row
+    cur_start_c = range_start_col
   end
 
   local cur_opts = m.opts
-  if cur_opts.end_line >= line + 1 then
+  if cmp_lt(range_end_row, range_end_col, cur_opts.end_line, cur_opts.end_col) then
     cur_opts = vim.deepcopy(cur_opts, true)
-    cur_opts.end_line = line + 1
-    cur_opts.end_col = 0
+    cur_opts.end_line = range_end_row
+    cur_opts.end_col = range_end_col
     table.insert(next_marks, m)
   end
 
-  local empty = cur_opts.end_line < cur_start_l
-    or (cur_opts.end_line == cur_start_l and cur_opts.end_col <= cur_start_c)
-  if cur_start_l <= line and not empty then
+  if cmp_lt(cur_start_l, cur_start_c, cur_opts.end_line, cur_opts.end_col) then
     api.nvim_buf_set_extmark(buf, ns, cur_start_l, cur_start_c, cur_opts)
   end
 end
@@ -359,17 +371,44 @@ end
 ---@param self vim.treesitter.highlighter
 ---@param win integer
 ---@param buf integer
----@param line integer
+---@param range_start_row integer
+---@param range_start_col integer
+---@param range_end_row integer
+---@param range_end_col integer
 ---@param on_spell boolean
 ---@param on_conceal boolean
-local function on_line_impl(self, win, buf, line, on_spell, on_conceal)
-  self._conceal_checked[line] = self._conceal_line and true or nil
+local function on_range_impl(
+  self,
+  win,
+  buf,
+  range_start_row,
+  range_start_col,
+  range_end_row,
+  range_end_col,
+  on_spell,
+  on_conceal
+)
+  if self._conceal_line then
+    range_start_col = 0
+    if range_end_col ~= 0 then
+      range_end_row = range_end_row + 1
+      range_end_col = 0
+    end
+  end
+  for i = range_start_row, range_end_row - 1 do
+    self._conceal_checked[i] = self._conceal_line or nil
+  end
   self:for_each_highlight_state(win, 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
+    ---@type { [1]: integer, [2]: integer, [3]: integer, [4]: integer }
+    local root_range = { root_node:range() }
+
+    if
+      not Range.intercepts(
+        root_range,
+        { range_start_row, range_start_col, range_end_row, range_end_col }
+      )
+    then
       return
     end
 
@@ -378,85 +417,114 @@ local function on_line_impl(self, win, buf, line, on_spell, on_conceal)
     local next_marks = {}
 
     for _, mark in ipairs(state.prev_marks) do
-      add_mark(mark, buf, line, next_marks)
+      add_mark(
+        mark,
+        buf,
+        range_start_row,
+        range_start_col,
+        range_end_row,
+        range_end_col,
+        next_marks
+      )
     end
 
-    if state.iter == nil or state.next_row < line 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_start_row, range_start_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_captures(
+        root_node,
+        self.bufnr,
+        range_start_row,
+        root_range[3],
+        { start_col = range_start_col, end_col = root_range[4] }
+      )
     end
 
     local captures = state.highlighter_query:query().captures
 
-    while line >= state.next_row do
-      local capture, node, metadata, match = state.iter(line)
+    while cmp_lt(next_row, next_col, range_end_row, range_end_col) do
+      local capture, node, metadata, match = state.iter(range_end_row, range_end_col)
+      if not node then
+        next_row = math.huge
+        next_col = math.huge
+        break
+      end
+
+      local outer_range = vim.treesitter.get_range(node, buf, metadata and metadata[capture])
+      if cmp_lt(next_row, next_col, outer_range[1], outer_range[2]) then
+        next_row = outer_range[1]
+        next_col = outer_range[2]
+      end
 
-      local outer_range = { root_end_row + 1, 0, root_end_row + 1, 0 }
-      if node then
-        outer_range = vim.treesitter.get_range(node, buf, metadata and metadata[capture])
+      if not capture then
+        break
       end
-      local outer_range_start_row = outer_range[1]
 
       for _, range in ipairs(tree_region) do
         local intersection = Range.intersection(range, outer_range)
         if intersection then
           local start_row, start_col, end_row, end_col = Range.unpack4(intersection)
 
-          if capture then
-            local hl = state.highlighter_query:get_hl_from_capture(capture)
-
-            local capture_name = 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 >= line and not on_conceal and (not on_spell or spell ~= nil) then
-              local opts = {
-                end_line = end_row,
-                end_col = end_col,
-                hl_group = hl,
-                ephemeral = true,
-                priority = priority,
-                conceal = conceal,
-                spell = spell,
-                url = url,
-              }
-              local mark = { start_line = start_row, start_col = start_col, opts = opts }
-              add_mark(mark, buf, line, next_marks)
-            end
-
-            if
-              (metadata.conceal_lines or metadata[capture] and metadata[capture].conceal_lines)
-              and #api.nvim_buf_get_extmarks(buf, ns, { start_row, 0 }, { start_row, 0 }, {}) == 0
-            then
-              api.nvim_buf_set_extmark(buf, ns, start_row, 0, {
-                end_line = end_row,
-                conceal_lines = '',
-              })
-            end
+          local hl = state.highlighter_query:get_hl_from_capture(capture)
+
+          local capture_name = 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 not on_conceal and (not on_spell or spell ~= nil) then
+            local opts = {
+              end_line = end_row,
+              end_col = end_col,
+              hl_group = hl,
+              ephemeral = true,
+              priority = priority,
+              conceal = conceal,
+              spell = spell,
+              url = url,
+            }
+            local mark = { start_line = start_row, start_col = start_col, opts = opts }
+            add_mark(
+              mark,
+              buf,
+              range_start_row,
+              range_start_col,
+              range_end_row,
+              range_end_col,
+              next_marks
+            )
           end
-        end
-      end
 
-      if outer_range_start_row > line then
-        state.next_row = outer_range_start_row
+          if
+            (metadata.conceal_lines or metadata[capture] and metadata[capture].conceal_lines)
+            and #api.nvim_buf_get_extmarks(buf, ns, { start_row, 0 }, { start_row, 0 }, {}) == 0
+          then
+            api.nvim_buf_set_extmark(buf, ns, start_row, 0, {
+              end_line = end_row,
+              conceal_lines = '',
+            })
+          end
+        end
       end
     end
 
+    state.next_row = next_row
+    state.next_col = next_col
     state.prev_marks = next_marks
   end)
 end
@@ -464,14 +532,17 @@ end
 ---@private
 ---@param win integer
 ---@param buf integer
----@param line integer
-function TSHighlighter._on_line(_, win, buf, line, _)
+---@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_line_impl(self, win, buf, line, false, false)
+  on_range_impl(self, win, buf, br, bc, er, ec, false, false)
 end
 
 ---@private
@@ -490,9 +561,7 @@ function TSHighlighter._on_spell_nav(_, win, buf, srow, _, erow, _)
   local highlight_states = self._highlight_states[win]
   self:prepare_highlight_states(win, srow, erow)
 
-  for row = srow, erow do
-    on_line_impl(self, win, buf, row, true, false)
-  end
+  on_range_impl(self, win, buf, srow, 0, erow, 0, true, false)
   self._highlight_states[win] = highlight_states
 end
 
@@ -510,7 +579,7 @@ function TSHighlighter._on_conceal_line(_, win, buf, row)
   local highlight_states = self._highlight_states[win]
   self.tree:parse({ row, row })
   self:prepare_highlight_states(win, row, row)
-  on_line_impl(self, win, buf, row, false, true)
+  on_range_impl(self, win, buf, row, 0, row + 1, 0, false, true)
   self._highlight_states[win] = highlight_states
 end
 
@@ -554,7 +623,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,
   _on_conceal_line = TSHighlighter._on_conceal_line,
 })
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()