From 5edbabdbec0ac3fba33be8afc008845130158583 Mon Sep 17 00:00:00 2001 From: vanaigr Date: Sun, 3 Aug 2025 21:39:55 -0500 Subject: perf: add on_range in treesitter highlighting --- runtime/lua/vim/_meta/api.lua | 12 +- runtime/lua/vim/_meta/api_keysets.lua | 1 + runtime/lua/vim/treesitter/_meta/misc.lua | 6 +- runtime/lua/vim/treesitter/highlighter.lua | 235 +++++++++++++++++++---------- runtime/lua/vim/treesitter/query.lua | 48 ++++-- 5 files changed, 204 insertions(+), 98 deletions(-) (limited to 'runtime/lua/vim') 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 local match_cache = {} - local function iter(end_line) + local function iter(end_line, end_col) local capture, captured_node, match = cursor:next_capture() if not capture then @@ -1005,9 +1015,22 @@ function Query:iter_captures(node, source, start, stop, opts) local predicates = processed_pattern.predicates if not self:_match_predicates(predicates, pattern_i, captures, source) then cursor:remove_match(match_id) - if end_line and captured_node:range() > end_line then + + local row, col = captured_node:range() + + local outside = false + if end_line then + if end_col then + outside = cmp_ge(row, col, end_line, end_col) + else + outside = row > end_line + end + end + + if outside then return nil, captured_node, nil, nil end + return iter(end_line) -- tail call: try next match end @@ -1072,7 +1095,14 @@ function Query:iter_matches(node, source, start, stop, opts) start, stop = value_or_node_range(start, stop, node) local tree = node:tree() - local cursor = vim._create_ts_querycursor(node, self.query, start, stop, opts) + local cursor = vim._create_ts_querycursor(node, self.query, { + start_row = start, + start_col = 0, + end_row = stop, + end_col = 0, + max_start_depth = opts.max_start_depth, + match_limit = opts.match_limit or 256, + }) local function iter() local match = cursor:next_match() -- cgit