Precompute trees' visible children

spec/runtime/document_spec.cc

@@ -86,6 +86,8 @@ describe("incremental parsing", [&]() {
         ts_document_set_input_string(doc, "  [12, 5, 345]");
 
         TSNode *array = ts_document_root_node(doc);
+        AssertThat(ts_node_string(array), Equals("(array (number) (number) (number))"));
+
         TSNode *number1 = ts_node_child(array, 0);
         TSNode *number2 = ts_node_child(array, 1);
         TSNode *number3 = ts_node_child(array, 2);
@@ -93,6 +95,7 @@ describe("incremental parsing", [&]() {
         AssertThat(ts_node_name(array), Equals("array"));
         AssertThat(ts_node_name(number1), Equals("number"));
         AssertThat(ts_node_name(number2), Equals("number"));
+        AssertThat(ts_node_name(number3), Equals("number"));
 
         AssertThat(ts_node_pos(array), Equals<size_t>(2));
         AssertThat(ts_node_size(array), Equals<size_t>(12));

src/runtime/document.c

@@ -95,22 +95,7 @@ void ts_document_set_input_string(TSDocument *document, const char *text) {
 }
 
 TSNode * ts_document_root_node(const TSDocument *document) {
-    const TSTree *tree = document->tree;
-    size_t position = 0;
-    while (ts_tree_is_wrapper(tree)) {
-        position = tree->offset;
-        tree = tree->children[0];
-    }
-
-    TSNode *result = malloc(sizeof(TSNode));
-    *result = (TSNode) {
-        .ref_count = 1,
-        .position = position,
-        .content = tree,
-        .parent = NULL,
-        .config = &document->parser->config,
-    };
-    return result;
+    return ts_node_make_root(document->tree, &document->parser->config);
 }
 
 TSNode * ts_document_get_node(const TSDocument *document, size_t pos) {

src/runtime/node.c

@@ -2,19 +2,25 @@
 #include "tree_sitter/parser.h"
 #include "runtime/tree.h"
 
-TSNode * ts_node_make(TSTree *tree, TSNode *parent, size_t position, TSParserConfig *config) {
+TSNode * ts_node_make(const TSTree *tree, TSNode *parent, size_t start_position, TSParserConfig *config) {
     if (parent) ts_node_retain(parent);
     TSNode *result = malloc(sizeof(TSNode));
     *result = (TSNode) {
         .ref_count = 1,
         .parent = parent,
         .content = tree,
-        .position = position,
+        .start_position = start_position,
         .config = config
     };
     return result;
 }
 
+TSNode * ts_node_make_root(const TSTree *tree, TSParserConfig *config) {
+    while (ts_tree_is_wrapper(tree))
+        tree = tree->children[0];
+    return ts_node_make(tree, NULL, 0, config);
+}
+
 void ts_node_retain(TSNode *node) {
     node->ref_count++;
 }
@@ -28,7 +34,7 @@ void ts_node_release(TSNode *node) {
 }
 
 size_t ts_node_pos(const TSNode *node) {
-    return node->position;
+    return node->start_position + node->content->offset;
 }
 
 size_t ts_node_size(const TSNode *node) {
@@ -39,55 +45,17 @@ const char * ts_node_name(const TSNode *node) {
     return node->config->symbol_names[node->content->symbol];
 }
 
-TSTree * ts_tree_child(const TSTree *tree, size_t goal_index, size_t *offset, size_t *children_seen) {
-    *offset = 0;
-    *children_seen = 0;
-    size_t child_count;
-    TSTree **children = ts_tree_children(tree, &child_count);
-
-    if (!children) {
-        *offset = tree->offset + tree->size;
-    }
-
-    for (size_t i = 0; i < child_count; i++) {
-        size_t delta_index = 0, delta_offset = 0;
-        TSTree *child = children[i];
-        TSTree *result = NULL;
-
-        if (ts_tree_is_visible(child)) {
-            delta_offset = child->offset;
-            if (*children_seen == goal_index) {
-                result = child;
-            } else {
-                delta_index = 1;
-                delta_offset += child->size;
-            }
-        } else {
-            result = ts_tree_child(child, (goal_index - *children_seen), &delta_offset, &delta_index);
-        }
-
-        *offset += delta_offset;
-        *children_seen += delta_index;
-
-        if (result) {
-            return result;
-        }
-    }
-
-    return NULL;
-}
-
 const char * ts_node_string(const TSNode *node) {
     return ts_tree_string(node->content, node->config->symbol_names);
 }
 
-TSNode * ts_node_child(TSNode *parent, size_t goal_index) {
-    size_t offset, index;
-    TSTree *child = ts_tree_child(parent->content, goal_index, &offset, &index);
-    if (child)
-        return ts_node_make(child, parent, offset, parent->config);
-    else
+TSNode * ts_node_child(TSNode *parent, size_t index) {
+    size_t child_count;
+    TSChildWithPosition *children = ts_tree_visible_children(parent->content, &child_count);
+    if (child_count <= index)
         return NULL;
+    size_t position = parent->start_position + children[index].position;
+    return ts_node_make(children[index].tree, parent, position, parent->config);
 }
 
 TSNode * ts_node_leaf_at_pos(TSNode *parent, size_t child_index) {

src/runtime/node.h

@@ -5,13 +5,14 @@
 
 struct TSNode {
     size_t ref_count;
-    size_t position;
+    size_t start_position;
     size_t index;
     const TSTree *content;
     struct TSNode *parent;
     TSParserConfig *config;
 };
 
-TSNode * ts_node_make(TSTree *tree, TSNode *parent, size_t position, TSParserConfig *config);
+TSNode * ts_node_make(const TSTree *tree, TSNode *parent, size_t start_position, TSParserConfig *config);
+TSNode * ts_node_make_root(const TSTree *tree, TSParserConfig *config);
 
 #endif

src/runtime/tree.c

@@ -23,7 +23,7 @@ TSTree * ts_tree_make_leaf(TSSymbol symbol, size_t size, size_t offset, int is_h
 }
 
 TSTree * ts_tree_make_node(TSSymbol symbol, size_t child_count, TSTree **children, int is_hidden) {
-    size_t size = 0, offset = 0;
+    size_t size = 0, offset = 0, visible_child_count = 0;
     for (size_t i = 0; i < child_count; i++) {
         TSTree *child = children[i];
         ts_tree_retain(child);
@@ -33,14 +33,55 @@ TSTree * ts_tree_make_node(TSSymbol symbol, size_t child_count, TSTree **childre
         } else {
             size += child->offset + child->size;
         }
+
+        if (ts_tree_is_visible(child))
+            visible_child_count++;
+        else
+            visible_child_count += ts_tree_visible_child_count(child);
     }
 
-    TSTree *result = ts_tree_make(symbol, size, offset, is_hidden);
-    result->child_count = child_count;
-    result->children = children;
-
+    TSTreeOptions options = 0;
+    if (is_hidden)
+        options |= TSTreeOptionsHidden;
     if (child_count == 1 && (ts_tree_is_visible(children[0]) || ts_tree_is_wrapper(children[0])))
-        result->options |= (TSTreeOptionsWrapper | TSTreeOptionsHidden);
+        options |= (TSTreeOptionsWrapper | TSTreeOptionsHidden);
+
+    TSTree *result = malloc(sizeof(TSTree) + (visible_child_count * sizeof(TSChildWithPosition)));
+    *result = (TSTree) {
+        .ref_count = 1,
+        .symbol = symbol,
+        .size = size,
+        .offset = offset,
+        .options = options,
+        .children = children,
+        .child_count = child_count,
+        .visible_child_count = visible_child_count,
+    };
+
+    TSChildWithPosition *visible_children = ts_tree_visible_children(result, NULL);
+
+    for (size_t i = 0, visible_i = 0, child_position = 0; i < child_count; i++) {
+        TSTree *child = children[i];
+        if (ts_tree_is_visible(child)) {
+            visible_children[visible_i] = (TSChildWithPosition) {
+                .tree = child,
+                .position = child_position
+            }; 
+            visible_i++;
+        } else {
+            size_t granchild_count = 0;
+            TSChildWithPosition *grandchildren = ts_tree_visible_children(child, &granchild_count);
+            for (size_t j = 0; j < granchild_count; j++) {
+                visible_children[visible_i] = (TSChildWithPosition) {
+                    .tree = grandchildren[j].tree,
+                    .position = grandchildren[j].position + child_position
+                };
+                visible_i++;
+            }
+        }
+
+        child_position += child->offset + child->size;
+    }
 
     return result;
 }

src/runtime/tree.h

@@ -23,6 +23,7 @@ struct TSTree {
         struct {
             size_t child_count;
             struct TSTree **children;
+            size_t visible_child_count;
         };
         struct {
             size_t expected_input_count;
@@ -32,6 +33,11 @@ struct TSTree {
     };
 };
 
+typedef struct {
+    TSTree *tree;
+    size_t position;
+} TSChildWithPosition;
+
 static inline int ts_tree_is_extra(const TSTree *tree) {
     return (tree->options & TSTreeOptionsExtra);
 }
@@ -48,6 +54,23 @@ static inline int ts_tree_is_wrapper(const TSTree *tree) {
     return (tree->options & TSTreeOptionsWrapper);
 }
 
+static inline size_t ts_tree_visible_child_count(const TSTree *tree) {
+    if (tree->symbol == ts_builtin_sym_error)
+        return 0;
+    else
+        return tree->visible_child_count;
+}
+
+static inline TSChildWithPosition * ts_tree_visible_children(const TSTree *tree, size_t *count) {
+    if (tree->symbol == ts_builtin_sym_error || tree->visible_child_count == 0) {
+        if (count) *count = 0;
+        return NULL;
+    } else {
+        if (count) *count = tree->visible_child_count;
+        return (TSChildWithPosition *)(tree + 1);
+    }
+}
+
 TSTree * ts_tree_make_leaf(TSSymbol symbol, size_t size, size_t offset, int is_hidden);
 TSTree * ts_tree_make_node(TSSymbol symbol, size_t child_count, TSTree **children, int is_hidden);
 TSTree * ts_tree_make_error(char lookahead_char, size_t expected_input_count, const TSSymbol *expected_inputs, size_t size, size_t offset);