From e7d3d40a59ee63c67b99f2b8522f1cb9614e8736 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Thu, 10 Mar 2016 11:51:38 -0800
Subject: [PATCH] Explicitly inform stack pop callback when the stack is
 exhausted

Also, pass non-extra tree count as a single value, rather than keeping
track of the extra count and the total separately.
---
 spec/runtime/stack_spec.cc |  4 ++--
 src/runtime/parser.c       |  6 +++---
 src/runtime/stack.c        | 35 ++++++++++++++++++-----------------
 src/runtime/stack.h        |  4 ++--
 4 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/spec/runtime/stack_spec.cc b/spec/runtime/stack_spec.cc
index 8d98b8e9..1e768bcb 100644
--- a/spec/runtime/stack_spec.cc
+++ b/spec/runtime/stack_spec.cc
@@ -71,9 +71,9 @@ vector<StackEntry> get_stack_entries(Stack *stack, int head_index) {
   ts_stack_pop_until(
     stack,
     head_index,
-    [](void *payload, TSStateId state, size_t depth, size_t extra_count) {
+    [](void *payload, TSStateId state, size_t tree_count, bool is_done) {
       auto entries = static_cast<vector<StackEntry> *>(payload);
-      StackEntry entry = {state, depth};
+      StackEntry entry = {state, tree_count};
       if (find(entries->begin(), entries->end(), entry) == entries->end())
         entries->push_back(entry);
       return StackIterateContinue;
diff --git a/src/runtime/parser.c b/src/runtime/parser.c
index 0877b200..9d212aea 100644
--- a/src/runtime/parser.c
+++ b/src/runtime/parser.c
@@ -491,8 +491,8 @@ error:
 
 static StackIterateAction ts_parser__error_repair_callback(void *payload,
                                                            TSStateId state,
-                                                           size_t depth,
-                                                           size_t extra_count) {
+                                                           size_t tree_count,
+                                                           bool is_done) {
   ErrorRepairSession *session = (ErrorRepairSession *)payload;
   const TSParser *self = session->parser;
   size_t count_above_error = session->count_above_error;
@@ -506,7 +506,7 @@ static StackIterateAction ts_parser__error_repair_callback(void *payload,
     TSSymbol symbol = repair->action.data.symbol;
     size_t child_count = repair->action.data.child_count;
 
-    if (depth + count_above_error - extra_count >= child_count &&
+    if (tree_count + count_above_error >= child_count &&
         repair->in_progress_state_count > 0) {
       TSParseAction new_action =
         ts_language_last_action(self->language, state, symbol);
diff --git a/src/runtime/stack.c b/src/runtime/stack.c
index 514a3856..2883d0e8 100644
--- a/src/runtime/stack.c
+++ b/src/runtime/stack.c
@@ -330,7 +330,7 @@ static inline ALWAYS_INLINE StackSliceArray stack__pop(
 
       StackNode *node = path->node;
       size_t successor_count = node->successor_count;
-      switch (callback(payload, node->state, depth, path->extra_count)) {
+      switch (callback(payload, node->state, depth - path->extra_count, node == self->base_node)) {
         case StackIteratePop:
           path->is_done = true;
           continue;
@@ -350,17 +350,17 @@ static inline ALWAYS_INLINE StackSliceArray stack__pop(
 
       all_paths_done = false;
 
-      for (size_t j = 0; j < successor_count; j++) {
-        StackLink successor = node->successors[j];
-
+      for (size_t j = 1; j <= successor_count; j++) {
         PopPath *next_path;
-        if (j == 0) {
-          next_path = path;
+        StackLink successor;
+        if (j == successor_count) {
+          successor = node->successors[0];
+          next_path = &self->pop_paths.contents[i];
         } else {
-          if (!array_push(&self->pop_paths, *path))
+          successor = node->successors[j];
+          if (!array_push(&self->pop_paths, self->pop_paths.contents[i]))
             goto error;
           next_path = array_back(&self->pop_paths);
-          next_path->trees.size--;
           next_path->trees = ts_tree_array_copy(&next_path->trees);
         }
 
@@ -422,22 +422,23 @@ error:
 }
 
 static inline ALWAYS_INLINE StackIterateAction stack__pop_count_callback(
-  void *payload, TSStateId state, size_t depth, size_t extra_count) {
+  void *payload, TSStateId state, size_t tree_count, bool is_done) {
   StackPopSession *pop_session = (StackPopSession *)payload;
   if (pop_session->found_error)
     return StackIterateAbort;
 
-  if (state == ts_parse_state_error && pop_session->goal_tree_count > 0) {
-    pop_session->found_error = true;
+  if (pop_session->goal_tree_count > 0) {
+    if ((int)tree_count == pop_session->goal_tree_count)
+      return StackIteratePop;
+
+    if (state == ts_parse_state_error) {
+      pop_session->found_error = true;
+      return StackIteratePop;
+    }
+  } else if (is_done) {
     return StackIteratePop;
   }
 
-  if ((int)(depth - extra_count) == pop_session->goal_tree_count)
-    return StackIteratePop;
-
-  if (state == 0 && depth > 0)
-    return StackIteratePop;
-
   return StackIterateContinue;
 }
 
diff --git a/src/runtime/stack.h b/src/runtime/stack.h
index c66086b8..e4961d25 100644
--- a/src/runtime/stack.h
+++ b/src/runtime/stack.h
@@ -40,8 +40,8 @@ typedef enum {
 } StackIterateAction;
 
 typedef StackIterateAction (*StackIterateCallback)(void *, TSStateId state,
-                                                   size_t depth,
-                                                   size_t extra_count);
+                                                   size_t tree_count,
+                                                   bool is_done);
 
 typedef int (*TreeSelectionFunction)(void *, TSTree *tree1, TSTree *tree2);