From 134c455b80782a9bd2ab76eb6e03e33dfb0768e8 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Mon, 12 Feb 2018 11:52:06 -0800
Subject: [PATCH] Simplify logic for terminating tree balancing

---
 src/runtime/tree.c | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/runtime/tree.c b/src/runtime/tree.c
index dc82960d..d1b551c8 100644
--- a/src/runtime/tree.c
+++ b/src/runtime/tree.c
@@ -206,15 +206,13 @@ Tree *ts_tree_make_copy(TreePool *pool, Tree *self) {
 static void ts_tree__compress(Tree *self, unsigned count, const TSLanguage *language) {
   Tree *tree = self;
   for (unsigned i = 0; i < count; i++) {
-    if (tree->child_count != 2) break;
     Tree *child = tree->children[0];
-    if (child->symbol != tree->symbol || child->child_count != 2) break;
-    if (tree->ref_count > 1) break;
-    if (child->ref_count > 1) break;
+    if (child->symbol != tree->symbol) break;
 
     Tree *grandchild = child->children[0];
-    if (grandchild->symbol != tree->symbol || grandchild->child_count != 2) break;
-    if (grandchild->ref_count > 1) break;
+    if (grandchild->symbol != tree->symbol) break;
+    if (grandchild->children[0]->symbol != tree->symbol) break;
+    if (child->ref_count > 1 || grandchild->ref_count > 1) break;
 
     tree->children[0] = grandchild;
     grandchild->context.parent = tree;