feat(treesitter): handle quantified fold captures

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.
 

runtime/lua/vim/treesitter/_fold.lua

@@ -149,11 +149,26 @@ 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
+    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(node, bufnr, metadata[id])
+          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
@@ -173,6 +188,7 @@ local function get_folds_levels(bufnr, info, srow, erow, parse_injections)
           end
         end
       end
+    end
   end)
 
   local nestmax = vim.wo.foldnestmax

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()