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 <tomtomjhj@gmail.com>

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
 

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<integer, TSTree>)?`)
                     Function invoked when parsing completes. When provided and
                     `vim.g._ts_force_sync_parsing` is not set, parsing will

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<integer, vim.treesitter.highlighter.State[]>
+---@field private _highlight_states vim.treesitter.highlighter.State[]
 ---@field private _queries table<string,vim.treesitter.highlighter.Query>
 ---@field  _conceal_line boolean?
 ---@field  _conceal_checked table<integer, boolean>
 ---@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<integer, boolean>
+---@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<integer, Range[]>
+  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,

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<TSCallbackName,function[]> Callback handlers (recursive)
 ---@field private _children table<string,vim.treesitter.LanguageTree> 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<integer, TSTree> 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<string, Range6[][]>
 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