From 7b0a52ec268ca9799ca82339349067d92478d98f Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 27 Aug 2014 12:56:36 -0700
Subject: [PATCH] Pretty-print single hidden tree nodes correctly

---
 spec/runtime/tree_spec.cc | 34 ++++++++++++++++++++++++++++------
 src/runtime/tree.c        | 28 ++++++++++++----------------
 2 files changed, 40 insertions(+), 22 deletions(-)

diff --git a/spec/runtime/tree_spec.cc b/spec/runtime/tree_spec.cc
index fabb97e6..734b72d2 100644
--- a/spec/runtime/tree_spec.cc
+++ b/spec/runtime/tree_spec.cc
@@ -11,15 +11,13 @@ enum {
 
 static const char *names[] = { "error", "end", "cat", "dog", "pig" };
 
-describe("trees", []() {
+describe("Tree", []() {
   TSTree *tree1, *tree2, *parent1;
 
   before_each([&]() {
     tree1 = ts_tree_make_leaf(cat, 5, 2, 0);
     tree2 = ts_tree_make_leaf(cat, 3, 1, 0);
-    parent1 = ts_tree_make_node(dog, 2, tree_array({
-        tree1, tree2, // children
-    }), 0);
+    parent1 = ts_tree_make_node(dog, 2, tree_array({ tree1, tree2, }), 0);
   });
 
   after_each([&]() {
@@ -77,14 +75,38 @@ describe("trees", []() {
 
   describe("serialization", [&]() {
     it("returns a readable string", [&]() {
-      auto string1 = ts_tree_string(tree1, names);
+      char *string1 = ts_tree_string(tree1, names);
       AssertThat(string(string1), Equals("(cat)"));
       free(string1);
 
-      auto string2 = ts_tree_string(parent1, names);
+      char *string2 = ts_tree_string(parent1, names);
       AssertThat(string(string2), Equals("(dog (cat) (cat))"));
       free(string2);
     });
+
+    it("hides invisible nodes", [&]() {
+      tree2->options = TSTreeOptionsHidden;
+
+      char *string1 = ts_tree_string(parent1, names);
+      AssertThat(string(string1), Equals("(dog (cat))"));
+      free(string1);
+    });
+
+    describe("when the root node is not visible", [&]() {
+      it("still serializes it", [&]() {
+        parent1->options = TSTreeOptionsHidden;
+
+        char *string1 = ts_tree_string(parent1, names);
+        AssertThat(string(string1), Equals("(dog (cat) (cat))"));
+        free(string1);
+
+        tree1->options = TSTreeOptionsHidden;
+
+        char *string2 = ts_tree_string(tree1, names);
+        AssertThat(string(string2), Equals("(cat)"));
+        free(string2);
+      });
+    });
   });
 });
 
diff --git a/src/runtime/tree.c b/src/runtime/tree.c
index bc462777..69be08f8 100644
--- a/src/runtime/tree.c
+++ b/src/runtime/tree.c
@@ -161,27 +161,23 @@ static size_t tree_write_to_string(const TSTree *tree,
   if (visible && !is_root)
     cursor += snprintf(*writer, limit, " ");
 
-  if (!tree)
-    return snprintf(*writer, limit, "(NULL)");
-  if (tree->symbol == ts_builtin_sym_error) {
+  if (!tree) {
+    cursor += snprintf(*writer, limit, "(NULL)");
+  } else if (tree->symbol == ts_builtin_sym_error) {
     cursor += snprintf(*writer, limit, "(ERROR ");
     cursor += write_lookahead_to_string(*writer, limit, tree->lookahead_char);
     cursor += snprintf(*writer, limit, ")");
-    return cursor - string;
+  } else {
+    if (visible || is_root)
+      cursor += snprintf(*writer, limit, "(%s", symbol_names[tree->symbol]);
+    for (size_t i = 0; i < tree->child_count; i++) {
+      TSTree *child = tree->children[i];
+      cursor += tree_write_to_string(child, symbol_names, *writer, limit, 0);
+    }
+    if (visible || is_root)
+      cursor += snprintf(*writer, limit, ")");
   }
 
-  if (visible) {
-    cursor += snprintf(*writer, limit, "(%s", symbol_names[tree->symbol]);
-    is_root = 0;
-  }
-
-  for (size_t i = 0; i < tree->child_count; i++)
-    cursor += tree_write_to_string(tree->children[i], symbol_names, *writer,
-                                   limit, is_root);
-
-  if (visible)
-    cursor += snprintf(*writer, limit, ")");
-
   return cursor - string;
 }