Encapsulate ParseStackNodes

spec/runtime/helpers/tree_helpers.cc

@@ -1,5 +1,10 @@
 #include "runtime/helpers/tree_helpers.h"
 
+const char *symbol_names[12] = {
+  "ERROR", "END", "DOCUMENT", "AMBIGUITY",
+  "zero", "one", "two", "three", "four", "five", "six", "seven",
+};
+
 TSTree ** tree_array(std::vector<TSTree *> trees) {
   TSTree ** result = (TSTree **)calloc(trees.size(), sizeof(TSTree *));
   for (size_t i = 0; i < trees.size(); i++)
@@ -18,3 +23,7 @@ std::ostream &operator<<(std::ostream &stream, const EqualsTree &matcher) {
   stream << std::string("equals tree: ") << std::string(ts_tree_string(matcher.expected, matcher.symbol_names));
   return stream;
 }
+
+std::ostream &operator<<(std::ostream &stream, const TSTree *tree) {
+  return stream << std::string(ts_tree_string(tree, symbol_names));;
+}

spec/runtime/helpers/tree_helpers.h

@@ -5,6 +5,8 @@
 #include <vector>
 #include <string>
 
+extern const char *symbol_names[12];
+
 TSTree ** tree_array(std::vector<TSTree *> trees);
 
 struct EqualsTree {
@@ -15,5 +17,6 @@ struct EqualsTree {
 };
 
 std::ostream &operator<<(std::ostream &stream, const EqualsTree &matcher);
+std::ostream &operator<<(std::ostream &stream, const TSTree *tree);
 
 #endif  // HELPERS_TREE_HELPERS_H_

spec/runtime/parse_stack_spec.cc

@@ -12,18 +12,17 @@ enum {
   symbol1, symbol2, symbol3, symbol4, symbol5, symbol6, symbol7
 };
 
-const char *symbol_names[] = {
-  "ERROR", "END", "DOCUMENT", "AMBIGUITY",
-  "zero", "one", "two", "three", "four", "five", "six", "seven",
-};
-
 START_TEST
 
 describe("ParseStack", [&]() {
   ParseStack *stack;
-  const size_t tree_count = 6;
+  const size_t tree_count = 8;
   TSTree *trees[tree_count];
 
+  bool merged;
+  ParseStackPopResultList pop;
+  int new_index;
+
   before_each([&]() {
     stack = ts_parse_stack_new();
     TSLength len = ts_length_make(2, 2);
@@ -37,214 +36,198 @@ describe("ParseStack", [&]() {
       ts_tree_release(trees[i]);
   });
 
-  it("starts with a single null head", [&]() {
+  it("starts with a single head with a null tree", [&]() {
     AssertThat(ts_parse_stack_head_count(stack), Equals(1));
-    AssertThat(ts_parse_stack_head(stack, 0), Equals<ParseStackNode *>(NULL));
+    AssertThat(ts_parse_stack_head(stack, 0), Equals<const ParseStackEntry *>(nullptr));
   });
 
-  describe("shift(head_index, state, tree)", [&]() {
+  it("can push and pop (state, tree) pairs to and from the stack", [&]() {
-    it("pushes a node with the given state and tree onto the given head", [&]() {
+    /*
-      ts_parse_stack_shift(stack, 0, 100, trees[0]);
+     *  A0.
-      ts_parse_stack_shift(stack, 0, 101, trees[1]);
+     */
-      ts_parse_stack_shift(stack, 0, 102, trees[2]);
+    ts_parse_stack_push(stack, 0, stateA, trees[0]);
+    const ParseStackEntry *entry1 = ts_parse_stack_head(stack, 0);
+    AssertThat(*entry1, Equals<ParseStackEntry>({trees[0], stateA}));
+    AssertThat(ts_parse_stack_entry_next_count(entry1), Equals(1));
+    AssertThat(ts_parse_stack_entry_next(entry1, 0), Equals<const ParseStackEntry *>(nullptr));
 
-      ParseStackNode *head = ts_parse_stack_head(stack, 0);
+    /*
-      AssertThat(head->state, Equals(102));
+     *  A0__B1.
-      AssertThat(head->tree, Equals(trees[2]));
+     */
-      AssertThat(head->successor_count, Equals(1));
+    ts_parse_stack_push(stack, 0, stateB, trees[1]);
+    const ParseStackEntry *entry2 = ts_parse_stack_head(stack, 0);
+    AssertThat(*entry2, Equals<ParseStackEntry>({trees[1], stateB}));
+    AssertThat(ts_parse_stack_entry_next_count(entry2), Equals(1));
+    AssertThat(ts_parse_stack_entry_next(entry2, 0), Equals(entry1));
 
-      head = head->successors[0];
+    /*
-      AssertThat(head->state, Equals(101));
+     *  A0__B1__C2.
-      AssertThat(head->tree, Equals(trees[1]));
+     */
-      AssertThat(head->successor_count, Equals(1));
+    ts_parse_stack_push(stack, 0, stateC, trees[2]);
+    const ParseStackEntry *entry3 = ts_parse_stack_head(stack, 0);
+    AssertThat(*entry3, Equals<ParseStackEntry>({trees[2], stateC}));
+    AssertThat(ts_parse_stack_entry_next_count(entry3), Equals(1));
+    AssertThat(ts_parse_stack_entry_next(entry3, 0), Equals(entry2));
 
-      head = head->successors[0];
+    /*
-      AssertThat(head->state, Equals(100));
+     *  A0.
-      AssertThat(head->tree, Equals(trees[0]));
+     */
-      AssertThat(head->successor_count, Equals(1));
+    pop = ts_parse_stack_pop(stack, 0, 2);
+    AssertThat(pop.size, Equals(1));
+    AssertThat(pop.contents[0].tree_count, Equals(2));
+    AssertThat(pop.contents[0].trees[0], Equals(trees[1]));
+    AssertThat(pop.contents[0].trees[1], Equals(trees[2]));
+    AssertThat(ts_parse_stack_head(stack, 0), Equals(entry1));
 
-      head = head->successors[0];
+    /*
-      AssertThat(head, Equals<ParseStackNode *>(NULL));
+     *  .
-    });
+     */
+    pop = ts_parse_stack_pop(stack, 0, 1);
+    AssertThat(pop.size, Equals(1));
+    AssertThat(pop.contents[0].tree_count, Equals(1));
+    AssertThat(pop.contents[0].trees[0], Equals(trees[0]));
+    AssertThat(ts_parse_stack_head(stack, 0), Equals<const ParseStackEntry *>(nullptr));
   });
 
-  describe("reduce(head_index, state, symbol, child_count)", [&]() {
+  it("does not count 'extra' trees toward the count given to pop()", [&]() {
+    ts_parse_stack_push(stack, 0, stateA, trees[0]);
+    ts_parse_stack_push(stack, 0, stateB, trees[1]);
+    ts_parse_stack_push(stack, 0, stateC, trees[2]);
+    ts_tree_set_extra(trees[1]);
+
+    pop = ts_parse_stack_pop(stack, 0, 2);
+    AssertThat(pop.size, Equals(1));
+    AssertThat(pop.contents[0].tree_count, Equals(3));
+    AssertThat(pop.contents[0].trees[0], Equals(trees[0]));
+    AssertThat(pop.contents[0].trees[1], Equals(trees[1]));
+    AssertThat(pop.contents[0].trees[2], Equals(trees[2]));
+    AssertThat(ts_parse_stack_head(stack, 0), Equals<const ParseStackEntry *>(nullptr));
+  });
+
+  it("can split a head and manipulate the two resulting heads independently", [&]() {
+    /*
+     *  A0__B1__C2.
+     */
+    ts_parse_stack_push(stack, 0, stateA, trees[0]);
+    ts_parse_stack_push(stack, 0, stateB, trees[1]);
+    ts_parse_stack_push(stack, 0, stateC, trees[2]);
+
+    new_index = ts_parse_stack_split(stack, 0);
+    AssertThat(ts_parse_stack_head_count(stack), Equals(2));
+    AssertThat(new_index, Equals(1));
+
+    /*
+     *  A0__B1__C2__D3.
+     *       \.
+     */
+    ts_parse_stack_push(stack, 0, stateD, trees[3]);
+    ts_parse_stack_pop(stack, 1, 1);
+
+    AssertThat(ts_parse_stack_head_count(stack), Equals(2));
+    AssertThat(*ts_parse_stack_head(stack, 0), Equals<ParseStackEntry>({trees[3], stateD}));
+    AssertThat(*ts_parse_stack_head(stack, 1), Equals<ParseStackEntry>({trees[1], stateB}));
+
+    /*
+     *  A0__B1__C2__D3.
+     *       \__E4__F3.
+     */
+    ts_parse_stack_push(stack, 1, stateE, trees[4]);
+    ts_parse_stack_push(stack, 1, stateF, trees[3]);
+
+    AssertThat(ts_parse_stack_head_count(stack), Equals(2));
+    AssertThat(*ts_parse_stack_head(stack, 0), Equals<ParseStackEntry>({trees[3], stateD}));
+    AssertThat(*ts_parse_stack_head(stack, 1), Equals<ParseStackEntry>({trees[3], stateF}));
+  });
+
+  describe("when same state is pushed onto two heads", [&]() {
     before_each([&]() {
-      ts_parse_stack_shift(stack, 0, 100, trees[0]);
+      /*
-      ts_parse_stack_shift(stack, 0, 101, trees[1]);
+       *  A0__B1__C2.
-      ts_parse_stack_shift(stack, 0, 102, trees[2]);
+       */
-    });
+      ts_parse_stack_push(stack, 0, stateA, trees[0]);
+      ts_parse_stack_push(stack, 0, stateB, trees[1]);
+      ts_parse_stack_push(stack, 0, stateC, trees[2]);
 
-    it("replaces the given number of nodes with a single parent node", [&]() {
+      /*
-      ts_parse_stack_reduce(stack, 0, 103, symbol4, 2);
+       *  A0__B1__C2__D3.
+       *       \__E4__F5.
+       */
+      ts_parse_stack_split(stack, 0);
+      ts_parse_stack_push(stack, 0, stateD, trees[3]);
+      ts_parse_stack_pop(stack, 1, 1);
+      ts_parse_stack_push(stack, 1, stateE, trees[4]);
+      ts_parse_stack_push(stack, 1, stateF, trees[5]);
 
-      ParseStackNode *head = ts_parse_stack_head(stack, 0);
-      AssertThat(head->state, Equals(103));
-      AssertThat(head->tree, Fulfills(EqualsTree(
-        ts_tree_make_node(symbol4, 2, tree_array({ trees[1], trees[2] }), false),
-        symbol_names)));
-      AssertThat(head->successor_count, Equals(1));
-
-      head = head->successors[0];
-      AssertThat(head->state, Equals(100));
-      AssertThat(head->tree, Equals(trees[0]));
-      AssertThat(head->successor_count, Equals(1));
-    });
-
-    describe("when one of the reduced children is an 'extra' tree", [&]() {
-      before_each([&]() {
-        ts_tree_set_extra(trees[1]);
-      });
-
-      it("does not count that child toward the number of children to replace", [&]() {
-        ts_parse_stack_reduce(stack, 0, 103, symbol4, 2);
-
-        ParseStackNode *head = ts_parse_stack_head(stack, 0);
-        AssertThat(head->state, Equals(103));
-        AssertThat(head->tree, Fulfills(EqualsTree(
-          ts_tree_make_node(symbol4, 3, tree_array({ trees[0], trees[1], trees[2] }), false),
-          symbol_names)));
-        AssertThat(head->successor_count, Equals(1));
-
-        head = head->successors[0];
-        AssertThat(head, Equals<ParseStackNode *>(NULL));
-      });
-    });
-  });
-
-  describe("split(head_index)", [&]() {
-    int new_index;
-    bool merged;
-
-    before_each([&]() {
-      //  A0__B1__C2
-      ts_parse_stack_shift(stack, 0, stateA, trees[0]);
-      ts_parse_stack_shift(stack, 0, stateB, trees[1]);
-      ts_parse_stack_shift(stack, 0, stateC, trees[2]);
-
-      new_index = ts_parse_stack_split(stack, 0);
-      AssertThat(new_index, Equals(1));
       AssertThat(ts_parse_stack_head_count(stack), Equals(2));
-
+      AssertThat(*ts_parse_stack_head(stack, 0), Equals<ParseStackEntry>({trees[3], stateD}));
-      //  A0__B1__C2__D3
+      AssertThat(*ts_parse_stack_head(stack, 1), Equals<ParseStackEntry>({trees[5], stateF}));
-      //   \______E4__F3
-      ts_parse_stack_shift(stack, 0, stateD, trees[3]);
-      ts_parse_stack_reduce(stack, 1, stateE, symbol4, 2);
-      merged = ts_parse_stack_shift(stack, 1, stateF, trees[3]);
     });
 
-    it("creates a new head pointing to the same node as the given head", [&]() {
+    describe("when the trees are identical", [&]() {
-      AssertThat(merged, IsFalse());
+      it("merges the heads, and removes nodes along both heads on subsequent pop operations", [&]() {
-      AssertThat(ts_parse_stack_head_count(stack), Equals(2));
+        /*
-
+         *  A0__B1__C2__D3__G6.
-      ParseStackNode *head1 = ts_parse_stack_head(stack, 0);
+         *       \__E4__F5__/
-      AssertThat(head1->state, Equals(stateD));
+         */
-      AssertThat(head1->tree, Equals(trees[3]));
+        merged = ts_parse_stack_push(stack, 0, stateG, trees[6]);
-      AssertThat(head1->successor_count, Equals(1));
-
-      ParseStackNode *head2 = ts_parse_stack_head(stack, 1);
-      AssertThat(head2->state, Equals(stateF));
-      AssertThat(head2->tree, Equals(trees[3]));
-      AssertThat(head2->successor_count, Equals(1));
-    });
-
-    describe("when the same state and tree are shifted onto both heads", [&]() {
-      before_each([&]() {
-        //  A0__B1__C2__D3__G5
-        //   \______E4__F3__/
-        merged = ts_parse_stack_shift(stack, 0, stateG, trees[5]);
         AssertThat(merged, IsFalse());
-        merged = ts_parse_stack_shift(stack, 1, stateG, trees[5]);
+        merged = ts_parse_stack_push(stack, 1, stateG, trees[6]);
         AssertThat(merged, IsTrue());
-      });
 
-      it("re-joins the heads", [&]() {
         AssertThat(ts_parse_stack_head_count(stack), Equals(1));
+        const ParseStackEntry *entry1 = ts_parse_stack_head(stack, 0);
+        AssertThat(*entry1, Equals<ParseStackEntry>({trees[6], stateG}));
+        AssertThat(ts_parse_stack_entry_next_count(entry1), Equals(2));
+        AssertThat(*ts_parse_stack_entry_next(entry1, 0), Equals<ParseStackEntry>({trees[3], stateD}));
+        AssertThat(*ts_parse_stack_entry_next(entry1, 1), Equals<ParseStackEntry>({trees[5], stateF}));
 
-        ParseStackNode *head = ts_parse_stack_head(stack, 0);
+        /*
-        AssertThat(head->state, Equals(stateG));
+         *  A0__B1__C2.
-        AssertThat(head->tree, Equals(trees[5]));
+         *       \__E4.
-        AssertThat(head->successor_count, Equals(2));
+         */
+        pop = ts_parse_stack_pop(stack, 0, 2);
 
-        ParseStackNode *successor1 = head->successors[0];
+        AssertThat(pop.size, Equals(2));
-        AssertThat(successor1->state, Equals(stateD));
+        AssertThat(pop.contents[0].tree_count, Equals(2));
-        AssertThat(successor1->tree, Equals(trees[3]))
+        AssertThat(pop.contents[0].trees[0], Equals(trees[3]));
-        AssertThat(successor1->successor_count, Equals(1));
+        AssertThat(pop.contents[0].trees[1], Equals(trees[6]));
+        AssertThat(pop.contents[1].tree_count, Equals(2));
+        AssertThat(pop.contents[1].trees[0], Equals(trees[5]));
+        AssertThat(pop.contents[1].trees[1], Equals(trees[6]));
 
-        ParseStackNode *successor2 = head->successors[1];
+        AssertThat(ts_parse_stack_head_count(stack), Equals(2));
-        AssertThat(successor2->state, Equals(stateF));
+        AssertThat(*ts_parse_stack_head(stack, 0), Equals<ParseStackEntry>({trees[2], stateC}));
-        AssertThat(successor2->tree, Equals(trees[3]))
+        AssertThat(*ts_parse_stack_head(stack, 1), Equals<ParseStackEntry>({trees[4], stateE}));
-        AssertThat(successor2->successor_count, Equals(1));
       });
     });
 
-    describe("when both heads are reduced into the same state with the same symbol and yield", [&]() {
+    describe("when the trees are different", [&]() {
-      before_each([&]() {
+      it("merges the heads by creating an ambiguity node", [&]() {
-        //  A0__G5
+        /*
-        merged = ts_parse_stack_reduce(stack, 0, stateG, symbol5, 3);
+         *  A0__B1__C2__D3__G(6|7)
+         *       \__E4__F5____/
+         */
+        merged = ts_parse_stack_push(stack, 0, stateG, trees[6]);
         AssertThat(merged, IsFalse());
-        merged = ts_parse_stack_reduce(stack, 1, stateG, symbol5, 2);
+        merged = ts_parse_stack_push(stack, 1, stateG, trees[7]);
         AssertThat(merged, IsTrue());
-      });
 
-      it("re-joins the heads, creating an ambiguity node", [&]() {
         AssertThat(ts_parse_stack_head_count(stack), Equals(1));
-
+        AssertThat(*ts_parse_stack_head(stack, 0), Equals<ParseStackEntry>({
-        ParseStackNode *head = ts_parse_stack_head(stack, 0);
+          ts_tree_make_ambiguity(2, tree_array({ trees[6], trees[7] })),
-        AssertThat(head->state, Equals(stateG));
+          stateG
-
+        }));
-        AssertThat(head->tree, Fulfills(EqualsTree(
-          ts_tree_make_ambiguity(2, tree_array({
-            ts_tree_make_node(symbol5, 3, tree_array({
-              trees[1],
-              trees[2],
-              trees[3],
-            }), false),
-            ts_tree_make_node(symbol5, 2, tree_array({
-              ts_tree_make_node(symbol4, 2, tree_array({
-                trees[1],
-                trees[2],
-              }), false),
-              trees[3]
-            }), false)
-          })),
-          symbol_names)));
-        AssertThat(head->successor_count, Equals(1));
-      });
-    });
-
-    describe("when a head with multiple paths is reduced", [&]() {
-      before_each([&]() {
-        //  A0__B1__C2__D3__G5
-        //   \______E4__F3__/
-        ts_parse_stack_shift(stack, 0, stateG, trees[5]);
-        ts_parse_stack_shift(stack, 1, stateG, trees[5]);
-      });
-
-      it("reduces along all paths, creating an ambiguity node", [&]() {
-        //  A0__B1__C2__H6
-        //   \______E4__/
-        ts_parse_stack_reduce(stack, 0, stateH, symbol6, 2);
-        AssertThat(ts_parse_stack_head_count(stack), Equals(1));
-
-        ParseStackNode *head = ts_parse_stack_head(stack, 0);
-        AssertThat(head->state, Equals(stateH));
-        AssertThat(head->tree, Fulfills(EqualsTree(
-          ts_tree_make_ambiguity(2, tree_array({
-            ts_tree_make_node(symbol6, 2, tree_array({
-              trees[3],
-              trees[5],
-            }), false),
-            ts_tree_make_node(symbol6, 2, tree_array({
-              trees[3],
-              trees[5],
-            }), false)
-          })),
-          symbol_names)));
-        AssertThat(head->successor_count, Equals(2));
       });
     });
   });
 });
 
 END_TEST
+
+bool operator==(const ParseStackEntry &left, const ParseStackEntry &right) {
+  return left.state == right.state && ts_tree_eq(left.tree, right.tree);
+}
+
+std::ostream &operator<<(std::ostream &stream, const ParseStackEntry &entry) {
+  return stream << "{" << entry.state << ", " << entry.tree << "}";
+}

src/runtime/parse_stack.c

@@ -5,16 +5,25 @@
 #include "runtime/length.h"
 #include <assert.h>
 
-static const size_t INITIAL_HEAD_CAPACITY = 3;
+#define MAX_POP_PATH_COUNT 8
+#define INITIAL_HEAD_CAPACITY 3
+
+typedef struct ParseStackNode {
+  ParseStackEntry entry;
+  struct ParseStackNode *successors[MAX_POP_PATH_COUNT];
+  short unsigned int successor_count;
+  short unsigned int ref_count;
+} ParseStackNode;
 
 struct ParseStack {
   ParseStackNode **heads;
   int head_count;
   int head_capacity;
+  ParseStackPopResult last_pop_results[MAX_POP_PATH_COUNT];
 };
 
 /*
- *  Section: Lifecycle
+ *  Section: Stack lifecycle
  */
 
 ParseStack *ts_parse_stack_new() {
@@ -28,163 +37,36 @@ ParseStack *ts_parse_stack_new() {
 }
 
 void ts_parse_stack_delete(ParseStack *this) {
-  if (this->heads)
+  free(this->heads);
-    free(this->heads);
   free(this);
 }
 
 /*
- *  Section: Reading
+ *  Section: Reading from the stack
  */
 
-ParseStackNode *ts_parse_stack_head(const ParseStack *this, int head_index) {
+const ParseStackEntry *ts_parse_stack_head(const ParseStack *this, int head) {
-  assert(head_index < this->head_count);
+  assert(head < this->head_count);
-  return this->heads[head_index];
+  ParseStackNode *node = this->heads[head];
+  return node ? &node->entry : NULL;
 }
 
 int ts_parse_stack_head_count(const ParseStack *this) {
   return this->head_count;
 }
 
-/*
+int ts_parse_stack_entry_next_count(const ParseStackEntry *entry) {
- *  Section: Updating
+  return ((const ParseStackNode *)entry)->successor_count;
- */
-
-static ParseStackNode *stack_node_new(ParseStackNode *, TSStateId, TSTree *);
-static void stack_node_retain(ParseStackNode *);
-static bool stack_node_release(ParseStackNode *);
-static void stack_node_add_successor(ParseStackNode *, ParseStackNode *);
-static void parse_stack_remove_head(ParseStack *, int);
-static bool parse_stack_merge_head(ParseStack *, int, TSStateId, TSTree *);
-
-bool ts_parse_stack_shift(ParseStack *this, int head_index, TSStateId state, TSTree *tree) {
-  assert(head_index < this->head_count);
-  if (parse_stack_merge_head(this, head_index, state, tree))
-    return true;
-  this->heads[head_index] = stack_node_new(this->heads[head_index], state, tree);
-  return false;
 }
 
-#define MAX_PATH_COUNT 8
+const ParseStackEntry *ts_parse_stack_entry_next(const ParseStackEntry *entry, int i) {
-
+  return &((const ParseStackNode *)entry)->successors[i]->entry;
-bool ts_parse_stack_reduce(ParseStack *this, int head_index, TSStateId state,
-                           TSSymbol symbol, int child_count) {
-  int path_count = 1;
-  ParseStackNode *nodes_by_path[MAX_PATH_COUNT] = {this->heads[head_index]};
-  TreeVector children_by_path[MAX_PATH_COUNT] = {tree_vector_new(child_count)};
-  size_t child_counts_by_path[MAX_PATH_COUNT] = {child_count};
-
-  /*
-   *  Reduce along every possible path in parallel. Stop when the given number
-   *  of child trees have been collected along every path.
-   */
-  bool all_paths_done = false;
-  while (!all_paths_done) {
-    all_paths_done = true;
-    int current_path_count = path_count;
-    for (int path = 0; path < current_path_count; path++) {
-      if (children_by_path[path].size == child_counts_by_path[path])
-        continue;
-      else
-        all_paths_done = false;
-
-      /*
-       *  Children that are 'extra' do not count towards the total child count.
-       */
-      ParseStackNode *node = nodes_by_path[path];
-      if (ts_tree_is_extra(node->tree))
-        child_counts_by_path[path]++;
-
-      /*
-       *  If a node has more than one successor, create new paths for each of
-       *  the additional successors.
-       */
-      tree_vector_push(&children_by_path[path], node->tree);
-
-      for (int i = 0; i < node->successor_count; i++) {
-        int next_path;
-        if (i > 0) {
-          if (path_count == MAX_PATH_COUNT) break;
-          next_path = path_count;
-          child_counts_by_path[next_path] = child_counts_by_path[path];
-          children_by_path[next_path] = tree_vector_copy(&children_by_path[path]);
-          path_count++;
-        } else {
-          next_path = path;
-        }
-
-        nodes_by_path[next_path] = node->successors[i];
-      }
-    }
-  }
-
-  TSTree *parent;
-  if (path_count > 1) {
-    TSTree **trees_by_path = malloc(path_count * sizeof(TSTree *));
-    for (int path = 0; path < path_count; path++) {
-      stack_node_retain(nodes_by_path[path]);
-      tree_vector_reverse(&children_by_path[path]);
-      trees_by_path[path] = ts_tree_make_node(
-        symbol,
-        child_counts_by_path[path],
-        children_by_path[path].contents,
-        false
-      );
-      parent = ts_tree_make_ambiguity(path_count, trees_by_path);
-    }
-  } else {
-    stack_node_retain(nodes_by_path[0]);
-    tree_vector_reverse(&children_by_path[0]);
-    parent = ts_tree_make_node(
-      symbol,
-      child_counts_by_path[0],
-      children_by_path[0].contents,
-      false
-    );
-  }
-
-  stack_node_release(this->heads[head_index]);
-  this->heads[head_index] = nodes_by_path[0];
-
-  if (parse_stack_merge_head(this, head_index, state, parent))
-    return true;
-
-  this->heads[head_index] = stack_node_new(nodes_by_path[0], state, parent);
-  for (int i = 1; i < path_count; i++) {
-    stack_node_add_successor(this->heads[head_index], nodes_by_path[i]);
-  }
-
-  return false;
-}
-
-int ts_parse_stack_split(ParseStack *this, int head_index) {
-  assert(head_index < this->head_count);
-  if (this->head_count == this->head_capacity) {
-    this->head_capacity += 3;
-    this->heads = realloc(this->heads, this->head_capacity * sizeof(ParseStackNode *));
-  }
-  int new_index = this->head_count++;
-  this->heads[new_index] = this->heads[head_index];
-  stack_node_retain(this->heads[new_index]);
-  return new_index;
 }
 
 /*
- *  Section: Private
+ *  Section: Manipulating nodes (Private)
  */
 
-static ParseStackNode *stack_node_new(ParseStackNode *next, TSStateId state, TSTree *tree) {
-  ParseStackNode *this = malloc(sizeof(ParseStackNode));
-  *this = (ParseStackNode) {
-    .ref_count = 1,
-    .successor_count = 1,
-    .successors = {next, NULL, NULL},
-    .state = state,
-    .tree = tree,
-  };
-  return this;
-}
-
 static void stack_node_retain(ParseStackNode *this) {
   if (!this) return;
   assert(this->ref_count != 0);
@@ -198,7 +80,7 @@ static bool stack_node_release(ParseStackNode *this) {
   if (this->ref_count == 0) {
     for (int i = 0; i < this->successor_count; i++)
       stack_node_release(this->successors[i]);
-    ts_tree_release(this->tree);
+    ts_tree_release(this->entry.tree);
     free(this);
     return true;
   } else {
@@ -206,33 +88,69 @@ static bool stack_node_release(ParseStackNode *this) {
   }
 }
 
+static ParseStackNode *stack_node_new(ParseStackNode *next, TSStateId state, TSTree *tree) {
+  ParseStackNode *this = malloc(sizeof(ParseStackNode));
+  ts_tree_retain(tree);
+  stack_node_retain(next);
+  *this = (ParseStackNode) {
+    .ref_count = 1,
+    .successor_count = 1,
+    .successors = {next, NULL, NULL},
+    .entry = {
+      .state = state,
+      .tree = tree,
+    },
+  };
+  return this;
+}
+
 static void stack_node_add_successor(ParseStackNode *this, ParseStackNode *successor) {
-  stack_node_retain(successor);
   for (int i = 0; i < this->successor_count; i++)
     if (this->successors[i] == successor)
       return;
+  stack_node_retain(successor);
   this->successors[this->successor_count] = successor;
   this->successor_count++;
 }
 
-static bool parse_stack_merge_head(ParseStack *this, int head_index, TSStateId state, TSTree *tree) {
+/*
+ *  Section: Mutating the stack (Private)
+ */
+
+static int ts_parse_stack_add_head(ParseStack *this, ParseStackNode *node) {
+  if (this->head_count == this->head_capacity) {
+    this->head_capacity += 3;
+    this->heads = realloc(this->heads, this->head_capacity * sizeof(ParseStackNode *));
+  }
+  int new_index = this->head_count++;
+  this->heads[new_index] = node;
+  stack_node_retain(node);
+  return new_index;
+}
+
+static void ts_parse_stack_remove_head(ParseStack *this, int head_index) {
+  stack_node_release(this->heads[head_index]);
+  for (int i = head_index; i < this->head_count - 1; i++) {
+    this->heads[head_index] = this->heads[head_index + 1];
+  }
+  this->head_count--;
+}
+
+static bool ts_parse_stack_merge_head(ParseStack *this, int head_index, TSStateId state, TSTree *tree) {
   for (int i =  0; i < head_index; i++) {
     ParseStackNode *head = this->heads[i];
-    if (head->state == state) {
+    if (head->entry.state == state) {
-      if (head->tree == tree) {
+      if (head->entry.tree == tree) {
         stack_node_add_successor(head, this->heads[head_index]);
-        parse_stack_remove_head(this, head_index);
+        ts_parse_stack_remove_head(this, head_index);
         return true;
-      }
+      } else {
-
-      if (head->tree->symbol == tree->symbol &&
-          ts_length_eq(head->tree->size, tree->size)) {
         TSTree **options = malloc(2 * sizeof(TSTree *));
-        options[0] = head->tree;
+        options[0] = head->entry.tree;
         options[1] = tree;
-        head->tree = ts_tree_make_ambiguity(2, options);
+        head->entry.tree = ts_tree_make_ambiguity(2, options);
         stack_node_add_successor(head, this->heads[head_index]);
-        parse_stack_remove_head(this, head_index);
+        ts_parse_stack_remove_head(this, head_index);
         return true;
       }
     }
@@ -240,10 +158,96 @@ static bool parse_stack_merge_head(ParseStack *this, int head_index, TSStateId s
   return false;
 }
 
-static void parse_stack_remove_head(ParseStack *this, int head_index) {
+/*
-  stack_node_release(this->heads[head_index]);
+ *  Section: Mutating the stack (Public)
-  for (int i = head_index; i < this->head_count - 1; i++) {
+ */
-    this->heads[head_index] = this->heads[head_index + 1];
+
-  }
+bool ts_parse_stack_push(ParseStack *this, int head_index, TSStateId state, TSTree *tree) {
-  this->head_count--;
+  assert(head_index < this->head_count);
+  if (ts_parse_stack_merge_head(this, head_index, state, tree))
+    return true;
+  this->heads[head_index] = stack_node_new(this->heads[head_index], state, tree);
+  return false;
+}
+
+int ts_parse_stack_split(ParseStack *this, int head_index) {
+  assert(head_index < this->head_count);
+  return ts_parse_stack_add_head(this, this->heads[head_index]);
+}
+
+ParseStackPopResultList ts_parse_stack_pop(ParseStack *this, int head_index, int child_count) {
+  ParseStackNode *previous_head = this->heads[head_index];
+
+  int path_count = 1;
+  size_t tree_counts_by_path[MAX_POP_PATH_COUNT] = {child_count};
+  TreeVector trees_by_path[MAX_POP_PATH_COUNT] = {tree_vector_new(child_count)};
+  ParseStackNode *nodes_by_path[MAX_POP_PATH_COUNT] = {previous_head};
+
+  /*
+   *  Reduce along every possible path in parallel. Stop when the given number
+   *  of child trees have been collected along every path.
+   */
+  bool all_paths_done = false;
+  while (!all_paths_done) {
+    all_paths_done = true;
+    int current_path_count = path_count;
+    for (int path = 0; path < current_path_count; path++) {
+      if (trees_by_path[path].size == tree_counts_by_path[path])
+        continue;
+      else
+        all_paths_done = false;
+
+      /*
+       *  Children that are 'extra' do not count towards the total child count.
+       */
+      ParseStackNode *node = nodes_by_path[path];
+      if (ts_tree_is_extra(node->entry.tree))
+        tree_counts_by_path[path]++;
+
+      /*
+       *  If a node has more than one successor, create new paths for each of
+       *  the additional successors.
+       */
+      tree_vector_push(&trees_by_path[path], node->entry.tree);
+
+      for (int i = 0; i < node->successor_count; i++) {
+        int next_path;
+        if (i > 0) {
+          if (path_count == MAX_POP_PATH_COUNT) break;
+          next_path = path_count;
+          tree_counts_by_path[next_path] = tree_counts_by_path[path];
+          trees_by_path[next_path] = tree_vector_copy(&trees_by_path[path]);
+          path_count++;
+        } else {
+          next_path = path;
+        }
+
+        nodes_by_path[next_path] = node->successors[i];
+      }
+    }
+  }
+
+  for (int path = 0; path < path_count; path++) {
+    tree_vector_reverse(&trees_by_path[path]);
+    int index;
+    if (path == 0) {
+      stack_node_retain(nodes_by_path[path]);
+      this->heads[head_index] = nodes_by_path[path];
+      index = head_index;
+    } else {
+      index = ts_parse_stack_add_head(this, nodes_by_path[path]);
+    }
+
+    this->last_pop_results[path] = (ParseStackPopResult) {
+      .index = index,
+      .tree_count = trees_by_path[path].size,
+      .trees = trees_by_path[path].contents,
+    };
+  }
+
+  stack_node_release(previous_head);
+  return (ParseStackPopResultList) {
+    .size = path_count,
+    .contents = this->last_pop_results,
+  };
 }

src/runtime/parse_stack.h

@@ -9,23 +9,73 @@ extern "C" {
 
 typedef struct ParseStack ParseStack;
 
-typedef struct ParseStackNode {
+typedef struct {
   TSTree *tree;
   TSStateId state;
-  struct ParseStackNode *successors[4];
+} ParseStackEntry;
-  short unsigned int successor_count;
-  short unsigned int ref_count;
-} ParseStackNode;
 
+typedef struct {
+  int index;
+  int tree_count;
+  TSTree **trees;
+} ParseStackPopResult;
+
+typedef struct {
+  int size;
+  ParseStackPopResult *contents;
+} ParseStackPopResultList;
+
+/*
+ *  Create a ParseStack.
+ */
 ParseStack *ts_parse_stack_new();
+
+/*
+ *  Release any resources reserved by a parse stack.
+ */
 void ts_parse_stack_delete(ParseStack *);
 
-ParseStackNode *ts_parse_stack_head(const ParseStack *, int);
+/*
+ *  Get the stack's current number of heads.
+ */
 int ts_parse_stack_head_count(const ParseStack *);
 
-bool ts_parse_stack_shift(ParseStack *, int, TSStateId, TSTree *);
+/*
-bool ts_parse_stack_reduce(ParseStack *, int, TSStateId, TSSymbol, int);
+ *  Get the tree and state that are at the top of the given stack head.
-int ts_parse_stack_split(ParseStack *, int);
+ */
+const ParseStackEntry *ts_parse_stack_head(const ParseStack *, int head);
+
+/*
+ *  Get the number of successors for a given parse stack entry.
+ */
+int ts_parse_stack_entry_next_count(const ParseStackEntry *);
+
+/*
+ *  Get the nth successor to a given parse stack entry.
+ */
+const ParseStackEntry *ts_parse_stack_entry_next(const ParseStackEntry *, int);
+
+/*
+ *  Push a (tree, state) pair onto the given head of the stack. Returns
+ *  a boolean indicating whether the stack head was merged with an
+ *  existing head.
+ */
+bool ts_parse_stack_push(ParseStack *, int head, TSStateId, TSTree *);
+
+/*
+ *  Pop the given number of entries from the given head of the stack. This
+ *  operation can increase the number of stack heads by revealing multiple heads
+ *  which had previously been merged. It returns a struct that indicates the
+ *  index of each revealed head and the trees removed from that head.
+ */
+ParseStackPopResultList ts_parse_stack_pop(ParseStack *, int head, int count);
+
+/*
+ *  Split the given stack head into two heads, so that the stack can be
+ *  transformed from its current state in multiple alternative ways. Returns
+ *  the index of the newly-created head.
+ */
+int ts_parse_stack_split(ParseStack *, int head);
 
 #ifdef __cplusplus
 }

src/runtime/tree.c

@@ -153,6 +153,12 @@ TSLength ts_tree_total_size(const TSTree *tree) {
 }
 
 bool ts_tree_eq(const TSTree *node1, const TSTree *node2) {
+  if (node1) {
+    if (!node2) return false;
+  } else {
+    return !node2;
+  }
+
   if (node1->symbol != node2->symbol)
     return false;
   if (node1->symbol == ts_builtin_sym_error)