From 098da1fc2c9ba65adb9c2f7ecb76b26f6e6b0383 Mon Sep 17 00:00:00 2001 From: Riley Bruins Date: Tue, 18 Nov 2025 10:09:49 -0800 Subject: [PATCH] perf(treesitter): parse multiple ranges in languagetree, eliminate flickering #36503 **Problem:** Whenever `LanguageTree:parse()` is called, injection trees from previously parsed ranges are dropped. **Solution:** Allow the function to accept a list of ranges, so it can return injection trees for all the given ranges. Co-authored-by: Jaehwang Jung --- runtime/doc/news.txt | 1 + runtime/doc/treesitter.txt | 15 ++- runtime/lua/vim/treesitter/highlighter.lua | 107 ++++++++------- runtime/lua/vim/treesitter/languagetree.lua | 138 ++++++++++++++------ 4 files changed, 166 insertions(+), 95 deletions(-) diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt index 5b5c1f78d1..221a07bd34 100644 --- a/runtime/doc/news.txt +++ b/runtime/doc/news.txt @@ -340,6 +340,7 @@ TREESITTER • |Query:iter_captures()| supports specifying starting and ending columns. • |:EditQuery| command gained tab-completion, works with injected languages. +• |LanguageTree:parse()| now accepts a list of ranges. TUI diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt index b0d7b82241..9af868e4f9 100644 --- a/runtime/doc/treesitter.txt +++ b/runtime/doc/treesitter.txt @@ -1345,7 +1345,9 @@ LanguageTree:is_valid({exclude_children}, {range}) Parameters: ~ • {exclude_children} (`boolean?`) whether to ignore the validity of children (default `false`) - • {range} (`Range?`) range to check for validity + • {range} (`Range|Range[]?`) range (or list of ranges, + sorted by starting point in ascending order) to + check for validity Return: ~ (`boolean`) @@ -1421,11 +1423,12 @@ LanguageTree:parse({range}, {on_parse}) *LanguageTree:parse()* if {range} is `true`). Parameters: ~ - • {range} (`boolean|Range?`) Parse this range in the parser's - source. Set to `true` to run a complete parse of the - source (Note: Can be slow!) Set to `false|nil` to only - parse regions with empty ranges (typically only the root - tree without injections). + • {range} (`boolean|Range|Range[]?`) Parse this range (or list of + ranges, sorted by starting point in ascending order) in + the parser's source. Set to `true` to run a complete parse + of the source (Note: Can be slow!) Set to `false|nil` to + only parse regions with empty ranges (typically only the + root tree without injections). • {on_parse} (`fun(err?: string, trees?: table)?`) Function invoked when parsing completes. When provided and `vim.g._ts_force_sync_parsing` is not set, parsing will diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua index 9bfbd04cba..472b605e38 100644 --- a/runtime/lua/vim/treesitter/highlighter.lua +++ b/runtime/lua/vim/treesitter/highlighter.lua @@ -67,14 +67,14 @@ end ---@field private orig_spelloptions string --- A map from window ID to highlight states. --- This state is kept during rendering across each line update. ----@field private _highlight_states table +---@field private _highlight_states vim.treesitter.highlighter.State[] ---@field private _queries table ---@field _conceal_line boolean? ---@field _conceal_checked table ---@field tree vim.treesitter.LanguageTree ---@field private redraw_count integer --- A map from window ID to whether we are currently parsing that window asynchronously ----@field parsing table +---@field parsing boolean local TSHighlighter = { active = {}, } @@ -144,7 +144,7 @@ function TSHighlighter.new(tree, opts) self._conceal_checked = {} self._queries = {} self._highlight_states = {} - self.parsing = {} + self.parsing = false -- Queries for a specific language can be overridden by a custom -- string query... if one is not provided it will be looked up by file. @@ -201,12 +201,11 @@ function TSHighlighter:destroy() end end ----@param win integer ---@param srow integer ---@param erow integer exclusive ---@private -function TSHighlighter:prepare_highlight_states(win, srow, erow) - self._highlight_states[win] = {} +function TSHighlighter:prepare_highlight_states(srow, erow) + self._highlight_states = {} self.tree:for_each_tree(function(tstree, tree) if not tstree then @@ -229,7 +228,7 @@ function TSHighlighter:prepare_highlight_states(win, srow, erow) -- _highlight_states should be a list so that the highlights are added in the same order as -- for_each_tree traversal. This ensures that parents' highlight don't override children's. - table.insert(self._highlight_states[win], { + table.insert(self._highlight_states, { tstree = tstree, next_row = 0, next_col = 0, @@ -239,11 +238,10 @@ function TSHighlighter:prepare_highlight_states(win, srow, erow) end) end ----@param win integer ---@param fn fun(state: vim.treesitter.highlighter.State) ---@package -function TSHighlighter:for_each_highlight_state(win, fn) - for _, state in ipairs(self._highlight_states[win] or {}) do +function TSHighlighter:for_each_highlight_state(fn) + for _, state in ipairs(self._highlight_states) do fn(state) end end @@ -327,7 +325,6 @@ local function get_spell(capture_name) end ---@param self vim.treesitter.highlighter ----@param win integer ---@param buf integer ---@param range_start_row integer ---@param range_start_col integer @@ -337,7 +334,6 @@ end ---@param on_conceal boolean local function on_range_impl( self, - win, buf, range_start_row, range_start_col, @@ -362,7 +358,7 @@ local function on_range_impl( local skip_until_col = 0 local subtree_counter = 0 - self:for_each_highlight_state(win, function(state) + self:for_each_highlight_state(function(state) subtree_counter = subtree_counter + 1 local root_node = state.tstree:root() ---@type { [1]: integer, [2]: integer, [3]: integer, [4]: integer } @@ -485,27 +481,25 @@ local function on_range_impl( end ---@private ----@param win integer ---@param buf integer ---@param br integer ---@param bc integer ---@param er integer ---@param ec integer -function TSHighlighter._on_range(_, win, buf, br, bc, er, ec, _) +function TSHighlighter._on_range(_, _, buf, br, bc, er, ec, _) local self = TSHighlighter.active[buf] if not self then return end - return on_range_impl(self, win, buf, br, bc, er, ec, false, false) + return on_range_impl(self, buf, br, bc, er, ec, false, false) end ---@private ----@param win integer ---@param buf integer ---@param srow integer ---@param erow integer -function TSHighlighter._on_spell_nav(_, win, buf, srow, _, erow, _) +function TSHighlighter._on_spell_nav(_, _, buf, srow, _, erow, _) local self = TSHighlighter.active[buf] if not self then return @@ -513,72 +507,85 @@ function TSHighlighter._on_spell_nav(_, win, buf, srow, _, erow, _) -- Do not affect potentially populated highlight state. Here we just want a temporary -- empty state so the C code can detect whether the region should be spell checked. - local highlight_states = self._highlight_states[win] - self:prepare_highlight_states(win, srow, erow) + local highlight_states = self._highlight_states + self:prepare_highlight_states(srow, erow) - on_range_impl(self, win, buf, srow, 0, erow, 0, true, false) - self._highlight_states[win] = highlight_states + on_range_impl(self, buf, srow, 0, erow, 0, true, false) + self._highlight_states = highlight_states end ---@private ----@param win integer ---@param buf integer ---@param row integer -function TSHighlighter._on_conceal_line(_, win, buf, row) +function TSHighlighter._on_conceal_line(_, _, buf, row) local self = TSHighlighter.active[buf] if not self or not self._conceal_line or self._conceal_checked[row] then return end -- Do not affect potentially populated highlight state. - local highlight_states = self._highlight_states[win] + local highlight_states = self._highlight_states self.tree:parse({ row, row }) - self:prepare_highlight_states(win, row, row) - on_range_impl(self, win, buf, row, 0, row + 1, 0, false, true) - self._highlight_states[win] = highlight_states + self:prepare_highlight_states(row, row) + on_range_impl(self, buf, row, 0, row + 1, 0, false, true) + self._highlight_states = highlight_states end ---@private ---@param buf integer ---@param topline integer ---@param botline integer -function TSHighlighter._on_win(_, win, buf, topline, botline) +function TSHighlighter._on_win(_, _, buf, topline, botline) local self = TSHighlighter.active[buf] if not self then return false end - self.parsing[win] = self.parsing[win] - or nil - == self.tree:parse({ topline, botline + 1 }, function(_, trees) - if trees and self.parsing[win] then - self.parsing[win] = false - if api.nvim_win_is_valid(win) then - api.nvim__redraw({ win = win, valid = false, flush = false }) - end - end - end) - if not self.parsing[win] then + if not self.parsing then self.redraw_count = self.redraw_count + 1 - self:prepare_highlight_states(win, topline, botline) + self:prepare_highlight_states(topline, botline) else - self:for_each_highlight_state(win, function(state) - -- TODO(ribru17): Inefficient. Eventually all marks should be applied in on_buf, and all - -- non-folded ranges of each open window should be merged, and iterators should only be - -- created over those regions. This would also fix #31777. - -- - -- Currently this is not possible because the parser discards previously parsed injection - -- trees upon parsing a different region. + self:for_each_highlight_state(function(state) state.iter = nil state.next_row = 0 state.next_col = 0 end) end - local hl_states = self._highlight_states[win] or {} - return #hl_states > 0 + return next(self._highlight_states) ~= nil +end + +function TSHighlighter._on_start() + local buf_ranges = {} ---@type table + for _, win in ipairs(api.nvim_tabpage_list_wins(0)) do + local buf = api.nvim_win_get_buf(win) + if TSHighlighter.active[buf] then + if not buf_ranges[buf] then + buf_ranges[buf] = {} + end + local topline, botline = vim.fn.line('w0', win) - 1, vim.fn.line('w$', win) + table.insert(buf_ranges[buf], { topline, botline }) + end + end + for buf, ranges in pairs(buf_ranges) do + local highlighter = TSHighlighter.active[buf] + if not highlighter.parsing then + table.sort(ranges, function(a, b) + return a[1] < b[1] + end) + highlighter.parsing = highlighter.parsing + or nil + == highlighter.tree:parse(ranges, function(_, trees) + if trees and highlighter.parsing then + highlighter.parsing = false + api.nvim__redraw({ buf = buf, valid = false, flush = false }) + end + end) + end + end end api.nvim_set_decoration_provider(ns, { on_win = TSHighlighter._on_win, + on_start = TSHighlighter._on_start, 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/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua index 9948e96f03..ddd9d337a8 100644 --- a/runtime/lua/vim/treesitter/languagetree.lua +++ b/runtime/lua/vim/treesitter/languagetree.lua @@ -48,10 +48,12 @@ local hrtime = vim.uv.hrtime -- Parse in 3ms chunks. local default_parse_timeout_ns = 3 * 1000000 ----@type Range2 +---@type Range2[] local entire_document_range = { - 0, - math.huge --[[@as integer]], + { + 0, + math.huge --[[@as integer]], + }, } ---@alias TSCallbackName @@ -85,7 +87,7 @@ local TSCallbackNames = { ---@field package _callbacks_rec table Callback handlers (recursive) ---@field private _children table Injected languages ---@field private _injection_query vim.treesitter.Query Queries defining injected languages ----@field private _processed_injection_range Range? Range for which injections have been processed +---@field private _processed_injection_region Range[]? Range for which injections have been processed ---@field private _opts table Options ---@field private _parser TSParser Parser for language ---Table of regions for which the tree is currently running an async parse @@ -161,7 +163,7 @@ function LanguageTree.new(source, lang, opts) _opts = opts, _injection_query = injections[lang] and query.parse(lang, injections[lang]) or query.get(lang, 'injections'), - _processed_injection_range = nil, + _processed_injection_region = nil, _valid_regions = {}, _num_valid_regions = 0, _num_regions = 1, @@ -306,7 +308,7 @@ function LanguageTree:lang() end --- @param region Range6[] ---- @param range? boolean|Range +--- @param range? boolean|Range|Range[] --- @return boolean local function intercepts_region(region, range) if #region == 0 then @@ -321,8 +323,16 @@ local function intercepts_region(region, range) return range end + local is_range_list = type(range[1]) == 'table' for _, r in ipairs(region) do - if Range.intercepts(r, range) then + if is_range_list then + for _, inner_range in ipairs(range) do + ---@cast inner_range Range + if Range.intercepts(r, inner_range) then + return true + end + end + elseif Range.intercepts(r, range) then return true end end @@ -330,10 +340,36 @@ local function intercepts_region(region, range) return false end +--- @param region1 Range6[] +--- @param region2 Range|Range[] +--- @return boolean +local function contains_region(region1, region2) + if type(region2[1]) ~= 'table' then + region2 = { region2 } + end + + -- TODO: Combine intersection ranges in region1 + local i, j, len1, len2 = 1, 1, #region1, #region2 + while i <= len1 and j <= len2 do + local r1 = { Range.unpack4(region1[i]) } + local r2 = { Range.unpack4(region2[j]) } + + if Range.contains(r1, r2) then + j = j + 1 + elseif Range.cmp_pos.lt(r1[3], r1[4], r2[1], r2[2]) then + i = i + 1 + else + return false -- r1 starts after r2 starts and thus can't cover it + end + end + + return j > len2 +end + --- Returns whether this LanguageTree is valid, i.e., |LanguageTree:trees()| reflects the latest --- state of the source. If invalid, user should call |LanguageTree:parse()|. ---@param exclude_children boolean? whether to ignore the validity of children (default `false`) ----@param range Range? range to check for validity +---@param range Range|Range[]? range (or list of ranges, sorted by starting point in ascending order) to check for validity ---@return boolean function LanguageTree:is_valid(exclude_children, range) local valid_regions = self._valid_regions @@ -358,8 +394,8 @@ function LanguageTree:is_valid(exclude_children, range) if not exclude_children then if - not self._processed_injection_range - or not Range.contains(self._processed_injection_range, range or entire_document_range) + not self._processed_injection_region + or not contains_region(self._processed_injection_region, range or entire_document_range) then return false end @@ -387,7 +423,7 @@ function LanguageTree:source() end --- @private ---- @param range boolean|Range? +--- @param range boolean|Range|Range[]? --- @param thread_state ParserThreadState --- @return Range6[] changes --- @return integer no_regions_parsed @@ -483,10 +519,24 @@ function LanguageTree:_add_injections(injections_by_lang) end end ---- @param range boolean|Range? +--- @param range boolean|Range|Range[]? --- @return string local function range_to_string(range) - return type(range) == 'table' and table.concat(range, ',') or tostring(range) + if type(range) ~= 'table' then + return tostring(range) + end + if type(range[1]) ~= 'table' then + return table.concat(range, ',') + end + ---@cast range Range[] + local str = '' + for i, r in ipairs(range) do + if i > 1 then + str = str .. '|' + end + str = str .. table.concat(r, ',') + end + return str end --- @private @@ -577,7 +627,8 @@ end --- Any region with empty range (`{}`, typically only the root tree) is always parsed; --- otherwise (typically injections) only if it intersects {range} (or if {range} is `true`). --- ---- @param range boolean|Range|nil: Parse this range in the parser's source. +--- @param range? boolean|Range|Range[]: Parse this range (or list of ranges, sorted by starting +--- point in ascending order) in the parser's source. --- Set to `true` to run a complete parse of the source (Note: Can be slow!) --- Set to `false|nil` to only parse regions with empty ranges (typically --- only the root tree without injections). @@ -609,7 +660,7 @@ function LanguageTree:_subtract_time(thread_state, time) end --- @private ---- @param range boolean|Range|nil +--- @param range? boolean|Range|Range[] --- @param thread_state ParserThreadState --- @return table trees --- @return boolean finished @@ -632,16 +683,16 @@ function LanguageTree:_parse(range, thread_state) -- Need to run injections when we parsed something if no_regions_parsed > 0 then - self._processed_injection_range = nil + self._processed_injection_region = nil end end if range and not ( - self._processed_injection_range - and Range.contains( - self._processed_injection_range, + self._processed_injection_region + and contains_region( + self._processed_injection_region, range ~= true and range or entire_document_range ) ) @@ -1044,12 +1095,12 @@ end --- TODO: Allow for an offset predicate to tailor the injection range --- instead of using the entire nodes range. --- @private ---- @param range Range|true +--- @param range Range|Range[]|true --- @param thread_state ParserThreadState --- @return table function LanguageTree:_get_injections(range, thread_state) if not self._injection_query or #self._injection_query.captures == 0 then - self._processed_injection_range = entire_document_range + self._processed_injection_region = entire_document_range return {} end @@ -1059,40 +1110,49 @@ function LanguageTree:_get_injections(range, thread_state) local result = {} local full_scan = range == true or self._injection_query.has_combined_injections + if not full_scan and type(range[1]) ~= 'table' then + ---@diagnostic disable-next-line: missing-fields, assign-type-mismatch + range = { range } + end + ---@cast range Range[] for tree_index, tree in pairs(self._trees) do ---@type vim.treesitter.languagetree.Injection local injections = {} local root_node = tree:root() local parent_ranges = self._regions and self._regions[tree_index] or nil - local start_line, end_line ---@type integer, integer + local scan_region ---@type Range4[] if full_scan then - start_line, _, end_line = root_node:range() + --- @diagnostic disable-next-line: missing-fields LuaLS varargs bug + scan_region = { { root_node:range() } } else - start_line, _, end_line = Range.unpack4(range --[[@as Range]]) + scan_region = range end - for pattern, match, metadata in - self._injection_query:iter_matches(root_node, self._source, start_line, end_line + 1) - do - local lang, combined, ranges = self:_get_injection(match, metadata) - if lang then - add_injection(injections, pattern, lang, combined, ranges, parent_ranges, result) - else - self:_log('match from injection query failed for pattern', pattern) - end + for _, r in ipairs(scan_region) do + local start_line, _, end_line, _ = Range.unpack4(r) + for pattern, match, metadata in + self._injection_query:iter_matches(root_node, self._source, start_line, end_line + 1) + do + local lang, combined, ranges = self:_get_injection(match, metadata) + if lang then + add_injection(injections, pattern, lang, combined, ranges, parent_ranges, result) + else + self:_log('match from injection query failed for pattern', pattern) + end - -- Check the current function duration against the timeout, if it exists. - local current_time = hrtime() - self:_subtract_time(thread_state, current_time - start) - start = hrtime() + -- Check the current function duration against the timeout, if it exists. + local current_time = hrtime() + self:_subtract_time(thread_state, current_time - start) + start = hrtime() + end end end if full_scan then - self._processed_injection_range = entire_document_range + self._processed_injection_region = entire_document_range else - self._processed_injection_range = range --[[@as Range]] + self._processed_injection_region = range end return result