fix(treesitter): correctly handle query quantifiers (#24738)

Query patterns can contain quantifiers (e.g. (foo)+ @bar), so a single
capture can map to multiple nodes. The iter_matches API can not handle
this situation because the match table incorrectly maps capture indices
to a single node instead of to an array of nodes.

The match table should be updated to map capture indices to an array of
nodes. However, this is a massively breaking change, so must be done
with a proper deprecation period.

`iter_matches`, `add_predicate` and `add_directive` must opt-in to the
correct behavior for backward compatibility. This is done with a new
"all" option. This option will become the default and removed after the
0.10 release.

Co-authored-by: Christian Clason <c.clason@uni-graz.at>
Co-authored-by: MDeiml <matthias@deiml.net>
Co-authored-by: Gregory Anders <greg@gpanders.com>

1 files changed, 17 insertions, 5 deletions

diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index c1816a8860..25a753b179 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -1364,9 +1364,16 @@ static int node_equal(lua_State *L)
 /// assumes the match table being on top of the stack
 static void set_match(lua_State *L, TSQueryMatch *match, int nodeidx)
 {
-  for (int i = 0; i < match->capture_count; i++) {
-    push_node(L, match->captures[i].node, nodeidx);
-    lua_rawseti(L, -2, (int)match->captures[i].index + 1);
+  // [match]
+  for (size_t i = 0; i < match->capture_count; i++) {
+    lua_rawgeti(L, -1, (int)match->captures[i].index + 1);  // [match, captures]
+    if (lua_isnil(L, -1)) {  // [match, nil]
+      lua_pop(L, 1);  // [match]
+      lua_createtable(L, 1, 0);  // [match, captures]
+    }
+    push_node(L, match->captures[i].node, nodeidx);  // [match, captures, node]
+    lua_rawseti(L, -2, (int)lua_objlen(L, -2) + 1);  // [match, captures]
+    lua_rawseti(L, -2, (int)match->captures[i].index + 1);  // [match]
   }
 }
 
@@ -1379,7 +1386,7 @@ static int query_next_match(lua_State *L)
   TSQueryMatch match;
   if (ts_query_cursor_next_match(cursor, &match)) {
     lua_pushinteger(L, match.pattern_index + 1);  // [index]
-    lua_createtable(L, (int)ts_query_capture_count(query), 2);  // [index, match]
+    lua_createtable(L, (int)ts_query_capture_count(query), 0);  // [index, match]
     set_match(L, &match, lua_upvalueindex(2));
     return 2;
   }
@@ -1421,7 +1428,8 @@ static int query_next_capture(lua_State *L)
     if (n_pred > 0 && (ud->max_match_id < (int)match.id)) {
       ud->max_match_id = (int)match.id;
 
-      lua_pushvalue(L, lua_upvalueindex(4));  // [index, node, match]
+      // Create a new cleared match table
+      lua_createtable(L, (int)ts_query_capture_count(query), 2);  // [index, node, match]
       set_match(L, &match, lua_upvalueindex(2));
       lua_pushinteger(L, match.pattern_index + 1);
       lua_setfield(L, -2, "pattern");
@@ -1431,6 +1439,10 @@ static int query_next_capture(lua_State *L)
         lua_pushboolean(L, false);
         lua_setfield(L, -2, "active");
       }
+
+      // Set current_match to the new match
+      lua_replace(L, lua_upvalueindex(4));  // [index, node]
+      lua_pushvalue(L, lua_upvalueindex(4));  // [index, node, match]
       return 3;
     }
     return 2;