Rework AST access API: reduce heap allocation

src/runtime/document.c

@@ -1,13 +1,9 @@
 #include "tree_sitter/parser.h"
 #include "runtime/node.h"
+#include "runtime/length.h"
 #include "runtime/parser.h"
 #include "runtime/string_input.h"
-
-struct TSDocument {
-  TSParser parser;
-  TSInput input;
-  TSNode *node;
-};
+#include "runtime/document.h"
 
 TSDocument *ts_document_make() {
   TSDocument *document = calloc(sizeof(TSDocument), 1);
@@ -19,19 +15,18 @@ void ts_document_free(TSDocument *document) {
   ts_parser_destroy(&document->parser);
   if (document->input.release_fn)
     document->input.release_fn(document->input.data);
-  if (document->node)
-    ts_node_release(document->node);
+  if (document->tree)
+    ts_tree_release(document->tree);
   free(document);
 }
 
 static void reparse(TSDocument *document, TSInputEdit *edit) {
   if (document->input.read_fn && document->parser.language) {
-    const TSTree *tree =
-      ts_parser_parse(&document->parser, document->input, edit);
-    if (document->node)
-      ts_node_release(document->node);
-    document->node =
-      ts_node_make_root(tree, document->parser.language->symbol_names);
+    TSTree *tree = ts_parser_parse(&document->parser, document->input, edit);
+    if (document->tree)
+      ts_tree_release(document->tree);
+    document->tree = tree;
+    ts_tree_retain(tree);
   }
 }
 
@@ -66,6 +61,9 @@ void ts_document_set_input_string(TSDocument *document, const char *text) {
   ts_document_set_input(document, ts_string_input_make(text));
 }
 
-TSNode *ts_document_root_node(const TSDocument *document) {
-  return document->node;
+TSNode ts_document_root_node(const TSDocument *document) {
+  if (document->tree)
+    return ts_node_make(document->tree, document->tree->padding);
+  else
+    return ts_node_null();
 }

src/runtime/document.h

@@ -0,0 +1,14 @@
+#ifndef RUNTIME_DOCUMENT_H_
+#define RUNTIME_DOCUMENT_H_
+
+#include "tree_sitter/parser.h"
+#include "runtime/parser.h"
+#include "runtime/tree.h"
+
+struct TSDocument {
+  TSParser parser;
+  TSInput input;
+  TSTree *tree;
+};
+
+#endif

src/runtime/node.c

@@ -2,105 +2,129 @@
 #include "runtime/node.h"
 #include "runtime/length.h"
 #include "runtime/tree.h"
+#include "runtime/document.h"
 
-TSNode *ts_node_make(const TSTree *tree, TSNode *parent, size_t index,
-                     TSLength position, const char **names) {
-  if (parent)
-    ts_node_retain(parent);
-  TSNode *result = malloc(sizeof(TSNode));
-  *result = (TSNode){
-    .ref_count = 1,
-    .parent = parent,
-    .index = index,
-    .content = tree,
-    .position = position,
-    .names = names,
+TSNode ts_node_make(const TSTree *tree, TSLength position) {
+  return (TSNode){.data = tree, .position = position };
+}
+
+static inline const TSTree *get_tree(TSNode tree) {
+  return tree.data;
+}
+
+TSLength ts_node_pos(TSNode tree) {
+  return tree.position;
+}
+
+TSLength ts_node_size(TSNode tree) {
+  return get_tree(tree)->size;
+}
+
+bool ts_node_eq(TSNode left, TSNode right) {
+  return ts_tree_eq(get_tree(left), get_tree(right)) &&
+         ts_length_eq(left.position, right.position);
+}
+
+const char *ts_node_name(TSNode tree, const TSDocument *document) {
+  return document->parser.language->symbol_names[get_tree(tree)->symbol];
+}
+
+const char *ts_node_string(TSNode tree, const TSDocument *document) {
+  return ts_tree_string(get_tree(tree), document->parser.language->symbol_names);
+}
+
+typedef struct {
+  TSNode ref;
+  size_t index;
+} ParentWithIndex;
+
+static inline ParentWithIndex ts_node_parent_with_index(TSNode ref) {
+  const TSTree *tree = get_tree(ref);
+  size_t index = 0;
+  TSLength position = ts_length_sub(ref.position, tree->padding);
+
+  do {
+    TSTree *parent = tree->parent;
+    if (!parent)
+      return (ParentWithIndex){ ts_node_null(), 0 };
+
+    for (size_t i = 0; i < parent->child_count; i++) {
+      TSTree *child = parent->children[i];
+      if (child == tree)
+        break;
+      if (ts_tree_is_visible(child)) {
+        index += 1;
+      } else {
+        size_t child_count;
+        ts_tree_visible_children(child, &child_count);
+        index += child_count;
+      }
+      position = ts_length_sub(position, ts_tree_total_size(child));
+    }
+    tree = parent;
+  } while (!ts_tree_is_visible(tree));
+
+  return (ParentWithIndex){
+    ts_node_make(tree, ts_length_add(position, tree->padding)), index
   };
-  return result;
 }
 
-TSNode *ts_node_make_root(const TSTree *tree, const char **names) {
-  return ts_node_make(tree, NULL, 0, tree->padding, names);
+TSNode ts_node_parent(TSNode ref) {
+  return ts_node_parent_with_index(ref).ref;
 }
 
-void ts_node_retain(TSNode *node) { node->ref_count++; }
-
-void ts_node_release(TSNode *node) {
-  node->ref_count--;
-  if (node->ref_count == 0) {
-    if (node->parent)
-      ts_node_release(node->parent);
-    free(node);
-  }
-}
-
-TSLength ts_node_pos(const TSNode *node) { return node->position; }
-
-TSLength ts_node_size(const TSNode *node) { return node->content->size; }
-
-bool ts_node_eq(const TSNode *left, const TSNode *right) {
-  return ts_tree_eq(left->content, right->content);
-}
-
-const char *ts_node_name(const TSNode *node) {
-  return node->names[node->content->symbol];
-}
-
-const char *ts_node_string(const TSNode *node) {
-  return ts_tree_string(node->content, node->names);
-}
-
-TSNode *ts_node_parent(TSNode *child) { return child->parent; }
-
-TSNode *ts_node_prev_sibling(TSNode *child) {
-  if (child->parent)
-    return ts_node_child(child->parent, child->index - 1);
+TSNode ts_node_prev_sibling(TSNode ref) {
+  ParentWithIndex parent = ts_node_parent_with_index(ref);
+  if (parent.ref.data && parent.index > 0)
+    return ts_node_child(parent.ref, parent.index - 1);
   else
-    return NULL;
+    return ts_node_null();
 }
 
-TSNode *ts_node_next_sibling(TSNode *child) {
-  if (child->parent)
-    return ts_node_child(child->parent, child->index + 1);
+TSNode ts_node_next_sibling(TSNode ref) {
+  ParentWithIndex parent = ts_node_parent_with_index(ref);
+  if (parent.ref.data)
+    return ts_node_child(parent.ref, parent.index + 1);
   else
-    return NULL;
+    return ts_node_null();
 }
 
-size_t ts_node_child_count(const TSNode *parent) {
+size_t ts_node_child_count(TSNode parent) {
   size_t result;
-  ts_tree_visible_children(parent->content, &result);
+  ts_tree_visible_children(get_tree(parent), &result);
   return result;
 }
 
-TSNode *ts_node_child(TSNode *parent, size_t i) {
+TSNode ts_node_child(TSNode ref, size_t i) {
   size_t count;
-  TSTreeChild *children = ts_tree_visible_children(parent->content, &count);
+  TSTreeChild *children = ts_tree_visible_children(get_tree(ref), &count);
   if (i >= count)
-    return NULL;
-  TSLength pos = ts_length_add(parent->position, children[i].offset);
-  return ts_node_make(children[i].tree, parent, i, pos, parent->names);
+    return ts_node_null();
+  TSLength position = ts_length_add(ref.position, children[i].offset);
+  return ts_node_make(children[i].tree, position);
 }
 
-TSNode *ts_node_find_for_range(TSNode *parent, size_t min, size_t max) {
-  size_t count;
-  TSTreeChild *children = ts_tree_visible_children(parent->content, &count);
-  for (size_t i = 0; i < count; i++) {
-    TSTreeChild child = children[i];
-    TSLength pos = ts_length_add(parent->position, child.offset);
-    if (pos.chars > min)
-      break;
-    if (pos.chars + child.tree->size.chars > max) {
-      TSNode *node = ts_node_make(child.tree, parent, i, pos, parent->names);
-      TSNode *result = ts_node_find_for_range(node, min, max);
-      ts_node_release(node);
-      return result;
+TSNode ts_node_find_for_range(TSNode ref, size_t min, size_t max) {
+  bool did_descend = true;
+  while (did_descend) {
+    did_descend = false;
+    size_t count;
+    TSTreeChild *children = ts_tree_visible_children(get_tree(ref), &count);
+
+    for (size_t i = 0; i < count; i++) {
+      TSTreeChild child = children[i];
+      TSLength position = ts_length_add(ref.position, child.offset);
+      if (position.chars > min)
+        break;
+      if (position.chars + child.tree->size.chars > max) {
+        ref = ts_node_make(child.tree, position);
+        did_descend = true;
+      }
     }
   }
-
-  ts_node_retain(parent);
-  return parent;
+  return ref;
 }
 
-TSNode *ts_node_find_for_pos(TSNode *parent, size_t position) {
-  return ts_node_find_for_range(parent, position, position);
+TSNode ts_node_find_for_pos(TSNode ref, size_t position) {
+  return ts_node_find_for_range(ref, position, position);
 }

src/runtime/node.h

@@ -2,19 +2,13 @@
 #define RUNTIME_NODE_H_
 
 #include "tree_sitter/parser.h"
+#include "runtime/length.h"
 #include "runtime/tree.h"
 
-struct TSNode {
-  size_t ref_count;
-  size_t index;
-  TSLength position;
-  const TSTree *content;
-  struct TSNode *parent;
-  const char **names;
-};
+TSNode ts_node_make(const TSTree *, TSLength);
 
-TSNode *ts_node_make(const TSTree *tree, TSNode *parent, size_t index,
-                     TSLength position, const char **names);
-TSNode *ts_node_make_root(const TSTree *tree, const char **names);
+static inline TSNode ts_node_null() {
+  return (TSNode){.data = NULL, .position = ts_length_zero() };
+}
 
 #endif

src/runtime/parser.c

@@ -431,8 +431,7 @@ ParserNextResult ts_parser_next(TSParser *parser, int head_to_advance) {
   return result;
 }
 
-const TSTree *ts_parser_parse(TSParser *parser, TSInput input,
-                              TSInputEdit *edit) {
+TSTree *ts_parser_parse(TSParser *parser, TSInput input, TSInputEdit *edit) {
   TSLength position;
   if (edit) {
     DEBUG("edit pos:%lu, inserted:%lu, deleted:%lu", edit->position,

src/runtime/parser.h

@@ -21,7 +21,7 @@ TSParser ts_parser_make();
 void ts_parser_destroy(TSParser *);
 TSDebugger ts_parser_get_debugger(const TSParser *);
 void ts_parser_set_debugger(TSParser *, TSDebugger);
-const TSTree *ts_parser_parse(TSParser *, TSInput, TSInputEdit *);
+TSTree *ts_parser_parse(TSParser *, TSInput, TSInputEdit *);
 
 #ifdef __cplusplus
 }

src/runtime/tree.c

@@ -64,7 +64,7 @@ TSTree *ts_tree_make_node(TSSymbol symbol, size_t child_count,
   } else {
     if (is_hidden)
       options |= TSTreeOptionsHidden;
-    if (child_count == 1 &&
+    if (child_count == 1 && symbol != ts_builtin_sym_document &&
         (ts_tree_is_visible(children[0]) || ts_tree_is_wrapper(children[0])))
       options |= (TSTreeOptionsWrapper | TSTreeOptionsHidden);
     if (child_count > 0) {
@@ -82,6 +82,7 @@ TSTree *ts_tree_make_node(TSSymbol symbol, size_t child_count,
   TSTree *result =
     malloc(sizeof(TSTree) + (visible_child_count * sizeof(TSTreeChild)));
   *result = (TSTree){.ref_count = 1,
+                     .parent = NULL,
                      .symbol = symbol,
                      .children = children,
                      .child_count = child_count,
@@ -98,6 +99,7 @@ TSTree *ts_tree_make_node(TSSymbol symbol, size_t child_count,
   TSLength offset = ts_length_zero();
   for (size_t i = 0, vis_i = 0; i < child_count; i++) {
     TSTree *child = children[i];
+    child->parent = result;
 
     if (i > 0)
       offset = ts_length_add(offset, child->padding);

src/runtime/tree.h

@@ -17,10 +17,7 @@ typedef enum {
 } TSTreeOptions;
 
 struct TSTree {
-  TSSymbol symbol;
-  TSTreeOptions options;
-  TSLength padding;
-  TSLength size;
+  struct TSTree *parent;
   size_t child_count;
   union {
     struct {
@@ -29,6 +26,10 @@ struct TSTree {
     };
     char lookahead_char;
   };
+  TSLength padding;
+  TSLength size;
+  TSSymbol symbol;
+  TSTreeOptions options;
   unsigned short int ref_count;
 };