feat(treesitter): async parsing

**Problem:** Parsing can be slow for large files, and it is a blocking
operation which can be disruptive and annoying.

**Solution:** Provide a function for asynchronous parsing, which accepts
a callback to be run after parsing completes.

Co-authored-by: Lewis Russell <lewis6991@gmail.com>

runtime/doc/treesitter.txt

@@ -1479,6 +1479,21 @@ anyway as they will be very frequent. Rather a plugin that does any kind of
 analysis on a tree should use a timer to throttle too frequent updates.
 
 
+LanguageTree:async_parse({range}, {opts})         *LanguageTree:async_parse()*
+    Like |LanguageTree:parse()|, but asynchronous.
+
+    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).
+      • {opts}   (`AsyncParseOpts?`) Options:
+                 • timeout_ms: (integer?) The maximum time (in milliseconds)
+                   that each parse segment should take.
+                 • callback: (fun(trees: table<integer, TSTree>)?) A callback
+                   to run after parsing completes.
+
 LanguageTree:children()                              *LanguageTree:children()*
     Returns a map of language to child tree.
 

runtime/lua/vim/treesitter.lua

@@ -61,7 +61,7 @@ function M._create_parser(bufnr, lang, opts)
     { on_bytes = bytes_cb, on_detach = detach_cb, on_reload = reload_cb, preview = true }
   )
 
-  self:parse()
+  self:async_parse()
 
   return self
 end

runtime/lua/vim/treesitter/highlighter.lua

@@ -82,7 +82,6 @@ TSHighlighter.__index = TSHighlighter
 ---@param tree vim.treesitter.LanguageTree parser object to use for highlighting
 ---@param opts (table|nil) Configuration of the highlighter:
 ---           - queries table overwrite queries used by the highlighter
----@return vim.treesitter.highlighter Created highlighter object
 function TSHighlighter.new(tree, opts)
   local self = setmetatable({}, TSHighlighter)
 
@@ -147,9 +146,7 @@ function TSHighlighter.new(tree, opts)
     vim.opt_local.spelloptions:append('noplainbuffer')
   end)
 
-  self.tree:parse()
-
-  return self
+  self.tree:async_parse()
 end
 
 --- @nodoc
@@ -391,12 +388,22 @@ function TSHighlighter._on_win(_, _win, buf, topline, botline)
   if not self then
     return false
   end
-  self.tree:parse({ topline, botline + 1 })
-  self:prepare_highlight_states(topline, botline + 1)
-  self.redraw_count = self.redraw_count + 1
+  local range = { topline, botline + 1 }
+  self.tree:async_parse(range, {
+    callback = function()
+      self:_async_parse_callback(range)
+    end,
+  })
   return true
 end
 
