From fbb9b24d7bf4b0d401d2a7332f79a5f673a40997 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Tue, 18 Mar 2014 12:47:26 -0700
Subject: [PATCH] Refactor ts_tree_children

---
 include/tree_sitter/runtime.h |  3 +--
 spec/runtime/json_spec.cc     |  6 +++---
 src/runtime/stack.c           | 15 +++++++++++----
 src/runtime/tree.c            | 24 +++++++++++-------------
 4 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/include/tree_sitter/runtime.h b/include/tree_sitter/runtime.h
index 64f39d80..807db435 100644
--- a/include/tree_sitter/runtime.h
+++ b/include/tree_sitter/runtime.h
@@ -37,8 +37,7 @@ void ts_tree_release(ts_tree *tree);
 int ts_tree_equals(const ts_tree *tree1, const ts_tree *tree2);
 char * ts_tree_string(const ts_tree *tree, const char **names);
 char * ts_tree_error_string(const ts_tree *tree, const char **names);
-size_t ts_tree_child_count(const ts_tree *tree);
-ts_tree ** ts_tree_children(const ts_tree *tree);
+ts_tree ** ts_tree_children(const ts_tree *tree, size_t *count);
 
 typedef struct {
     void *data;
diff --git a/spec/runtime/json_spec.cc b/spec/runtime/json_spec.cc
index b271390f..539f88c9 100644
--- a/spec/runtime/json_spec.cc
+++ b/spec/runtime/json_spec.cc
@@ -54,9 +54,9 @@ describe("json", []() {
             ts_document_set_input_string(doc, "  [12, 5]");
 
             const ts_tree *tree = ts_document_tree(doc);
-            const ts_tree *array = ts_tree_children(tree)[0];
-            const ts_tree *number1 = ts_tree_children(array)[0];
-            const ts_tree *number2 = ts_tree_children(array)[1];
+            const ts_tree *array = ts_tree_children(tree, NULL)[0];
+            const ts_tree *number1 = ts_tree_children(array, NULL)[0];
+            const ts_tree *number2 = ts_tree_children(array, NULL)[1];
 
             AssertThat(number1->offset, Equals<size_t>(0));
             AssertThat(number1->size, Equals<size_t>(2));
diff --git a/src/runtime/stack.c b/src/runtime/stack.c
index c1dd7666..63a267ba 100644
--- a/src/runtime/stack.c
+++ b/src/runtime/stack.c
@@ -50,7 +50,13 @@ ts_tree * ts_stack_reduce(ts_stack *stack, ts_symbol symbol, int immediate_child
     int child_count = 0;
     for (int i = 0; i < immediate_child_count; i++) {
         ts_tree *child = stack->entries[new_stack_size + i].node;
-        child_count += collapse_flags[i] ? ts_tree_child_count(child) : 1;
+        if (collapse_flags[i]) {
+            size_t grandchild_count;
+            ts_tree_children(child, &grandchild_count);
+            child_count += grandchild_count;
+        } else {
+            child_count++;
+        }
     }
 
     int child_index = 0;
@@ -66,11 +72,12 @@ ts_tree * ts_stack_reduce(ts_stack *stack, ts_symbol symbol, int immediate_child
         }
 
         if (collapse_flags[i]) {
-            size_t grandchild_count = ts_tree_child_count(child);
-            memcpy(children + child_index, ts_tree_children(child), (grandchild_count * sizeof(ts_tree *)));
+            size_t grandchild_count;
+            ts_tree ** grandchildren = ts_tree_children(child, &grandchild_count);
+            memcpy(children + child_index, grandchildren, (grandchild_count * sizeof(ts_tree *)));
             child_index += grandchild_count;
         } else {
-            memcpy(children + child_index, &child, sizeof(ts_tree *));
+            children[child_index] = child;
             child_index++;
         }
     }
diff --git a/src/runtime/tree.c b/src/runtime/tree.c
index 48f35449..f3387d2c 100644
--- a/src/runtime/tree.c
+++ b/src/runtime/tree.c
@@ -42,12 +42,11 @@ void ts_tree_retain(ts_tree *tree) {
 void ts_tree_release(ts_tree *tree) {
     tree->ref_count--;
     if (tree->ref_count == 0) {
-        ts_tree **children = ts_tree_children(tree);
-        if (children) {
-            for (size_t i = 0; i < ts_tree_child_count(tree); i++)
-                ts_tree_release(children[i]);
-            free(children);
-        }
+        size_t count;
+        ts_tree **children = ts_tree_children(tree, &count);
+        for (size_t i = 0; i < count; i++)
+            ts_tree_release(children[i]);
+        free(children);
         free(tree);
     }
 }
@@ -69,16 +68,15 @@ int ts_tree_equals(const ts_tree *node1, const ts_tree *node2) {
     return 1;
 }
 
-ts_tree ** ts_tree_children(const ts_tree *tree) {
-    if (tree->symbol == ts_builtin_sym_error) return NULL;
+ts_tree ** ts_tree_children(const ts_tree *tree, size_t *count) {
+    if (tree->symbol == ts_builtin_sym_error) {
+        if (count) *count = 0;
+        return NULL;
+    }
+    if (count) *count = tree->data.children.count;
     return tree->data.children.contents;
 }
 
-size_t ts_tree_child_count(const ts_tree *tree) {
-    if (tree->symbol == ts_builtin_sym_error) return 0;
-    return tree->data.children.count;
-}
-
 static size_t tree_write_to_string(const ts_tree *tree, const char **symbol_names, char *string, size_t limit) {
     static const char *NULL_TREE_STRING = "(NULL)";
     static const char *ERROR_TREE_STRING = "(ERROR)";