perf(treesitter): calculate folds asynchronously

**Problem:** The treesitter `foldexpr` runs synchronous parses to
calculate fold levels, which eliminates async parsing performance in the
highlighter.

**Solution:** Migrate the `foldexpr` to also calculate and apply fold
levels asynchronously.

runtime/doc/news.txt

@@ -305,6 +305,7 @@ PERFORMANCE
   queries.
 • Treesitter highlighting is now asynchronous. To force synchronous parsing,
   use `vim.g._ts_force_sync_parsing = true`.
+• Treesitter folding is now calculated asynchronously.
 
 PLUGINS
 

runtime/lua/vim/treesitter/_fold.lua

@@ -69,7 +69,8 @@ end
 ---@param info TS.FoldInfo
 ---@param srow integer?
 ---@param erow integer? 0-indexed, exclusive
-local function compute_folds_levels(bufnr, info, srow, erow)
+---@param callback function?
+local function compute_folds_levels(bufnr, info, srow, erow, callback)
   srow = srow or 0
   erow = erow or api.nvim_buf_line_count(bufnr)
 
@@ -78,7 +79,10 @@ local function compute_folds_levels(bufnr, info, srow, erow)
     return
   end
 
-  parser:parse()
+  parser:parse(nil, function(_, trees)
+    if not trees then
+      return
+    end
 
     local enter_counts = {} ---@type table<integer, integer>
     local leave_counts = {} ---@type table<integer, integer>
@@ -176,6 +180,11 @@ local function compute_folds_levels(bufnr, info, srow, erow)
       leave_prev = leave_line
       level0_prev = adjusted
     end
+
+    if callback then
+      callback()
+    end
+  end)
 end
 
 local M = {}
@@ -266,6 +275,8 @@ local function on_changedtree(bufnr, foldinfo, tree_changes)
   schedule_if_loaded(bufnr, function()
     local srow_upd, erow_upd ---@type integer?, integer?
     local max_erow = api.nvim_buf_line_count(bufnr)
+    -- TODO(ribru17): Replace this with a proper .all() awaiter once #19624 is resolved
+    local iterations = 0
     for _, change in ipairs(tree_changes) do
       local srow, _, erow, ecol = Range.unpack4(change)
       -- If a parser doesn't have any ranges explicitly set, treesitter will
@@ -279,15 +290,17 @@ local function on_changedtree(bufnr, foldinfo, tree_changes)
       end
       -- Start from `srow - foldminlines`, because this edit may have shrunken the fold below limit.
       srow = math.max(srow - vim.wo.foldminlines, 0)
-      compute_folds_levels(bufnr, foldinfo, srow, erow)
       srow_upd = srow_upd and math.min(srow_upd, srow) or srow
       erow_upd = erow_upd and math.max(erow_upd, erow) or erow
-    end
-    if #tree_changes > 0 then
+      compute_folds_levels(bufnr, foldinfo, srow, erow, function()
+        iterations = iterations + 1
+        if iterations == #tree_changes then
           foldinfo:foldupdate(bufnr, srow_upd, erow_upd)
         end
       end)
     end
+  end)
+end
 
 ---@param bufnr integer
 ---@param foldinfo TS.FoldInfo
@@ -342,9 +355,10 @@ local function on_bytes(bufnr, foldinfo, start_row, start_col, old_row, old_col,
       foldinfo.on_bytes_range = nil
       -- Start from `srow - foldminlines`, because this edit may have shrunken the fold below limit.
       srow = math.max(srow - vim.wo.foldminlines, 0)
-      compute_folds_levels(bufnr, foldinfo, srow, erow)
+      compute_folds_levels(bufnr, foldinfo, srow, erow, function()
         foldinfo:foldupdate(bufnr, srow, erow)
       end)
+    end)
   end
 end
 
@@ -400,9 +414,10 @@ api.nvim_create_autocmd('OptionSet', {
     for _, bufnr in ipairs(bufs) do
       foldinfos[bufnr] = FoldInfo.new(bufnr)
       api.nvim_buf_call(bufnr, function()
-        compute_folds_levels(bufnr, foldinfos[bufnr])
-      end)
+        compute_folds_levels(bufnr, foldinfos[bufnr], nil, nil, function()
           foldinfos[bufnr]:foldupdate(bufnr, 0, api.nvim_buf_line_count(bufnr))
+        end)
+      end)
     end
   end,
 })