Store the lookahead subtree of paused stack versions, not just the lookahead symbol

lib/src/parser.c

@@ -1063,7 +1063,7 @@ static bool ts_parser__do_all_potential_reductions(
 static void ts_parser__handle_error(
   TSParser *self,
   StackVersion version,
-  TSSymbol lookahead_symbol
+  Subtree lookahead
 ) {
   uint32_t previous_version_count = ts_stack_version_count(self->stack);
 
@@ -1093,7 +1093,7 @@ static void ts_parser__handle_error(
         if (ts_language_has_reduce_action(
           self->language,
           state_after_missing_symbol,
-          lookahead_symbol
+          ts_subtree_leaf_symbol(lookahead)
         )) {
           // In case the parser is currently outside of any included range, the lexer will
           // snap to the beginning of the next included range. The missing token's padding
@@ -1114,7 +1114,7 @@ static void ts_parser__handle_error(
 
           if (ts_parser__do_all_potential_reductions(
             self, version_with_missing_tree,
-            lookahead_symbol
+            ts_subtree_leaf_symbol(lookahead)
           )) {
             LOG(
               "recover_with_missing symbol:%s, state:%u",
@@ -1138,6 +1138,7 @@ static void ts_parser__handle_error(
   }
 
   ts_stack_record_summary(self->stack, version, MAX_SUMMARY_DEPTH);
+  ts_subtree_release(&self->tree_pool, lookahead);
   LOG_STACK();
 }
 
@@ -1523,7 +1524,7 @@ static bool ts_parser__advance(
     }
 
     if (!lookahead.ptr) {
-      ts_stack_pause(self->stack, version, ts_builtin_sym_end);
+      ts_stack_pause(self->stack, version, lookahead);
       return true;
     }
 
@@ -1576,8 +1577,7 @@ static bool ts_parser__advance(
     // version advances successfully, then this version can simply be removed.
     // But if all versions end up paused, then error recovery is needed.
     LOG("detect_error");
-    ts_stack_pause(self->stack, version, ts_subtree_leaf_symbol(lookahead));
-    ts_subtree_release(&self->tree_pool, lookahead);
+    ts_stack_pause(self->stack, version, lookahead);
     return true;
   }
 }
@@ -1660,8 +1660,8 @@ static unsigned ts_parser__condense_stack(TSParser *self) {
         if (!has_unpaused_version && self->accept_count < MAX_VERSION_COUNT) {
           LOG("resume version:%u", i);
           min_error_cost = ts_stack_error_cost(self->stack, i);
-          TSSymbol lookahead_symbol = ts_stack_resume(self->stack, i);
-          ts_parser__handle_error(self, i, lookahead_symbol);
+          Subtree lookahead = ts_stack_resume(self->stack, i);
+          ts_parser__handle_error(self, i, lookahead);
           has_unpaused_version = true;
         } else {
           ts_stack_remove_version(self->stack, i);

lib/src/stack.c

@@ -53,10 +53,10 @@ typedef enum {
 
 typedef struct {
   StackNode *node;
-  Subtree last_external_token;
   StackSummary *summary;
   unsigned node_count_at_last_error;
-  TSSymbol lookahead_when_paused;
+  Subtree last_external_token;
+  Subtree lookahead_when_paused;
   StackStatus status;
 } StackHead;
 
@@ -256,6 +256,9 @@ static void stack_head_delete(
     if (self->last_external_token.ptr) {
       ts_subtree_release(subtree_pool, self->last_external_token);
     }
+    if (self->lookahead_when_paused.ptr) {
+      ts_subtree_release(subtree_pool, self->lookahead_when_paused);
+    }
     if (self->summary) {
       array_delete(self->summary);
       ts_free(self->summary);
@@ -274,7 +277,7 @@ static StackVersion ts_stack__add_version(
     .node_count_at_last_error = self->heads.contents[original_version].node_count_at_last_error,
     .last_external_token = self->heads.contents[original_version].last_external_token,
     .status = StackStatusActive,
-    .lookahead_when_paused = 0,
+    .lookahead_when_paused = NULL_SUBTREE,
   };
   array_push(&self->heads, head);
   stack_node_retain(node);
@@ -703,7 +706,7 @@ void ts_stack_halt(Stack *self, StackVersion version) {
   array_get(&self->heads, version)->status = StackStatusHalted;
 }
 
-void ts_stack_pause(Stack *self, StackVersion version, TSSymbol lookahead) {
+void ts_stack_pause(Stack *self, StackVersion version, Subtree lookahead) {
   StackHead *head = array_get(&self->heads, version);
   head->status = StackStatusPaused;
   head->lookahead_when_paused = lookahead;
@@ -722,12 +725,12 @@ bool ts_stack_is_paused(const Stack *self, StackVersion version) {
   return array_get(&self->heads, version)->status == StackStatusPaused;
 }
 
-TSSymbol ts_stack_resume(Stack *self, StackVersion version) {
+Subtree ts_stack_resume(Stack *self, StackVersion version) {
   StackHead *head = array_get(&self->heads, version);
   assert(head->status == StackStatusPaused);
-  TSSymbol result = head->lookahead_when_paused;
+  Subtree result = head->lookahead_when_paused;
   head->status = StackStatusActive;
-  head->lookahead_when_paused = 0;
+  head->lookahead_when_paused = NULL_SUBTREE;
   return result;
 }
 
@@ -739,9 +742,9 @@ void ts_stack_clear(Stack *self) {
   array_clear(&self->heads);
   array_push(&self->heads, ((StackHead) {
     .node = self->base_node,
-    .last_external_token = NULL_SUBTREE,
     .status = StackStatusActive,
-    .lookahead_when_paused = 0,
+    .last_external_token = NULL_SUBTREE,
+    .lookahead_when_paused = NULL_SUBTREE,
   }));
 }
 

lib/src/stack.h

@@ -99,9 +99,9 @@ bool ts_stack_merge(Stack *, StackVersion, StackVersion);
 // Determine whether the given two stack versions can be merged.
 bool ts_stack_can_merge(Stack *, StackVersion, StackVersion);
 
-TSSymbol ts_stack_resume(Stack *, StackVersion);
+Subtree ts_stack_resume(Stack *, StackVersion);
 
-void ts_stack_pause(Stack *, StackVersion, TSSymbol);
+void ts_stack_pause(Stack *, StackVersion, Subtree);
 
 void ts_stack_halt(Stack *, StackVersion);