fix(treesitter): improve vim.treesitter.foldexpr

* Collect on_bytes and flush at the invocation of the scheduled callback
  to take account of commands that triggers multiple on_bytes.
* More accurately track movement of folds so that foldexpr returns
  reasonable values even when the scheduled computation is not run yet.
* Start computing folds from the line above (+ foldminlines) the changed
  lines to handle the folds that are removed due to the size limit.
* Shrink folds that end at the line at which another fold starts to
  assign proper level to that line.
* Use level '=' for lines that are not computed yet.

1 files changed, 110 insertions, 74 deletions

diff --git a/runtime/lua/vim/treesitter/_fold.lua b/runtime/lua/vim/treesitter/_fold.lua
index d5626d0391..735627d29f 100644
--- a/runtime/lua/vim/treesitter/_fold.lua
+++ b/runtime/lua/vim/treesitter/_fold.lua
@@ -5,35 +5,20 @@ local Range = require('vim.treesitter._range')
 local api = vim.api
 
 ---@class TS.FoldInfo
----@field levels string[] the foldexpr value for each line
+---@field levels string[] the foldexpr result for each line
 ---@field levels0 integer[] the raw fold levels
----@field private start_counts table<integer,integer>
----@field private stop_counts table<integer,integer>
+---@field edits? {[1]: integer, [2]: integer} line range edited since the last invocation of the callback scheduled in on_bytes. 0-indexed, end-exclusive.
 local FoldInfo = {}
 FoldInfo.__index = FoldInfo
 
 ---@private
 function FoldInfo.new()
   return setmetatable({
-    start_counts = {},
-    stop_counts = {},
     levels0 = {},
     levels = {},
   }, FoldInfo)
 end
 
----@package
----@param srow integer
----@param erow integer 0-indexed, exclusive
-function FoldInfo:invalidate_range(srow, erow)
-  for i = srow + 1, erow do
-    self.start_counts[i] = nil
-    self.stop_counts[i] = nil
-    self.levels0[i] = nil
-    self.levels[i] = nil
-  end
-end
-
 --- Efficiently remove items from middle of a list a list.
 ---
 --- Calling table.remove() in a loop will re-index the tail of the table on
@@ -59,8 +44,6 @@ end
 function FoldInfo:remove_range(srow, erow)
   list_remove(self.levels, srow + 1, erow)
   list_remove(self.levels0, srow + 1, erow)
-  list_remove(self.start_counts, srow + 1, erow)
-  list_remove(self.stop_counts, srow + 1, erow)
 end
 
 --- Efficiently insert items into the middle of a list.
@@ -93,44 +76,35 @@ end
 ---@param srow integer
 ---@param erow integer 0-indexed, exclusive
 function FoldInfo:add_range(srow, erow)
-  list_insert(self.levels, srow + 1, erow, '-1')
+  list_insert(self.levels, srow + 1, erow, '=')
   list_insert(self.levels0, srow + 1, erow, -1)
-  list_insert(self.start_counts, srow + 1, erow, nil)
-  list_insert(self.stop_counts, srow + 1, erow, nil)
-end
-
----@package
----@param lnum integer
-function FoldInfo:add_start(lnum)
-  self.start_counts[lnum] = (self.start_counts[lnum] or 0) + 1
 end
 
 ---@package
----@param lnum integer
-function FoldInfo:add_stop(lnum)
-  self.stop_counts[lnum] = (self.stop_counts[lnum] or 0) + 1
-end
-
----@package
----@param lnum integer
----@return integer
-function FoldInfo:get_start(lnum)
-  return self.start_counts[lnum] or 0
+---@param srow integer
+---@param erow_old integer
+---@param erow_new integer 0-indexed, exclusive
+function FoldInfo:edit_range(srow, erow_old, erow_new)
+  if self.edits then
+    self.edits[1] = math.min(srow, self.edits[1])
+    if erow_old <= self.edits[2] then
+      self.edits[2] = self.edits[2] + (erow_new - erow_old)
+    end
+    self.edits[2] = math.max(self.edits[2], erow_new)
+  else
+    self.edits = { srow, erow_new }
+  end
 end
 
 ---@package
----@param lnum integer
----@return integer
-function FoldInfo:get_stop(lnum)
-  return self.stop_counts[lnum] or 0
-end
-
-local function trim_level(level)
-  local max_fold_level = vim.wo.foldnestmax
-  if level > max_fold_level then
-    return max_fold_level
+---@return integer? srow
+---@return integer? erow 0-indexed, exclusive
+function FoldInfo:flush_edit()
+  if self.edits then
+    local srow, erow = self.edits[1], self.edits[2]
+    self.edits = nil
+    return srow, erow
   end
