Start work on exposing APIs for node descendant counts and indices

lib/src/node.c

@@ -458,6 +458,10 @@ bool ts_node_has_error(TSNode self) {
   return ts_subtree_error_cost(ts_node__subtree(self)) > 0;
 }
 
+uint32_t ts_node_descendant_count(TSNode self) {
+  return ts_subtree_visible_descendant_count(ts_node__subtree(self)) + 1;
+}
+
 TSNode ts_node_parent(TSNode self) {
   TSNode node = ts_tree_root_node(self.tree);
   uint32_t end_byte = ts_node_end_byte(self);

lib/src/subtree.c

@@ -376,7 +376,7 @@ void ts_subtree_summarize_children(
   self.ptr->visible_child_count = 0;
   self.ptr->error_cost = 0;
   self.ptr->repeat_depth = 0;
-  self.ptr->node_count = 1;
+  self.ptr->visible_descendant_count = 0;
   self.ptr->has_external_tokens = false;
   self.ptr->depends_on_column = false;
   self.ptr->has_external_scanner_state_change = false;
@@ -435,14 +435,16 @@ void ts_subtree_summarize_children(
     }
 
     self.ptr->dynamic_precedence += ts_subtree_dynamic_precedence(child);
-    self.ptr->node_count += ts_subtree_node_count(child);
+    self.ptr->visible_descendant_count += ts_subtree_visible_descendant_count(child);
 
     if (alias_sequence && alias_sequence[structural_index] != 0 && !ts_subtree_extra(child)) {
+      self.ptr->visible_descendant_count++;
       self.ptr->visible_child_count++;
       if (ts_language_symbol_metadata(language, alias_sequence[structural_index]).named) {
         self.ptr->named_child_count++;
       }
     } else if (ts_subtree_visible(child)) {
+      self.ptr->visible_descendant_count++;
       self.ptr->visible_child_count++;
       if (ts_subtree_named(child)) self.ptr->named_child_count++;
     } else if (grandchild_count > 0) {
@@ -529,7 +531,7 @@ MutableSubtree ts_subtree_new_node(
     .fragile_right = fragile,
     .is_keyword = false,
     {{
-      .node_count = 0,
+      .visible_descendant_count = 0,
       .production_id = production_id,
       .first_leaf = {.symbol = 0, .parse_state = 0},
     }}

lib/src/subtree.h

@@ -135,7 +135,7 @@ typedef struct {
     struct {
       uint32_t visible_child_count;
       uint32_t named_child_count;
-      uint32_t node_count;
+      uint32_t visible_descendant_count;
       int32_t dynamic_precedence;
       uint16_t repeat_depth;
       uint16_t production_id;
@@ -297,8 +297,16 @@ static inline uint32_t ts_subtree_is_repetition(Subtree self) {
     : !self.ptr->named && !self.ptr->visible && self.ptr->child_count != 0;
 }
 
+static inline uint32_t ts_subtree_visible_descendant_count(Subtree self) {
+  return (self.data.is_inline || self.ptr->child_count == 0)
+    ? 0
+    : self.ptr->visible_descendant_count;
+}
+
 static inline uint32_t ts_subtree_node_count(Subtree self) {
-  return (self.data.is_inline || self.ptr->child_count == 0) ? 1 : self.ptr->node_count;
+  return
+    ts_subtree_visible_descendant_count(self) +
+    (ts_subtree_visible(self) ? 1 : 0);
 }
 
 static inline uint32_t ts_subtree_visible_child_count(Subtree self) {

lib/src/tree_cursor.c

@@ -10,26 +10,50 @@ typedef struct {
   Length position;
   uint32_t child_index;
   uint32_t structural_child_index;
+  uint32_t descendant_index;
   const TSSymbol *alias_sequence;
 } CursorChildIterator;
 
 // CursorChildIterator
 
+static inline bool ts_tree_cursor_is_entry_visible(const TreeCursor *self, uint32_t i) {
+  TreeCursorEntry *entry = &self->stack.contents[i];
+  if (ts_subtree_visible(*entry->subtree)) {
+    return true;
+  } else if (i > 0 && !ts_subtree_extra(*entry->subtree)) {
+    TreeCursorEntry *parent_entry = &self->stack.contents[i - 1];
+    return ts_language_alias_at(
+      self->tree->language,
+      parent_entry->subtree->ptr->production_id,
+      entry->structural_child_index
+    );
+  } else {
+    return false;
+  }
+}
+
 static inline CursorChildIterator ts_tree_cursor_iterate_children(const TreeCursor *self) {
   TreeCursorEntry *last_entry = array_back(&self->stack);
   if (ts_subtree_child_count(*last_entry->subtree) == 0) {
-    return (CursorChildIterator) {NULL_SUBTREE, self->tree, length_zero(), 0, 0, NULL};
+    return (CursorChildIterator) {NULL_SUBTREE, self->tree, length_zero(), 0, 0, 0, NULL};
   }
   const TSSymbol *alias_sequence = ts_language_alias_sequence(
     self->tree->language,
     last_entry->subtree->ptr->production_id
   );
+
+  uint32_t descendant_index = last_entry->descendant_index;
+  if (ts_tree_cursor_is_entry_visible(self, self->stack.size - 1)) {
+    descendant_index += 1;
+  }
+
   return (CursorChildIterator) {
     .tree = self->tree,
     .parent = *last_entry->subtree,
     .position = last_entry->position,
     .child_index = 0,
     .structural_child_index = 0,
+    .descendant_index = descendant_index,
     .alias_sequence = alias_sequence,
   };
 }
@@ -46,14 +70,22 @@ static inline bool ts_tree_cursor_child_iterator_next(
     .position = self->position,
     .child_index = self->child_index,
     .structural_child_index = self->structural_child_index,
+    .descendant_index = self->descendant_index,
   };
   *visible = ts_subtree_visible(*child);
   bool extra = ts_subtree_extra(*child);
-  if (!extra && self->alias_sequence) {
-    *visible |= self->alias_sequence[self->structural_child_index];
+  if (!extra) {
+    if (self->alias_sequence) {
+      *visible |= self->alias_sequence[self->structural_child_index];
+    }
     self->structural_child_index++;
   }
 
+  self->descendant_index += ts_subtree_visible_descendant_count(*child);
+  if (*visible) {
+    self->descendant_index += 1;
+  }
+
   self->position = length_add(self->position, ts_subtree_size(*child));
   self->child_index++;
 
@@ -88,6 +120,7 @@ void ts_tree_cursor_init(TreeCursor *self, TSNode node) {
     },
     .child_index = 0,
     .structural_child_index = 0,
+    .descendant_index = 0,
   }));
 }
 
@@ -190,6 +223,7 @@ TreeCursorStep ts_tree_cursor_goto_next_sibling_internal(TSTreeCursor *_self) {
     iterator.child_index = entry.child_index;
     iterator.structural_child_index = entry.structural_child_index;
     iterator.position = entry.position;
+    iterator.descendant_index = entry.descendant_index;
 
     bool visible = false;
     ts_tree_cursor_child_iterator_next(&iterator, &entry, &visible);
@@ -247,6 +281,62 @@ bool ts_tree_cursor_goto_parent(TSTreeCursor *_self) {
   return false;
 }
 
+void ts_tree_cursor_goto_descendant(
+  TSTreeCursor *_self,
+  uint32_t goal_descendant_index
+) {
+  TreeCursor *self = (TreeCursor *)_self;
+
+  // Ascend to the lowest ancestor that contains the goal descendant.
+  for (;;) {
+    TreeCursorEntry *entry = &self->stack.contents[self->stack.size - 1];
+    if (
+      (entry->descendant_index <= goal_descendant_index) &&
+      (entry->descendant_index + ts_subtree_visible_descendant_count(*entry->subtree) > goal_descendant_index)
+    ) {
+      break;
+    } else if (self->stack.size <= 1) {
+      return;
+    } else {
+      array_pop(&self->stack);
+    }
+  }
+
+  // Descend to the goal descendant.
+  bool did_descend = true;
+  do {
+    did_descend = false;
+    bool visible;
+    TreeCursorEntry entry;
+    CursorChildIterator iterator = ts_tree_cursor_iterate_children(self);
+
+    // If the goal descendant is the current node, then we're done.
+    if (iterator.descendant_index > goal_descendant_index) {
+      return;
+    }
+
+    while (ts_tree_cursor_child_iterator_next(&iterator, &entry, &visible)) {
+      uint32_t next_descendant_index = entry.descendant_index + ts_subtree_visible_descendant_count(*entry.subtree);
+      if (visible) next_descendant_index += 1;
+      if (next_descendant_index > goal_descendant_index) {
+        array_push(&self->stack, entry);
+        if (visible && entry.descendant_index == goal_descendant_index) {
+          return;
+        } else {
+          did_descend = true;
+          break;
+        }
+      }
+    }
+  } while (did_descend);
+}
+
+uint32_t ts_tree_cursor_descendant_index(TSTreeCursor *_self) {
+  const TreeCursor *self = (const TreeCursor *)_self;
+  TreeCursorEntry *last_entry = array_back(&self->stack);
+  return last_entry->descendant_index;
+}
+
 TSNode ts_tree_cursor_current_node(const TSTreeCursor *_self) {
   const TreeCursor *self = (const TreeCursor *)_self;
   TreeCursorEntry *last_entry = array_back(&self->stack);

lib/src/tree_cursor.h

@@ -8,6 +8,7 @@ typedef struct {
   Length position;
   uint32_t child_index;
   uint32_t structural_child_index;
+  uint32_t descendant_index;
 } TreeCursorEntry;
 
 typedef struct {