From 2109f0ed74606f9dad3b03b26048fbfdb5b48bfe Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Tue, 14 Jun 2016 14:46:49 -0700
Subject: [PATCH] Handle allocation failures when copying tree arrays

---
 src/runtime/array.h |  9 ---------
 src/runtime/stack.c |  9 ++++++---
 src/runtime/tree.c  | 20 +++++++++++++++-----
 src/runtime/tree.h  |  2 +-
 4 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/src/runtime/array.h b/src/runtime/array.h
index 7770711c..c68522cf 100644
--- a/src/runtime/array.h
+++ b/src/runtime/array.h
@@ -58,15 +58,6 @@ extern "C" {
 #define array_reverse(self) \
   array__reverse((VoidArray *)(self), array__elem_size(self))
 
-#define array_copy(self)                                               \
-  {                                                                    \
-    (self)->contents                                                   \
-      ? memcpy(ts_calloc((self)->capacity, array__elem_size(self)),    \
-               (self)->contents, (self)->size *array__elem_size(self)) \
-      : NULL,                                                          \
-      (self)->size, (self)->capacity,                                  \
-  }
-
 // Private
 
 typedef Array(void) VoidArray;
diff --git a/src/runtime/stack.c b/src/runtime/stack.c
index 3d25a895..ae3bd51a 100644
--- a/src/runtime/stack.c
+++ b/src/runtime/stack.c
@@ -198,8 +198,10 @@ INLINE StackPopResult stack__iter(Stack *self, StackVersion version,
       bool should_stop = action & StackIterateStop || node->link_count == 0;
 
       if (should_pop) {
-        TreeArray trees =
-          should_stop ? path->trees : ts_tree_array_copy(&path->trees);
+        TreeArray trees = path->trees;
+        if (!should_stop)
+          if (!ts_tree_array_copy(trees, &trees))
+            goto error;
         array_reverse(&trees);
         if (!ts_stack__add_slice(self, node, &trees))
           goto error;
@@ -224,7 +226,8 @@ INLINE StackPopResult stack__iter(Stack *self, StackVersion version,
           if (!array_push(&self->pop_paths, self->pop_paths.contents[i]))
             goto error;
           next_path = array_back(&self->pop_paths);
-          next_path->trees = ts_tree_array_copy(&next_path->trees);
+          if (!ts_tree_array_copy(next_path->trees, &next_path->trees))
+            goto error;
         }
 
         next_path->node = link.node;
diff --git a/src/runtime/tree.c b/src/runtime/tree.c
index b73f64ad..400e0c42 100644
--- a/src/runtime/tree.c
+++ b/src/runtime/tree.c
@@ -40,11 +40,21 @@ TSTree *ts_tree_make_leaf(TSSymbol sym, TSLength padding, TSLength size,
   return result;
 }
 
-TreeArray ts_tree_array_copy(TreeArray *self) {
-  TreeArray result = array_copy(self);
-  for (size_t i = 0; i < result.size; i++)
-    ts_tree_retain(result.contents[i]);
-  return result;
+bool ts_tree_array_copy(TreeArray self, TreeArray *dest) {
+  TSTree **contents = NULL;
+  if (self.capacity > 0) {
+    contents = ts_calloc(self.capacity, sizeof(TSTree *));
+    if (!contents)
+      return false;
+    memcpy(contents, self.contents, self.size * sizeof(TSTree *));
+    for (size_t i = 0; i < self.size; i++)
+      ts_tree_retain(contents[i]);
+  }
+
+  dest->size = self.size;
+  dest->capacity = self.capacity;
+  dest->contents = contents;
+  return true;
 }
 
 void ts_tree_array_delete(TreeArray *self) {
diff --git a/src/runtime/tree.h b/src/runtime/tree.h
index 9cf7c15a..c0899956 100644
--- a/src/runtime/tree.h
+++ b/src/runtime/tree.h
@@ -46,7 +46,7 @@ typedef struct TSTree {
 } TSTree;
 
 typedef Array(TSTree *) TreeArray;
-TreeArray ts_tree_array_copy(TreeArray *);
+bool ts_tree_array_copy(TreeArray, TreeArray *);
 void ts_tree_array_delete(TreeArray *);
 size_t ts_tree_array_essential_count(const TreeArray *);