-  return level
 end
 
 --- If a parser doesn't have any ranges explicitly set, treesitter will
@@ -158,22 +132,24 @@ local function get_folds_levels(bufnr, info, srow, erow, parse_injections)
   srow = srow or 0
   erow = normalise_erow(bufnr, erow)
 
-  info:invalidate_range(srow, erow)
-
-  local prev_start = -1
-  local prev_stop = -1
-
   local parser = ts.get_parser(bufnr)
 
   parser:parse(parse_injections and { srow, erow } or nil)
 
+  local enter_counts = {} ---@type table<integer, integer>
+  local leave_counts = {} ---@type table<integer, integer>
+  local prev_start = -1
+  local prev_stop = -1
+
   parser:for_each_tree(function(tree, ltree)
     local query = ts.query.get(ltree:lang(), 'folds')
     if not query then
       return
     end
 
-    for id, node, metadata in query:iter_captures(tree:root(), bufnr, srow, erow) do
+    -- Collect folds starting from srow - 1, because we should first subtract the folds that end at
+    -- srow - 1 from the level of srow - 1 to get accurate level of srow.
+    for id, node, metadata in query:iter_captures(tree:root(), bufnr, math.max(srow - 1, 0), erow) do
       if query.captures[id] == 'fold' then
         local range = ts.get_range(node, bufnr, metadata[id])
         local start, _, stop, stop_col = Range.unpack4(range)
@@ -190,8 +166,8 @@ local function get_folds_levels(bufnr, info, srow, erow, parse_injections)
         if
           fold_length > vim.wo.foldminlines and not (start == prev_start and stop == prev_stop)
         then
-          info:add_start(start + 1)
-          info:add_stop(stop + 1)
+          enter_counts[start + 1] = (enter_counts[start + 1] or 0) + 1
+          leave_counts[stop + 1] = (leave_counts[stop + 1] or 0) + 1
           prev_start = start
           prev_stop = stop
         end
@@ -199,16 +175,15 @@ local function get_folds_levels(bufnr, info, srow, erow, parse_injections)
     end
   end)
 
-  local current_level = info.levels0[srow] or 0
+  local nestmax = vim.wo.foldnestmax
+  local level0_prev = info.levels0[srow] or 0
+  local leave_prev = leave_counts[srow] or 0
 
   -- We now have the list of fold opening and closing, fill the gaps and mark where fold start
   for lnum = srow + 1, erow do
-    local last_trimmed_level = trim_level(current_level)
-    current_level = current_level + info:get_start(lnum)
-    info.levels0[lnum] = current_level
-
-    local trimmed_level = trim_level(current_level)
-    current_level = current_level - info:get_stop(lnum)
+    local enter_line = enter_counts[lnum] or 0
+    local leave_line = leave_counts[lnum] or 0
+    local level0 = level0_prev - leave_prev + enter_line
 
     -- Determine if it's the start/end of a fold
     -- NB: vim's fold-expr interface does not have a mechanism to indicate that
@@ -216,14 +191,36 @@ local function get_folds_levels(bufnr, info, srow, erow, parse_injections)
     --  ( \n ( \n )) \n (( \n ) \n )
     -- versus
     --  ( \n ( \n ) \n ( \n ) \n )
-    -- If it did have such a mechanism, (trimmed_level - last_trimmed_level)
+    -- Both are represented by ['>1', '>2', '2', '>2', '2', '1'], and
+    -- vim interprets as the second case.
+    -- If it did have such a mechanism, (clamped - clamped_prev)
     -- would be the correct number of starts to pass on.
+    local adjusted = level0 ---@type integer
     local prefix = ''
-    if trimmed_level - last_trimmed_level > 0 then
+    if enter_line > 0 then
       prefix = '>'
+      if leave_line > 0 then
+        -- If this line ends a fold f1 and starts a fold f2, then move f1's end to the previous line
+        -- so that f2 gets the correct level on this line. This may reduce the size of f1 below
+        -- foldminlines, but we don't handle it for simplicity.
+        adjusted = level0 - leave_line
+        leave_line = 0
+      end
+    end
+
+    -- Clamp at foldnestmax.
+    local clamped = adjusted
+    if adjusted > nestmax then
+      prefix = ''
+      clamped = nestmax
     end
 
