From 2f3e92c9beb4981fdd6f4caab54cb641786940ec Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Mon, 4 Apr 2016 11:44:45 -0700
Subject: [PATCH] Add function for popping all nodes from the stack

---
 src/runtime/parser.c | 33 +++++++--------------------------
 src/runtime/stack.c  |  9 +++++++++
 src/runtime/stack.h  |  2 ++
 3 files changed, 18 insertions(+), 26 deletions(-)

diff --git a/src/runtime/parser.c b/src/runtime/parser.c
index 6e221010..99502c79 100644
--- a/src/runtime/parser.c
+++ b/src/runtime/parser.c
@@ -86,21 +86,6 @@ typedef enum {
   BreakdownAborted,
 } BreakdownResult;
 
-static StackSlice ts_parser__pop_one(TSParser *self, int head_index, int count) {
-  StackPopResult pop = ts_stack_pop_count(self->stack, head_index, count);
-  if (pop.status != StackPopSucceeded)
-    return (StackSlice){
-      .head_index = -1, .trees = array_new(),
-    };
-  assert(pop.slices.size > 0);
-  assert(pop.slices.contents[0].head_index == head_index);
-  for (size_t i = pop.slices.size - 1; i > 0; i--) {
-    ts_tree_array_delete(&pop.slices.contents[i].trees);
-    ts_stack_remove_head(self->stack, pop.slices.contents[i].head_index);
-  }
-  return pop.slices.contents[0];
-}
-
 static BreakdownResult ts_parser__breakdown_top_of_stack(TSParser *self,
                                                          int head) {
   TSTree *last_child = NULL;
@@ -586,12 +571,10 @@ static ParseActionResult ts_parser__start(TSParser *self, TSInput input,
 }
 
 static ParseActionResult ts_parser__accept(TSParser *self, int head) {
-  StackSlice slice = ts_parser__pop_one(self, head, -1);
-  if (!slice.trees.contents)
+  TreeArray trees = ts_stack_pop_all(self->stack, head);
+  if (!trees.contents)
     goto error;
 
-  TreeArray trees = slice.trees;
-
   for (size_t i = trees.size - 1; i + 1 > 0; i--) {
     if (!trees.contents[i]->extra) {
       TSTree *root = trees.contents[i];
@@ -601,7 +584,7 @@ static ParseActionResult ts_parser__accept(TSParser *self, int head) {
       if (!trees.size)
         array_delete(&trees);
 
-      ts_stack_remove_head(self->stack, slice.head_index);
+      ts_stack_remove_head(self->stack, head);
       int comparison = ts_parser__select_tree(self, self->finished_tree, root);
       if (comparison > 0) {
         ts_tree_release(self->finished_tree);
@@ -617,9 +600,7 @@ static ParseActionResult ts_parser__accept(TSParser *self, int head) {
   return ParseActionRemoved;
 
 error:
-  for (size_t i = 0; i < slice.trees.size; i++)
-    ts_tree_release(slice.trees.contents[i]);
-  array_delete(&slice.trees);
+  ts_tree_array_delete(&trees);
   return ParseActionFailed;
 }
 
@@ -638,11 +619,11 @@ static ParseActionResult ts_parser__handle_error(TSParser *self, int head,
       if (!ts_stack_push(self->stack, head, error, false, 0))
         goto error;
 
-      StackSlice slice = ts_parser__pop_one(self, head, -1);
-      if (!slice.trees.contents)
+      TreeArray trees = ts_stack_pop_all(self->stack, head);
+      if (!trees.contents)
         goto error;
       TSTree *parent = ts_tree_make_node(
-        ts_builtin_sym_start, slice.trees.size, slice.trees.contents,
+        ts_builtin_sym_start, trees.size, trees.contents,
         ts_language_symbol_metadata(self->language, ts_builtin_sym_start));
 
       if (!ts_stack_push(self->stack, head, parent, false, 0))
diff --git a/src/runtime/stack.c b/src/runtime/stack.c
index 687a1f13..f79db61c 100644
--- a/src/runtime/stack.c
+++ b/src/runtime/stack.c
@@ -515,6 +515,15 @@ void ts_stack_clear(Stack *self) {
   array_push(&self->heads, self->base_node);
 }
 
+TreeArray ts_stack_pop_all(Stack *self, int head_index) {
+  StackPopResult pop = ts_stack_pop_count(self, head_index, -1);
+  if (pop.status != StackPopSucceeded)
+    return (TreeArray)array_new();
+  assert(pop.slices.size == 1);
+  assert(pop.slices.contents[0].head_index == head_index);
+  return pop.slices.contents[0].trees;
+}
+
 void ts_stack_set_tree_selection_callback(Stack *self, void *payload,
                                           TreeSelectionFunction function) {
   self->tree_selection_payload = payload;
diff --git a/src/runtime/stack.h b/src/runtime/stack.h
index f352af69..8b479421 100644
--- a/src/runtime/stack.h
+++ b/src/runtime/stack.h
@@ -94,6 +94,8 @@ StackPopResult ts_stack_pop_until(Stack *, int head_index, StackIterateCallback,
 
 StackPopResult ts_stack_pop_pending(Stack *, int head_index);
 
+TreeArray ts_stack_pop_all(Stack *, int head_index);
+
 /*
  *  Remove the given number of entries from the given head of the stack.
  */