Add ts_tree_cursor_goto_first_child_for_byte method

Atom needs this for efficiently seeking to the leaf node at a given position,
visiting all of its ancestors along the way.

include/tree_sitter/runtime.h

@@ -122,6 +122,7 @@ TSNode ts_node_named_descendant_for_point_range(TSNode, TSPoint, TSPoint);
 TSTreeCursor *ts_tree_cursor_new(const TSTree *);
 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);
 bool ts_tree_cursor_goto_next_sibling(TSTreeCursor *);
 bool ts_tree_cursor_goto_parent(TSTreeCursor *);
 TSNode ts_tree_cursor_current_node(TSTreeCursor *);

src/runtime/tree_cursor.c

@@ -62,6 +62,57 @@ bool ts_tree_cursor_goto_first_child(TSTreeCursor *self) {
   return false;
 }
 
+int64_t ts_tree_cursor_goto_first_child_for_byte(TSTreeCursor *self, uint32_t goal_byte) {
+  uint32_t initial_size = self->stack.size;
+  TreeCursorEntry *last_entry = array_back(&self->stack);
+  const Subtree *tree = last_entry->subtree;
+  Length position = last_entry->position;
+  uint32_t visible_child_index = 0;
+
+  bool did_descend;
+  do {
+    did_descend = false;
+
+    uint32_t structural_child_index = 0;
+    for (uint32_t i = 0; i < tree->children.size; i++) {
+      const Subtree *child = tree->children.contents[i];
+      Length next_position = length_add(position, ts_subtree_total_size(child));
+      bool at_goal = next_position.bytes > goal_byte;
+
+      if (at_goal) {
+        if (child->visible || child->visible_child_count > 0) {
+          array_push(&self->stack, ((TreeCursorEntry) {
+            .subtree = child,
+            .child_index = i,
+            .structural_child_index = structural_child_index,
+            .position = position,
+          }));
+
+          if (child->visible) {
+            return visible_child_index;
+          } else {
+            tree = child;
+            did_descend = true;
+            break;
+          }
+        }
+      } else {
+        if (child->visible) {
+          visible_child_index++;
+        } else {
+          visible_child_index += child->visible_child_count;
+        }
+      }
+
+      if (!child->extra) structural_child_index++;
+      position = next_position;
+    }
+  } while (did_descend);
+
+  self->stack.size = initial_size;
+  return -1;
+}
+
 bool ts_tree_cursor_goto_next_sibling(TSTreeCursor *self) {
   TreeCursorEntry *child_entry = array_back(&self->stack);
 

test/runtime/node_test.cc

@@ -649,6 +649,56 @@ describe("TreeCursor", [&]() {
     AssertThat(ts_node_type(node), Equals("value"));
     AssertThat(ts_node_start_byte(node), Equals(array_index));
   });
+
+  it("can find the first child of a given node which spans the given byte offset", [&]() {
+    int64_t child_index = ts_tree_cursor_goto_first_child_for_byte(cursor, 1);
+    TSNode node = ts_tree_cursor_current_node(cursor);
+    AssertThat(ts_node_type(node), Equals("array"));
+    AssertThat(ts_node_start_byte(node), Equals(array_index));
+    AssertThat(child_index, Equals(0));
+
+    child_index = ts_tree_cursor_goto_first_child_for_byte(cursor, array_index);
+    node = ts_tree_cursor_current_node(cursor);
+    AssertThat(ts_node_type(node), Equals("["));
+    AssertThat(ts_node_start_byte(node), Equals(array_index));
+    AssertThat(child_index, Equals(0));
+
+    ts_tree_cursor_goto_parent(cursor);
+    child_index = ts_tree_cursor_goto_first_child_for_byte(cursor, array_index + 1);
+    node = ts_tree_cursor_current_node(cursor);
+    AssertThat(ts_node_type(node), Equals("number"));
+    AssertThat(ts_node_start_byte(node), Equals(number_index));
+    AssertThat(child_index, Equals(1));
+
+    ts_tree_cursor_goto_parent(cursor);
+    child_index = ts_tree_cursor_goto_first_child_for_byte(cursor, number_index + 1);
+    node = ts_tree_cursor_current_node(cursor);
+    AssertThat(ts_node_type(node), Equals("number"));
+    AssertThat(ts_node_start_byte(node), Equals(number_index));
+    AssertThat(child_index, Equals(1));
+
+    ts_tree_cursor_goto_parent(cursor);
+    child_index = ts_tree_cursor_goto_first_child_for_byte(cursor, false_index - 1);
+    node = ts_tree_cursor_current_node(cursor);
+    AssertThat(ts_node_type(node), Equals("false"));
+    AssertThat(ts_node_start_byte(node), Equals(false_index));
+    AssertThat(child_index, Equals(3));
+
+    ts_tree_cursor_goto_parent(cursor);
+    child_index = ts_tree_cursor_goto_first_child_for_byte(cursor, object_end_index - 1);
+    node = ts_tree_cursor_current_node(cursor);
+    AssertThat(ts_node_type(node), Equals("object"));
+    AssertThat(ts_node_start_byte(node), Equals(object_index));
+    AssertThat(child_index, Equals(5));
+
+    // There is no child past the end of the array
+    ts_tree_cursor_goto_parent(cursor);
+    child_index = ts_tree_cursor_goto_first_child_for_byte(cursor, array_end_index);
+    node = ts_tree_cursor_current_node(cursor);
+    AssertThat(ts_node_type(node), Equals("array"));
+    AssertThat(ts_node_start_byte(node), Equals(array_index));
+    AssertThat(child_index, Equals(-1));
+  });
 });
 
 END_TEST