-    info.levels[lnum] = prefix .. tostring(trimmed_level)
+    -- Record the "real" level, so that it can be used as "base" of later get_folds_levels().
+    info.levels0[lnum] = adjusted
+    info.levels[lnum] = prefix .. tostring(clamped)
+
+    leave_prev = leave_line
+    level0_prev = adjusted
   end
 end
 
@@ -297,7 +294,8 @@ local function on_changedtree(bufnr, foldinfo, tree_changes)
       if ecol > 0 then
         erow = erow + 1
       end
-      get_folds_levels(bufnr, foldinfo, srow, erow)
+      -- Start from `srow - foldminlines`, because this edit may have shrunken the fold below limit.
+      get_folds_levels(bufnr, foldinfo, math.max(srow - vim.wo.foldminlines, 0), erow)
     end
     if #tree_changes > 0 then
       foldupdate(bufnr)
@@ -309,20 +307,46 @@ end
 ---@param foldinfo TS.FoldInfo
 ---@param start_row integer
 ---@param old_row integer
+---@param old_col integer
 ---@param new_row integer
-local function on_bytes(bufnr, foldinfo, start_row, old_row, new_row)
+---@param new_col integer
+local function on_bytes(bufnr, foldinfo, start_row, start_col, old_row, old_col, new_row, new_col)
   -- extend the end to fully include the range
   local end_row_old = start_row + old_row + 1
   local end_row_new = start_row + new_row + 1
 
   if new_row ~= old_row then
+    -- foldexpr can be evaluated before the scheduled callback is invoked. So it may observe the
+    -- outdated levels, which may spuriously open the folds that didn't change. So we should shift
+    -- folds as accurately as possible. For this to be perfectly accurate, we should track the
+    -- actual TSNodes that account for each fold, and compare the node's range with the edited
+    -- range. But for simplicity, we just check whether the start row is completely removed (e.g.,
+    -- `dd`) or shifted (e.g., `o`).
     if new_row < old_row then
-      foldinfo:remove_range(end_row_new, end_row_old)
+      if start_col == 0 and new_row == 0 and new_col == 0 then
+        foldinfo:remove_range(start_row, start_row + (end_row_old - end_row_new))
+      else
+        foldinfo:remove_range(end_row_new, end_row_old)
+      end
     else
-      foldinfo:add_range(start_row, end_row_new)
+      if start_col == 0 and old_row == 0 and old_col == 0 then
+        foldinfo:add_range(start_row, start_row + (end_row_new - end_row_old))
+      else
+        foldinfo:add_range(end_row_old, end_row_new)
+      end
     end
+    foldinfo:edit_range(start_row, end_row_old, end_row_new)
+
+    -- This callback must not use on_bytes arguments, because they can be outdated when the callback
+    -- is invoked. For example, `J` with non-zero count triggers multiple on_bytes before executing
+    -- the scheduled callback. So we should collect the edits.
     schedule_if_loaded(bufnr, function()
-      get_folds_levels(bufnr, foldinfo, start_row, end_row_new)
+      local srow, erow = foldinfo:flush_edit()
+      if not srow then
+        return
+      end
+      -- Start from `srow - foldminlines`, because this edit may have shrunken the fold below limit.
+      get_folds_levels(bufnr, foldinfo, math.max(srow - vim.wo.foldminlines, 0), erow)
       foldupdate(bufnr)
     end)
   end
@@ -349,8 +373,8 @@ function M.foldexpr(lnum)
         on_changedtree(bufnr, foldinfos[bufnr], tree_changes)
       end,
 
-      on_bytes = function(_, _, start_row, _, _, old_row, _, _, new_row, _, _)
-        on_bytes(bufnr, foldinfos[bufnr], start_row, old_row, new_row)
+      on_bytes = function(_, _, start_row, start_col, _, old_row, old_col, _, new_row, new_col, _)
+        on_bytes(bufnr, foldinfos[bufnr], start_row, start_col, old_row, old_col, new_row, new_col)
       end,
 
       on_detach = function()
@@ -362,6 +386,18 @@ function M.foldexpr(lnum)
   return foldinfos[bufnr].levels[lnum] or '0'
 end
 
+api.nvim_create_autocmd('OptionSet', {
+  pattern = { 'foldminlines', 'foldnestmax' },
+  desc = 'Refresh treesitter folds',
+  callback = function()
+    for _, bufnr in ipairs(vim.tbl_keys(foldinfos)) do
+      foldinfos[bufnr] = FoldInfo.new()
+      get_folds_levels(bufnr, foldinfos[bufnr])
+      foldupdate(bufnr)
+    end
+  end,
+})
+
 ---@package
 ---@return { [1]: string, [2]: string[] }[]|string
 function M.foldtext()