From fcf9293d35ff7d87ce72488a3fa78e050426199e Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 14 Sep 2016 09:46:41 -0700
Subject: [PATCH] Use explicit stack for assigning trees' parent pointers

---
 src/runtime/parser.c |  5 ++++-
 src/runtime/parser.h |  1 +
 src/runtime/tree.c   | 35 ++++++++++++++++++-----------------
 src/runtime/tree.h   |  2 +-
 4 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/src/runtime/parser.c b/src/runtime/parser.c
index 7e3c9a0d..09a8757a 100644
--- a/src/runtime/parser.c
+++ b/src/runtime/parser.c
@@ -1224,6 +1224,7 @@ bool parser_init(Parser *self) {
   self->finished_tree = NULL;
   self->stack = NULL;
   array_init(&self->reduce_actions);
+  array_init(&self->tree_stack);
 
   self->stack = ts_stack_new();
   if (!self->stack)
@@ -1249,6 +1250,8 @@ void parser_destroy(Parser *self) {
     ts_stack_delete(self->stack);
   if (self->reduce_actions.contents)
     array_delete(&self->reduce_actions);
+  if (self->tree_stack.contents)
+    array_delete(&self->tree_stack);
 }
 
 TSTree *parser_parse(Parser *self, TSInput input, TSTree *old_tree) {
@@ -1294,7 +1297,7 @@ TSTree *parser_parse(Parser *self, TSInput input, TSTree *old_tree) {
   LOG_TREE();
   ts_stack_clear(self->stack);
   parser__clear_cached_token(self);
-  ts_tree_assign_parents(self->finished_tree);
+  CHECK(ts_tree_assign_parents(self->finished_tree, &self->tree_stack));
   return self->finished_tree;
 
 error:
diff --git a/src/runtime/parser.h b/src/runtime/parser.h
index 146ee6a4..4bc22697 100644
--- a/src/runtime/parser.h
+++ b/src/runtime/parser.h
@@ -26,6 +26,7 @@ typedef struct {
   TSTree *cached_token;
   size_t cached_token_byte_index;
   ReusableNode reusable_node;
+  TreeArray tree_stack;
 } Parser;
 
 bool parser_init(Parser *);
diff --git a/src/runtime/tree.c b/src/runtime/tree.c
index 85c545d7..2bf46d0f 100644
--- a/src/runtime/tree.c
+++ b/src/runtime/tree.c
@@ -91,28 +91,29 @@ TSTree *ts_tree_make_copy(TSTree *self) {
   return result;
 }
 
-void ts_tree_assign_parents(TSTree *self) {
-  TSLength offset;
-
-recur:
-  offset = ts_length_zero();
-  for (size_t i = 0; i < self->child_count; i++) {
-    TSTree *child = self->children[i];
-    if (child->context.parent != self || child->context.index != i) {
-      child->context.parent = self;
-      child->context.index = i;
-      child->context.offset = offset;
-      if (i == self->child_count - 1) {
-        self = child;
-        goto recur;
+bool ts_tree_assign_parents(TSTree *self, TreeArray *stack) {
+  array_clear(stack);
+  if (!array_push(stack, self))
+    return false;
+  while (stack->size > 0) {
+    TSTree *tree = array_pop(stack);
+    TSLength offset = ts_length_zero();
+    for (size_t i = 0; i < tree->child_count; i++) {
+      TSTree *child = tree->children[i];
+      if (child->context.parent != tree || child->context.index != i) {
+        child->context.parent = tree;
+        child->context.index = i;
+        child->context.offset = offset;
+        if (!array_push(stack, child))
+          return false;
       }
-
-      ts_tree_assign_parents(child);
+      offset = ts_length_add(offset, ts_tree_total_size(child));
     }
-    offset = ts_length_add(offset, ts_tree_total_size(child));
   }
+  return true;
 }
 
+
 void ts_tree_set_children(TSTree *self, size_t child_count, TSTree **children) {
   if (self->child_count > 0)
     ts_free(self->children);
diff --git a/src/runtime/tree.h b/src/runtime/tree.h
index af8d50d9..146897d8 100644
--- a/src/runtime/tree.h
+++ b/src/runtime/tree.h
@@ -67,7 +67,7 @@ int ts_tree_compare(const TSTree *tree1, const TSTree *tree2);
 size_t ts_tree_start_column(const TSTree *self);
 size_t ts_tree_end_column(const TSTree *self);
 void ts_tree_set_children(TSTree *, size_t, TSTree **);
-void ts_tree_assign_parents(TSTree *);
+bool ts_tree_assign_parents(TSTree *, TreeArray *);
 void ts_tree_edit(TSTree *, const TSInputEdit *edit);
 char *ts_tree_string(const TSTree *, const TSLanguage *, bool include_all);
 void ts_tree_print_dot_graph(const TSTree *, const TSLanguage *, FILE *);