Cache parent nodes to speed up ts_node_parent in common cases

src/runtime/tree.c

@@ -5,10 +5,15 @@
 #include "runtime/tree_cursor.h"
 #include "runtime/tree.h"
 
+static const unsigned PARENT_CACHE_CAPACITY = 32;
+
 TSTree *ts_tree_new(const Subtree *root, const TSLanguage *language) {
   TSTree *result = ts_malloc(sizeof(TSTree));
   result->root = root;
   result->language = language;
+  result->parent_cache = NULL;
+  result->parent_cache_start = 0;
+  result->parent_cache_size = 0;
   return result;
 }
 
@@ -21,6 +26,7 @@ void ts_tree_delete(TSTree *self) {
   SubtreePool pool = ts_subtree_pool_new(0);
   ts_subtree_release(&pool, self->root);
   ts_subtree_pool_delete(&pool);
+  if (self->parent_cache) ts_free(self->parent_cache);
   ts_free(self);
 }
 
@@ -51,3 +57,38 @@ TSRange *ts_tree_get_changed_ranges(const TSTree *self, const TSTree *other, uin
 void ts_tree_print_dot_graph(const TSTree *self, FILE *file) {
   ts_subtree_print_dot_graph(self->root, self->language, file);
 }
+
+TSNode ts_tree_get_cached_parent(const TSTree *self, const TSNode *node) {
+  for (uint32_t i = 0; i < self->parent_cache_size; i++) {
+    uint32_t index = (self->parent_cache_start + i) % PARENT_CACHE_CAPACITY;
+    ParentCacheEntry *entry = &self->parent_cache[index];
+    if (entry->child == node->id) {
+      return ts_node_new(self, entry->parent, entry->position, entry->alias_symbol);
+    }
+  }
+  return ts_node_new(NULL, NULL, length_zero(), 0);
+}
+
+void ts_tree_set_cached_parent(const TSTree *_self, const TSNode *node, const TSNode *parent) {
+  TSTree *self = (TSTree *)_self;
+  if (!self->parent_cache) {
+    self->parent_cache = ts_calloc(PARENT_CACHE_CAPACITY, sizeof(ParentCacheEntry));
+  }
+
+  uint32_t index = (self->parent_cache_start + self->parent_cache_size) % PARENT_CACHE_CAPACITY;
+  self->parent_cache[index] = (ParentCacheEntry) {
+    .child = node->id,
+    .parent = parent->id,
+    .position = {
+      parent->context[0],
+      {parent->context[1], parent->context[2]}
+    },
+    .alias_symbol = parent->context[3],
+  };
+
+  if (self->parent_cache_size == PARENT_CACHE_CAPACITY) {
+    self->parent_cache_start++;
+  } else {
+    self->parent_cache_size++;
+  }
+}