From 6e551d6d9fda8331cdb73b17e5fc7ef44426e7dd Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Mon, 14 Jul 2014 20:46:20 -0700
Subject: [PATCH] Simplify handling of multiple top-level nodes after parsing

---
 spec/runtime/languages/arithmetic/errors.txt  |  4 +--
 .../runtime/languages/javascript/literals.txt |  2 +-
 spec/runtime/languages/json/errors.txt        |  2 +-
 src/runtime/parser.c                          | 33 +++++--------------
 4 files changed, 13 insertions(+), 28 deletions(-)

diff --git a/spec/runtime/languages/arithmetic/errors.txt b/spec/runtime/languages/arithmetic/errors.txt
index 160141b4..ab181bb9 100644
--- a/spec/runtime/languages/arithmetic/errors.txt
+++ b/spec/runtime/languages/arithmetic/errors.txt
@@ -3,7 +3,7 @@ recovers from errors at the top level
 =====================================================
 x * * y
 ---
-(ERROR '*')
+(expression (variable) (ERROR '*'))
 
 =====================================================
 recovers from errors inside parenthesized expressions
@@ -12,4 +12,4 @@ x + (y * + z) * 5
 ---
 (expression (sum
     (variable)
-    (product (group (ERROR '+')) (number))))
\ No newline at end of file
+    (product (group (ERROR '+')) (number))))
diff --git a/spec/runtime/languages/javascript/literals.txt b/spec/runtime/languages/javascript/literals.txt
index aa7e0fe4..812a0bf9 100644
--- a/spec/runtime/languages/javascript/literals.txt
+++ b/spec/runtime/languages/javascript/literals.txt
@@ -39,7 +39,7 @@ var x = {
             (statement_block (var_declaration (identifier) (identifier)))))))
 
 ==========================================
-parses comments
+parses comments z
 ==========================================
 // this is the beginning of the script.
 // here we go.
diff --git a/spec/runtime/languages/json/errors.txt b/spec/runtime/languages/json/errors.txt
index 3e1beb27..3bde5aa2 100644
--- a/spec/runtime/languages/json/errors.txt
+++ b/spec/runtime/languages/json/errors.txt
@@ -3,7 +3,7 @@ recovers from top-level errors
 ==========================================
 [}
 ---
-(ERROR '}')
+(value (ERROR <EOF>) (ERROR '}'))
 
 ==========================================
 recovers from unexpected tokens
diff --git a/src/runtime/parser.c b/src/runtime/parser.c
index 0cbd1678..a610f96f 100644
--- a/src/runtime/parser.c
+++ b/src/runtime/parser.c
@@ -173,32 +173,17 @@ int ts_parser_handle_error(TSParser *parser) {
 
 TSTree * ts_parser_tree_root(TSParser *parser) {
     TSStack *stack = &parser->stack;
-    TSTree *top_node = ts_stack_top_node(stack);
-    if (stack->size <= 1)
-        return top_node;
-    if (ts_tree_symbol(top_node) == ts_builtin_sym_error)
-        return top_node;
-
-    size_t immediate_child_count;
-    TSTree **immedate_children = ts_tree_immediate_children(top_node, &immediate_child_count);
-
-    stack->size--;
-    for (size_t i = 0; i < immediate_child_count; i++) {
-        TSTree *child = immedate_children[i];
-        child->is_extra = 0;
-        ts_tree_retain(child);
-        TSStateId state = ts_stack_top_state(stack);
-        TSStateId next_state = actions_for_state(parser->config, state)[ts_tree_symbol(child)].data.to_state;
-        ts_stack_push(stack, next_state, child);
+    size_t node_count = 0;
+    for (size_t i = 0; i < stack->size; i++) {
+        TSTree *node = stack->entries[i].node;
+        if (!parser->config.hidden_symbol_flags[node->symbol])
+            node_count++;
     }
 
-    TSTree *new_node = ts_stack_reduce(stack,
-                                        top_node->symbol,
-                                        stack->size,
-                                        parser->config.hidden_symbol_flags,
-                                        0);
-    ts_tree_release(top_node);
-    return new_node;
+    if (node_count > 1)
+        return ts_stack_reduce(stack, 2, stack->size, parser->config.hidden_symbol_flags, 0);
+    else
+        return ts_stack_top_node(stack);
 }
 
 TSParseAction ts_parser_next_action(TSParser *parser) {