Allow creating a tree cursor starting at any node, not just the root

Co-Authored-By: Ashi Krishnan <queerviolet@github.com>

include/tree_sitter/runtime.h

@@ -131,7 +131,7 @@ TSNode ts_node_named_descendant_for_byte_range(TSNode, uint32_t, uint32_t);
 TSNode ts_node_descendant_for_point_range(TSNode, TSPoint, TSPoint);
 TSNode ts_node_named_descendant_for_point_range(TSNode, TSPoint, TSPoint);
 
-TSTreeCursor ts_tree_cursor_new(const TSTree *);
+TSTreeCursor ts_tree_cursor_new(TSNode);
 void ts_tree_cursor_delete(TSTreeCursor *);
 bool ts_tree_cursor_goto_first_child(TSTreeCursor *);
 int64_t ts_tree_cursor_goto_first_child_for_byte(TSTreeCursor *, uint32_t);

src/runtime/tree.c

@@ -43,8 +43,9 @@ void ts_tree_edit(TSTree *self, const TSInputEdit *edit) {
 TSRange *ts_tree_get_changed_ranges(const TSTree *self, const TSTree *other, uint32_t *count) {
   TSRange *result;
   TreeCursor cursor1, cursor2;
-  ts_tree_cursor_init(&cursor1, self);
-  ts_tree_cursor_init(&cursor2, self);
+  TSNode root = ts_tree_root_node(self);
+  ts_tree_cursor_init(&cursor1, root);
+  ts_tree_cursor_init(&cursor2, root);
   *count = ts_subtree_get_changed_ranges(
     self->root, other->root, &cursor1, &cursor2,
     self->language, &result

src/runtime/tree_cursor.c

@@ -46,26 +46,36 @@ static inline bool ts_tree_cursor_child_iterator_next(ChildIterator *self,
   if (!child->extra && self->alias_sequence) {
     *visible |= self->alias_sequence[self->structural_child_index];
   }
-  self->position = length_add(self->position, ts_subtree_total_size(child));
+
+  self->position = length_add(self->position, child->size);
   self->child_index++;
   if (!child->extra) self->structural_child_index++;
+
+  if (self->child_index < self->parent->children.size) {
+    const Subtree *child = self->parent->children.contents[self->child_index];
+    self->position = length_add(self->position, child->padding);
+  }
+
   return true;
 }
 
 // TSTreeCursor - lifecycle
 
-TSTreeCursor ts_tree_cursor_new(const TSTree *tree) {
+TSTreeCursor ts_tree_cursor_new(TSNode node) {
   TSTreeCursor self;
-  ts_tree_cursor_init((TreeCursor *)&self, tree);
+  ts_tree_cursor_init((TreeCursor *)&self, node);
   return self;
 }
 
-void ts_tree_cursor_init(TreeCursor *self, const TSTree *tree) {
-  self->tree = tree;
+void ts_tree_cursor_init(TreeCursor *self, TSNode node) {
+  self->tree = node.tree;
   array_init(&self->stack);
   array_push(&self->stack, ((TreeCursorEntry) {
-    .subtree = tree->root,
-    .position = length_zero(),
+    .subtree = (const Subtree *)node.id,
+    .position = {
+      ts_node_start_byte(node),
+      ts_node_start_point(node)
+    },
     .child_index = 0,
     .structural_child_index = 0,
   }));
@@ -118,7 +128,8 @@ int64_t ts_tree_cursor_goto_first_child_for_byte(TSTreeCursor *_self, uint32_t g
     TreeCursorEntry entry;
     ChildIterator iterator = ts_tree_cursor_iterate_children(self);
     while (ts_tree_cursor_child_iterator_next(&iterator, &entry, &visible)) {
-      bool at_goal = iterator.position.bytes > goal_byte;
+      uint32_t end_byte = entry.position.bytes + entry.subtree->size.bytes;
+      bool at_goal = end_byte > goal_byte;
       uint32_t visible_child_count = entry.subtree->children.size > 0
         ? entry.subtree->visible_child_count
         : 0;
@@ -208,7 +219,7 @@ TSNode ts_tree_cursor_current_node(const TSTreeCursor *_self) {
   return ts_node_new(
     self->tree,
     last_entry->subtree,
-    length_add(last_entry->position, last_entry->subtree->padding),
+    last_entry->position,
     alias_symbol
   );
 }

src/runtime/tree_cursor.h

@@ -15,6 +15,6 @@ typedef struct {
   const TSTree *tree;
 } TreeCursor;
 
-void ts_tree_cursor_init(TreeCursor *, const TSTree *);
+void ts_tree_cursor_init(TreeCursor *, TSNode);
 
 #endif  // RUNTIME_TREE_CURSOR_H_

test/runtime/node_test.cc

@@ -556,7 +556,7 @@ describe("TreeCursor", [&]() {
     parser = ts_parser_new();
     ts_parser_set_language(parser, load_real_language("json"));
     tree = ts_parser_parse_string(parser, nullptr, json_string.c_str(), json_string.size());
-    cursor = ts_tree_cursor_new(tree);
+    cursor = ts_tree_cursor_new(ts_tree_root_node(tree));
   });
 
   after_each([&]() {
@@ -728,7 +728,7 @@ describe("TreeCursor", [&]() {
     ts_tree_delete(tree);
 
     tree = ts_parser_parse_string(parser, nullptr, "b ... b ... c", 13);
-    cursor = ts_tree_cursor_new(tree);
+    cursor = ts_tree_cursor_new(ts_tree_root_node(tree));
 
     TSNode node = ts_tree_cursor_current_node(&cursor);
     AssertThat(ts_node_type(node), Equals("a"));