+--- @param range [integer, integer]
+function TSHighlighter:_async_parse_callback(range)
+  self:prepare_highlight_states(unpack(range))
+  self.redraw_count = self.redraw_count + 1
+  api.nvim__redraw({ buf = self.bufnr, range = range, flush = false })
+end
+
 api.nvim_set_decoration_provider(ns, {
   on_win = TSHighlighter._on_win,
   on_line = TSHighlighter._on_line,

runtime/lua/vim/treesitter/languagetree.lua

@@ -58,6 +58,11 @@ local Range = require('vim.treesitter._range')
 ---| 'on_child_added'
 ---| 'on_child_removed'
 
+---@nodoc
+---@class AsyncParseOpts
+---@field timeout_ms integer?
+---@field callback fun(trees: table<integer, TSTree>)?
+
 --- @type table<TSCallbackNameOn,TSCallbackName>
 local TSCallbackNames = {
   on_changedtree = 'changedtree',
@@ -98,6 +103,12 @@ local LanguageTree = {}
 
 LanguageTree.__index = LanguageTree
 
+--- @return integer
+function LanguageTree:_changedtick()
+  --- @type integer
+  return vim.b[self._source].changedtick
+end
+
 --- @nodoc
 ---
 --- LanguageTree contains a tree of parsers: the root treesitter parser for {lang} and any
@@ -334,13 +345,17 @@ end
 
 --- @private
 --- @param range boolean|Range?
+--- @param opts AsyncParseOpts?
 --- @return Range6[] changes
 --- @return integer no_regions_parsed
 --- @return number total_parse_time
-function LanguageTree:_parse_regions(range)
+--- @return boolean is_unfinished whether async parsing still needs time
+function LanguageTree:_parse_regions(range, opts)
   local changes = {}
+  opts = opts or {}
   local no_regions_parsed = 0
   local total_parse_time = 0
+  local timeout = opts.timeout_ms
 
   if type(self._valid) ~= 'table' then
     self._valid = {}
@@ -357,9 +372,18 @@ function LanguageTree:_parse_regions(range)
       )
     then
       self._parser:set_included_ranges(ranges)
+      if timeout then
+        self._parser:set_timeout(timeout * 1000) -- ms -> micros
+      else
+        self._parser:set_timeout(0)
+      end
       local parse_time, tree, tree_changes =
         tcall(self._parser.parse, self._parser, self._trees[i], self._source, true)
 
+      if not tree then
+        return changes, no_regions_parsed, total_parse_time, true
+      end
+
       -- Pass ranges if this is an initial parse
       local cb_changes = self._trees[i] and tree_changes or tree:included_ranges(true)
 
@@ -373,7 +397,7 @@ function LanguageTree:_parse_regions(range)
     end
   end
 
-  return changes, no_regions_parsed, total_parse_time
+  return changes, no_regions_parsed, total_parse_time, false
 end
 
 --- @private
@@ -409,6 +433,54 @@ function LanguageTree:_add_injections()
   return query_time
 end
 
+--- Like |LanguageTree:parse()|, but asynchronous.
+---
+--- @param range boolean|Range|nil: 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).
+--- @param opts AsyncParseOpts? Options:
+---   - timeout_ms: (integer?) The maximum time (in milliseconds) that each parse segment should
+---     take.
+---   - callback: (fun(trees: table<integer, TSTree>)?) A callback to run after parsing completes.
+function LanguageTree:async_parse(range, opts)
+  coroutine.resume(
+    --- @param co_range boolean|Range|nil
+    --- @param co_opts AsyncParseOpts?
+    coroutine.create(function(co_range, co_opts)
+      local co = coroutine.running()
+      co_opts = co_opts or {}
+      co_opts.timeout_ms = co_opts.timeout_ms or 3
+      local ct = self:_changedtick()
+      local unfinished = true
+      --- @type table<integer, TSTree>
+      local trees
+
+      while unfinished do
+        -- If buffer was changed in the middle of parsing, reset parse state
+        if self:_changedtick() ~= ct then
+          self._parser:reset()
+          self:invalidate()
+          coroutine.yield()
+        end
+
+        trees, unfinished = self:_parse(co_range, co_opts)
+
+        if unfinished then
+          vim.schedule(function()
+            coroutine.resume(co)
+          end)
+          coroutine.yield()
+        elseif co_opts.callback then
+          co_opts.callback(trees)
+        end
+      end
+    end),
+    range,
+    opts
+  )
+end
+
 --- Recursively parse all regions in the language tree using |treesitter-parsers|
 --- for the corresponding languages and run injection queries on the parsed trees
 --- to determine whether child trees should be created and parsed.
@@ -422,9 +494,21 @@ end
 ---     only the root tree without injections).
 --- @return table<integer, TSTree>
 function LanguageTree:parse(range)
+  local trees, _ = self:_parse(range)
+  return trees
+end
+
+--- @param range boolean|Range|nil: 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).
+--- @param opts AsyncParseOpts? See |TSParseOpts|
+--- @return table<integer, TSTree>
+--- @return boolean
+function LanguageTree:_parse(range, opts)
   if self:is_valid() then
     self:_log('valid')
-    return self._trees
+    return self._trees, false
   end
 
   local changes --- @type Range6[]?
@@ -433,10 +517,11 @@ function LanguageTree:parse(range)
   local no_regions_parsed = 0
   local query_time = 0
   local total_parse_time = 0
+  local is_unfinished = false
 
   -- At least 1 region is invalid
   if not self:is_valid(true) then
-    changes, no_regions_parsed, total_parse_time = self:_parse_regions(range)
+    changes, no_regions_parsed, total_parse_time, is_unfinished = self:_parse_regions(range, opts)
     -- Need to run injections when we parsed something
     if no_regions_parsed > 0 then
       self._injections_processed = false
@@ -457,10 +542,10 @@ function LanguageTree:parse(range)
   })
 
   for _, child in pairs(self._children) do
-    child:parse(range)
+    child:_parse(range, opts)
   end
 
-  return self._trees
+  return self._trees, is_unfinished
 end
 
 --- Invokes the callback for each |LanguageTree| recursively.

src/nvim/lua/treesitter.c

@@ -488,7 +488,11 @@ static int parser_parse(lua_State *L)
   // Sometimes parsing fails (timeout, or wrong parser ABI)
   // In those case, just return an error.
   if (!new_tree) {
-    return luaL_error(L, "An error occurred when parsing.");
+    if (ts_parser_timeout_micros(p) == 0) {
+      // No timeout set, must have had an error
+      return luaL_error(L, "An error occurred when parsing.");
+    }
+    return 0;
   }
 
   // The new tree will be pushed to the stack, without copy, ownership is now to the lua GC.