From f5d861a01930c48d289c7e4542002b37723c7d19 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 28 Oct 2015 12:10:45 -0700
Subject: [PATCH] Fix bug where ts_stack_pop results were backwards for some
 stack configurations

---
 spec/runtime/stack_spec.cc | 17 +++++++++++++++--
 src/runtime/lexer.c        |  3 ++-
 src/runtime/stack.c        |  9 +++++++--
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/spec/runtime/stack_spec.cc b/spec/runtime/stack_spec.cc
index 9a8092cc..721881d5 100644
--- a/spec/runtime/stack_spec.cc
+++ b/spec/runtime/stack_spec.cc
@@ -344,17 +344,30 @@ describe("Stack", [&]() {
 
     describe("when there is one path that leads to two different heads", [&]() {
       it("returns two entries with the same array of trees", [&]() {
+        /*
+         *  A0__B1__C2__D3__G6__H7.
+         *       \__E4__F5__/
+         */
+        ts_stack_push(stack, 0, stateH, trees[7]);
+
+
         /*
          *  A0__B1__C2__D3.
          *       \__E4__F5.
          */
-        StackPopResultList pop = ts_stack_pop(stack, 0, 1, false);
+        StackPopResultList pop = ts_stack_pop(stack, 0, 2, false);
         AssertThat(ts_stack_head_count(stack), Equals(2));
 
         AssertThat(pop.size, Equals(2));
         AssertThat(pop.contents[0].index, Equals(0));
+        AssertThat(pop.contents[0].tree_count, Equals(2));
+        AssertThat(pop.contents[0].trees[0], Equals(trees[6]));
+        AssertThat(pop.contents[0].trees[1], Equals(trees[7]));
+
         AssertThat(pop.contents[1].index, Equals(1));
-        AssertThat(pop.contents[0].trees, Equals(pop.contents[1].trees));
+        AssertThat(pop.contents[1].tree_count, Equals(2));
+        AssertThat(pop.contents[1].trees[0], Equals(trees[6]));
+        AssertThat(pop.contents[1].trees[1], Equals(trees[7]));
       });
     });
 
diff --git a/src/runtime/lexer.c b/src/runtime/lexer.c
index 6adb89c6..b2b21819 100644
--- a/src/runtime/lexer.c
+++ b/src/runtime/lexer.c
@@ -22,7 +22,8 @@ static const char *empty_chunk = "";
 
 static void ts_lexer__get_chunk(TSLexer *self) {
   TSInput input = self->input;
-  if (!self->chunk || self->current_position.bytes != self->chunk_start + self->chunk_size)
+  if (!self->chunk ||
+      self->current_position.bytes != self->chunk_start + self->chunk_size)
     input.seek_fn(input.payload, self->current_position);
 
   self->chunk_start = self->current_position.bytes;
diff --git a/src/runtime/stack.c b/src/runtime/stack.c
index 3f84f782..bda889b1 100644
--- a/src/runtime/stack.c
+++ b/src/runtime/stack.c
@@ -105,6 +105,7 @@ static bool stack_node_release(StackNode *self) {
 
 static StackNode *stack_node_new(StackNode *next, TSStateId state, TSTree *tree) {
   StackNode *self = malloc(sizeof(StackNode));
+  assert(tree->ref_count > 0);
   ts_tree_retain(tree);
   stack_node_retain(next);
   *self = (StackNode){
@@ -126,10 +127,12 @@ static void ts_stack__add_node_successor(Stack *self, StackNode *node,
     if (successor == new_successor)
       return;
     if (successor->entry.state == new_successor->entry.state) {
-      if (successor->entry.tree != new_successor->entry.tree)
+      if (successor->entry.tree != new_successor->entry.tree) {
         successor->entry.tree = self->tree_selection_callback.callback(
           self->tree_selection_callback.data, successor->entry.tree,
           new_successor->entry.tree);
+        ts_tree_retain(successor->entry.tree);
+      }
       for (int j = 0; j < new_successor->successor_count; j++)
         ts_stack__add_node_successor(self, successor,
                                      new_successor->successors[j]);
@@ -181,6 +184,7 @@ static bool ts_stack__merge_head(Stack *self, int head_index, TSStateId state,
       if (head->entry.tree != tree) {
         head->entry.tree = self->tree_selection_callback.callback(
           self->tree_selection_callback.data, head->entry.tree, tree);
+        ts_tree_retain(head->entry.tree);
       }
       ts_stack__add_node_successor(self, head, self->heads[head_index]);
       ts_stack_remove_head(self, head_index);
@@ -275,7 +279,8 @@ StackPopResultList ts_stack_pop(Stack *self, int head_index, int child_count,
   }
 
   for (int path = 0; path < path_count; path++) {
-    tree_vector_reverse(&trees_by_path[path]);
+    if (!is_shared_by_path[path])
+      tree_vector_reverse(&trees_by_path[path]);
     int index = -1;
     if (path == 0) {
       stack_node_retain(nodes_by_path[path]);