Folds: fix fold deduplication and improve start/stop logic

1 files changed, 42 insertions, 19 deletions

diff --git a/lua/nvim-treesitter/fold.lua b/lua/nvim-treesitter/fold.lua
index d12f1d06c..1d635bfd2 100644
--- a/lua/nvim-treesitter/fold.lua
+++ b/lua/nvim-treesitter/fold.lua
@@ -9,6 +9,13 @@ local M = {}
 -- Especially not for every line in the file when `zx` is hit
 local folds_levels = tsutils.memoize_by_buf_tick(function(bufnr)
   local max_fold_level = api.nvim_win_get_option(0, "foldnestmax")
+  local trim_level = function(level)
+    if level > max_fold_level then
+      return max_fold_level
+    end
+    return level
+  end
+
   local parser = parsers.get_parser(bufnr)
 
   if not parser then
@@ -23,22 +30,31 @@ local folds_levels = tsutils.memoize_by_buf_tick(function(bufnr)
     end
   end)
 
-  local levels_tmp = {}
+  -- start..stop is an inclusive range
+  local start_counts = {}
+  local stop_counts = {}
+
+  local prev_start = -1
+  local prev_stop = -1
 
   for _, node in ipairs(matches) do
     local start, _, stop, stop_col = node.node:range()
 
-    if stop_col > 0 then
-      stop = stop + 1
+    if stop_col == 0 then
+      stop = stop - 1
     end
 
-    local should_fold = start + 1 < stop -- Only fold for 2+ lines
+    local fold_length = stop - start + 1
+    local should_fold = fold_length >= 2
 
-    -- This can be folded
     -- Fold only multiline nodes that are not exactly the same as previously met folds
-    if should_fold and not (levels_tmp[start] and levels_tmp[stop]) then
-      levels_tmp[start] = (levels_tmp[start] or 0) + 1
-      levels_tmp[stop] = (levels_tmp[stop] or 0) - 1
+    -- Checking against just the previously found fold is sufficient if nodes
+    -- are returned in preorder or postorder when traversing tree
+    if should_fold and not (start == prev_start and stop == prev_stop) then
+      start_counts[start] = (start_counts[start] or 0) + 1
+      stop_counts[stop] = (stop_counts[stop] or 0) + 1
+      prev_start = start
+      prev_stop = stop
     end
   end
 
@@ -48,21 +64,28 @@ local folds_levels = tsutils.memoize_by_buf_tick(function(bufnr)
   -- We now have the list of fold opening and closing, fill the gaps and mark where fold start
   for lnum = 0, api.nvim_buf_line_count(bufnr) do
     local prefix = ""
-    local shift = levels_tmp[lnum] or 0
 
-    -- Determine if it's the start of a fold
-    if levels_tmp[lnum] and shift >= 0 then
+    local last_trimmed_level = trim_level(current_level)
+    current_level = current_level + (start_counts[lnum] or 0)
+    local trimmed_level = trim_level(current_level)
+    current_level = current_level - (stop_counts[lnum] or 0)
+    local next_trimmed_level = trim_level(current_level)
+
+    -- Determine if it's the start/end of a fold
+    -- NB: vim's fold-expr interface does not have a mechanism to indicate that
+    -- two (or more) folds start at this line, so it cannot distinguish between
+    --  ( \n ( \n )) \n (( \n ) \n )
+    -- versus
+    --  ( \n ( \n ) \n ( \n ) \n )
+    -- If it did have such a mechansim, (trimmed_level - last_trimmed_level)
+    -- would be the correct number of starts to pass on.
+    if trimmed_level - last_trimmed_level > 0 then
       prefix = ">"
+    elseif trimmed_level - next_trimmed_level > 0 then
+      prefix = "<"
     end
 
-    current_level = current_level + shift
-
-    -- Ignore folds greater than max_fold_level
-    if current_level > max_fold_level then
-      levels[lnum + 1] = max_fold_level
-    else
-      levels[lnum + 1] = prefix .. tostring(current_level)
-    end
+    levels[lnum + 1] = prefix .. tostring(trimmed_level)
   end
 
   return levels