Remove 'document' wrapper node

src/runtime/document.c

@@ -1,5 +1,6 @@
 #include "tree_sitter/parser.h"
 #include "runtime/node.h"
+#include "runtime/tree.h"
 #include "runtime/length.h"
 #include "runtime/parser.h"
 #include "runtime/string_input.h"
@@ -62,5 +63,8 @@ void ts_document_set_input_string(TSDocument *document, const char *text) {
 }
 
 TSNode ts_document_root_node(const TSDocument *document) {
-  return ts_node_make(document->tree, ts_length_zero());
+  TSTree *tree = document->tree;
+  while (tree && ts_tree_is_singleton(tree))
+    tree = tree->children[0];
+  return ts_node_make(tree, ts_length_zero());
 }

src/runtime/parser.c

@@ -305,7 +305,17 @@ static bool ts_parser__handle_error(TSParser *parser, int head) {
 }
 
 static TSTree *ts_parser__finish(TSParser *parser) {
-  return ts_parser__reduce(parser, 0, ts_builtin_sym_document, -1, false, true);
+  ParseStackPopResult pop_result =
+    ts_parse_stack_pop(parser->stack, 0, -1, true).contents[0];
+
+  TSTree **trees = pop_result.trees;
+  size_t extra_count = pop_result.tree_count - 1;
+  TSTree *root = trees[extra_count];
+  assert(root->child_count > 0);
+
+  ts_tree_prepend_children(root, extra_count, trees);
+  ts_parse_stack_push(parser->stack, 0, 0, root);
+  return root;
 }
 
 typedef enum {

src/runtime/tree.c

@@ -67,8 +67,8 @@ TSTree *ts_tree_make_node(TSSymbol symbol, size_t child_count,
   } else {
     if (is_hidden)
       options |= TSTreeOptionsHidden;
-    if (child_count == 1 && symbol != ts_builtin_sym_document &&
-        (ts_tree_is_visible(children[0]) || ts_tree_is_wrapper(children[0])))
+    if (child_count == 1 &&
+        (ts_tree_is_visible(children[0]) || ts_tree_is_singleton(children[0])))
       options |= (TSTreeOptionsSingleton | TSTreeOptionsHidden);
     if (child_count > 0) {
       if (ts_tree_is_fragile_left(children[0]))
@@ -183,3 +183,32 @@ char *ts_tree_string(const TSTree *tree, const char **symbol_names) {
   ts_tree__write_to_string(tree, symbol_names, result, size, 1);
   return result;
 }
+
+void ts_tree_prepend_children(TSTree *tree, size_t count, TSTree **children) {
+  if (count == 0)
+    return;
+
+  tree->size = ts_length_add(tree->size, tree->padding);
+
+  size_t visible_count = 0;
+  for (size_t i = 0; i < count; i++) {
+    if (i == 0)
+      tree->padding = children[i]->padding;
+    else
+      tree->size = ts_length_add(tree->size, children[i]->padding);
+    tree->size = ts_length_add(tree->size, children[i]->size);
+    if (ts_tree_is_visible(children[i]))
+      visible_count++;
+  }
+
+  size_t new_child_count = count + tree->child_count;
+  TSTree **new_children = realloc(children, new_child_count * sizeof(TSTree *));
+  memcpy(new_children + count, tree->children,
+         tree->child_count * sizeof(TSTree *));
+  free(tree->children);
+
+  ts_tree_unset_singleton(tree);
+  tree->children = new_children;
+  tree->visible_child_count += visible_count;
+  tree->child_count += count;
+}

src/runtime/tree.h

@@ -47,10 +47,14 @@ static inline bool ts_tree_is_visible(const TSTree *tree) {
   return !(tree->options & TSTreeOptionsHidden);
 }
 
-static inline bool ts_tree_is_wrapper(const TSTree *tree) {
+static inline bool ts_tree_is_singleton(const TSTree *tree) {
   return !!(tree->options & TSTreeOptionsSingleton);
 }
 
+static inline void ts_tree_unset_singleton(TSTree *tree) {
+  tree->options = (TSTreeOptions)(tree->options & ~TSTreeOptionsSingleton);
+}
+
 static inline void ts_tree_set_options(TSTree *tree, TSTreeOptions options) {
   tree->options = (TSTreeOptions)(tree->options | options);
 }
@@ -84,6 +88,7 @@ bool ts_tree_eq(const TSTree *tree1, const TSTree *tree2);
 char *ts_tree_string(const TSTree *tree, const char **names);
 char *ts_tree_error_string(const TSTree *tree, const char **names);
 TSLength ts_tree_total_size(const TSTree *tree);
+void ts_tree_prepend_children(TSTree *, size_t, TSTree **);
 
 static inline bool ts_tree_is_empty(TSTree *tree) {
   return ts_tree_total_size(tree).bytes == 0;