feat(treesitter): handle quantified fold captures

3 files changed, 61 insertions, 21 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 2c940887c0..7da6fb4ff8 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -267,6 +267,8 @@ The following new APIs and features were added.
   • The `#set!` directive can set the "url" property of a node to have the
     node emit a hyperlink. Hyperlinks are UI specific: in the TUI, the OSC 8
     control sequence is used.
+  • |vim.treesitter.foldexpr()| now recognizes folds captured using a
+    quantified query pattern.
 
 • |:Man| now supports `:hide` modifier to open page in the current window.
 
diff --git a/runtime/lua/vim/treesitter/_fold.lua b/runtime/lua/vim/treesitter/_fold.lua
index d96cc966de..d8b9f4a261 100644
--- a/runtime/lua/vim/treesitter/_fold.lua
+++ b/runtime/lua/vim/treesitter/_fold.lua
@@ -149,27 +149,43 @@ local function get_folds_levels(bufnr, info, srow, erow, parse_injections)
 
     -- 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)
-
-        if stop_col == 0 then
-          stop = stop - 1
-        end
-
-        local fold_length = stop - start + 1
-
-        -- Fold only multiline nodes that are not exactly the same as previously met folds
-        -- Checking against just the previously found fold is sufficient if nodes
-        -- are returned in preorder or postorder when traversing tree
-        if
-          fold_length > vim.wo.foldminlines and not (start == prev_start and stop == prev_stop)
-        then
-          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
+    for _, match, metadata in
+      query:iter_matches(tree:root(), bufnr, math.max(srow - 1, 0), erow, { all = true })
+    do
+      for id, nodes in pairs(match) do
+        if query.captures[id] == 'fold' then
+          local range = ts.get_range(nodes[1], bufnr, metadata[id])
+          local start, _, stop, stop_col = Range.unpack4(range)
+
+          for i = 2, #nodes, 1 do
+            local node_range = ts.get_range(nodes[i], bufnr, metadata[id])
+            local node_start, _, node_stop, node_stop_col = Range.unpack4(node_range)
+            if node_start < start then
+              start = node_start
+            end
+            if node_stop > stop then
+              stop = node_stop
+              stop_col = node_stop_col
+            end
+          end
+
+          if stop_col == 0 then
+            stop = stop - 1
+          end
+
+          local fold_length = stop - start + 1
+
+          -- Fold only multiline nodes that are not exactly the same as previously met folds
+          -- Checking against just the previously found fold is sufficient if nodes
+          -- are returned in preorder or postorder when traversing tree
+          if
+            fold_length > vim.wo.foldminlines and not (start == prev_start and stop == prev_stop)
+          then
+            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
         end
       end
     end
diff --git a/test/functional/treesitter/fold_spec.lua b/test/functional/treesitter/fold_spec.lua
index fc0996dd7f..06f97dc9bb 100644
--- a/test/functional/treesitter/fold_spec.lua
+++ b/test/functional/treesitter/fold_spec.lua
@@ -404,6 +404,28 @@ t3]])
     }, get_fold_levels())
   end)
 
+  it('handles quantified patterns', function()
+    insert([[
+import hello
+import hello
+import hello
+import hello
+import hello
+import hello]])
+
+    exec_lua([[vim.treesitter.query.set('python', 'folds', '(import_statement)+ @fold')]])
+    parse('python')
+
+    eq({
+      [1] = '>1',
+      [2] = '1',
+      [3] = '1',
+      [4] = '1',
+      [5] = '1',
+      [6] = '1',
+    }, get_fold_levels())
+  end)
+
   it('updates folds in all windows', function()
     local screen = Screen.new(60, 48)
     screen:attach()