From 49f19092fb985c955913b2e6ea7565c4a45bbc98 Mon Sep 17 00:00:00 2001
From: Asger F <asgerf@github.com>
Date: Thu, 7 May 2026 13:55:36 +0200
Subject: [PATCH] Yeast: add reachable_node_ids()

---
 shared/yeast/src/lib.rs    | 33 +++++++++++++++++++++++++++++++++
 shared/yeast/tests/test.rs | 22 ++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/shared/yeast/src/lib.rs b/shared/yeast/src/lib.rs
index 281f44a98b2..c06029a8626 100644
--- a/shared/yeast/src/lib.rs
+++ b/shared/yeast/src/lib.rs
@@ -193,10 +193,43 @@ impl Ast {
         AstCursor::new(self)
     }
 
+    /// Return all nodes currently allocated in the AST arena.
+    ///
+    /// This includes nodes that are no longer reachable from `get_root()`
+    /// after desugaring rewrites. Use `reachable_node_ids()` for output-level
+    /// validation/traversal semantics.
     pub fn nodes(&self) -> &[Node] {
         &self.nodes
     }
 
+    /// Return node ids reachable from `get_root()` by following child edges.
+    ///
+    /// This reflects the effective AST after desugaring and excludes orphaned
+    /// arena nodes left behind by rewrite operations.
+    pub fn reachable_node_ids(&self) -> Vec<usize> {
+        let mut reachable = Vec::new();
+        let mut stack = vec![self.root];
+        let mut seen = vec![false; self.nodes.len()];
+
+        while let Some(id) = stack.pop() {
+            if id >= self.nodes.len() || seen[id] {
+                continue;
+            }
+            seen[id] = true;
+            reachable.push(id);
+
+            if let Some(node) = self.get_node(id) {
+                for children in node.fields.values() {
+                    for &child in children {
+                        stack.push(child);
+                    }
+                }
+            }
+        }
+
+        reachable
+    }
+
     pub fn get_root(&self) -> Id {
         self.root
     }
diff --git a/shared/yeast/tests/test.rs b/shared/yeast/tests/test.rs
index ed4202493a4..e058e6b1eb0 100644
--- a/shared/yeast/tests/test.rs
+++ b/shared/yeast/tests/test.rs
@@ -166,6 +166,28 @@ fn test_query_no_match() {
     assert!(!matched);
 }
 
+#[test]
+fn test_reachable_nodes_excludes_orphaned_rewrite_nodes() {
+    let lang: tree_sitter::Language = tree_sitter_ruby::LANGUAGE.into();
+    let schema = yeast::node_types_yaml::schema_from_yaml_with_language(OUTPUT_SCHEMA_YAML, &lang)
+        .unwrap();
+    let rules = vec![yeast::rule!((integer) => (identifier "replaced"))];
+    let runner = Runner::with_schema(lang, &schema, &rules);
+
+    let input = "x = 1";
+    let ast = runner.run(input).unwrap();
+    let reachable_ids = ast.reachable_node_ids();
+
+    assert!(
+        ast.nodes().len() > reachable_ids.len(),
+        "expected rewrite to leave orphaned arena nodes"
+    );
+
+    let dump = dump_ast(&ast, ast.get_root(), input);
+    assert!(dump.contains("identifier \"replaced\""));
+    assert!(!dump.contains("integer \"1\""));
+}
+
 #[test]
 fn test_query_repeated_capture() {
     let runner = Runner::new(tree_sitter_ruby::LANGUAGE.into(), &[]);