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;