YEAST: add NodeRef type, YeastDisplay trait, and source text storage

Introduce NodeRef as a typed wrapper around node arena IDs. Captures in
desugaring rules are now bound as NodeRef instead of raw usize, which
prevents accidental misuse and enables source-text-aware rendering.

Add the YeastDisplay trait as an alternative to Display: its
yeast_to_string method receives the Ast, allowing NodeRef to resolve to
the captured node's source text instead of printing a numeric ID.

Store the original source bytes in the Ast so that NodeContent::Range
values (from synthesized literal nodes) can be resolved back to text.

Update yeast-macros to emit NodeRef-typed capture bindings and use
Into::<usize>::into where raw IDs are needed. The #{expr} template
syntax now uses YeastDisplay instead of Display.

The effect is visible in the corpus tests: operator nodes now correctly
render as e.g. operator "+" instead of operator "3".

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: asgerf

shared/tree-sitter-extractor/src/extractor/mod.rs

@@ -326,7 +326,7 @@ pub fn extract(
 
     if let Some(yeast_runner) = yeast_runner {
         let ast = yeast_runner
-            .run_from_tree(&tree)
+            .run_from_tree(&tree, source)
             .unwrap_or_else(|e| panic!("Desugaring failed for {path_str}: {e}"));
         traverse_yeast(&ast, &mut visitor);
     } else {

shared/yeast-macros/src/parse.rs

@@ -299,7 +299,7 @@ fn parse_direct_node(tokens: &mut Tokens, ctx: &Ident) -> Result<TokenStream> {
         Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Brace => {
             let group = expect_group(tokens, Delimiter::Brace)?;
             let expr = group.stream();
-            Ok(quote! { #expr })
+            Ok(quote! { ::std::convert::Into::<usize>::into(#expr) })
         }
         Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Parenthesis => {
             let group = expect_group(tokens, Delimiter::Parenthesis)?;
@@ -329,12 +329,17 @@ fn parse_direct_node_inner(tokens: &mut Tokens, ctx: &Ident) -> Result<TokenStre
         return Ok(quote! { #ctx.literal(#kind_str, #lit) });
     }
 
-    // Check for (kind #{expr}) — computed literal, expr converted via .to_string()
+    // Check for (kind #{expr}) — computed literal, expr converted via YeastDisplay
     if peek_is_hash(tokens) {
         tokens.next(); // consume #
         let group = expect_group(tokens, Delimiter::Brace)?;
         let expr = group.stream();
-        return Ok(quote! { #ctx.literal(#kind_str, &(#expr).to_string()) });
+        return Ok(quote! {
+            {
+                let __value = yeast::YeastDisplay::yeast_to_string(&(#expr), &*#ctx.ast);
+                #ctx.literal(#kind_str, &__value)
+            }
+        });
     }
 
     // Check for (kind $fresh)
@@ -374,15 +379,19 @@ fn parse_direct_node_inner(tokens: &mut Tokens, ctx: &Ident) -> Result<TokenStre
                     inner.next(); // consume first .
                     inner.next(); // consume second .
                     let expr: proc_macro2::TokenStream = inner.collect();
-                    stmts.push(quote! { let #temp: Vec<usize> = #expr; });
+                    stmts.push(quote! {
+                        let #temp: Vec<usize> = (#expr).into_iter()
+                            .map(::std::convert::Into::<usize>::into)
+                            .collect();
+                    });
                     field_args.push(quote! { (#field_str, #temp) });
                     continue;
                 }
             }
         }
 
         let value = parse_direct_node(tokens, ctx)?;
-        stmts.push(quote! { let #temp = #value; });
+        stmts.push(quote! { let #temp: usize = #value; });
         field_args.push(quote! { (#field_str, vec![#temp]) });
     }
 
@@ -427,10 +436,16 @@ fn parse_direct_list(tokens: &mut Tokens, ctx: &Ident) -> Result<Vec<TokenStream
                 inner.next(); // consume first .
                 inner.next(); // consume second .
                 let expr: TokenStream = inner.collect();
-                items.push(quote! { __nodes.extend(#expr); });
+                items.push(quote! {
+                    __nodes.extend(
+                        (#expr).into_iter().map(::std::convert::Into::<usize>::into)
+                    );
+                });
             } else {
                 let expr = group.stream();
-                items.push(quote! { __nodes.push(#expr); });
+                items.push(quote! {
+                    __nodes.push(::std::convert::Into::<usize>::into(#expr));
+                });
             }
             continue;
         }
@@ -580,13 +595,24 @@ pub fn parse_rule_top(input: TokenStream) -> Result<TokenStream> {
             let name_str = &cap.name;
             match cap.multiplicity {
                 CaptureMultiplicity::Repeated => {
-                    quote! { let #name: Vec<usize> = __captures.get_all(#name_str); }
+                    quote! {
+                        let #name: Vec<yeast::NodeRef> = __captures.get_all(#name_str)
+                            .into_iter()
+                            .map(yeast::NodeRef)
+                            .collect();
+                    }
                 }
                 CaptureMultiplicity::Optional => {
-                    quote! { let #name: Option<usize> = __captures.get_opt(#name_str); }
+                    quote! {
+                        let #name: Option<yeast::NodeRef> =
+                            __captures.get_opt(#name_str).map(yeast::NodeRef);
+                    }
                 }
                 CaptureMultiplicity::Single => {
-                    quote! { let #name: usize = __captures.get_var(#name_str).unwrap(); }
+                    quote! {
+                        let #name: yeast::NodeRef =
+                            yeast::NodeRef(__captures.get_var(#name_str).unwrap());
+                    }
                 }
             }
         })
@@ -613,19 +639,26 @@ pub fn parse_rule_top(input: TokenStream) -> Result<TokenStream> {
                     CaptureMultiplicity::Repeated => quote! {
                         let __field_id = #ctx_ident.ast.field_id_for_name(#name_str)
                             .unwrap_or_else(|| panic!("field '{}' not found", #name_str));
-                        __fields.insert(__field_id, #name);
+                        __fields.insert(
+                            __field_id,
+                            #name.into_iter()
+                                .map(::std::convert::Into::<usize>::into)
+                                .collect(),
+                        );
                     },
                     CaptureMultiplicity::Optional => quote! {
                         let __field_id = #ctx_ident.ast.field_id_for_name(#name_str)
                             .unwrap_or_else(|| panic!("field '{}' not found", #name_str));
                         if let Some(__id) = #name {
-                            __fields.entry(__field_id).or_insert_with(Vec::new).push(__id);
+                            __fields.entry(__field_id).or_insert_with(Vec::new)
+                                .push(::std::convert::Into::<usize>::into(__id));
                         }
                     },
                     CaptureMultiplicity::Single => quote! {
                         let __field_id = #ctx_ident.ast.field_id_for_name(#name_str)
                             .unwrap_or_else(|| panic!("field '{}' not found", #name_str));
-                        __fields.entry(__field_id).or_insert_with(Vec::new).push(#name);
+                        __fields.entry(__field_id).or_insert_with(Vec::new)
+                            .push(::std::convert::Into::<usize>::into(#name));
                     },
                 }
             })

shared/yeast/src/lib.rs

@@ -23,12 +23,73 @@ pub use cursor::Cursor;
 use query::QueryNode;
 
 /// Node ids are indexes into the arena
-type Id = usize;
+pub type Id = usize;
 
 /// Field and Kind ids are provided by tree-sitter
 type FieldId = u16;
 type KindId = u16;
 
+/// A typed reference to a node in an [`Ast`] arena. Wraps an [`Id`] but
+/// deliberately does not implement [`std::fmt::Display`]: rendering a node
+/// requires the [`Ast`] it lives in (to resolve [`NodeContent::Range`] back
+/// to source text). Use [`YeastDisplay::yeast_to_string`] to format it.
+#[derive(Copy, Clone, Eq, PartialEq, Debug, Hash)]
+pub struct NodeRef(pub Id);
+
+impl NodeRef {
+    pub fn id(self) -> Id {
+        self.0
+    }
+}
+
+impl From<NodeRef> for Id {
+    fn from(value: NodeRef) -> Self {
+        value.0
+    }
+}
+
+/// Like [`std::fmt::Display`], but the formatting routine is given access to
+/// the [`Ast`] so that node references can resolve to their source text.
+///
+/// All standard primitive and string types implement [`YeastDisplay`] via
+/// the [`impl_yeast_display_via_display`] macro below. Coherence prevents a
+/// blanket `impl<T: Display>`, so additional types must be added explicitly.
+pub trait YeastDisplay {
+    fn yeast_to_string(&self, ast: &Ast) -> String;
+}
+
+impl YeastDisplay for NodeRef {
+    fn yeast_to_string(&self, ast: &Ast) -> String {
+        ast.source_text(self.0)
+    }
+}
+
+macro_rules! impl_yeast_display_via_display {
+    ($($t:ty),* $(,)?) => {
+        $(
+            impl YeastDisplay for $t {
+                fn yeast_to_string(&self, _ast: &Ast) -> String {
+                    ::std::string::ToString::to_string(self)
+                }
+            }
+        )*
+    };
+}
+
+impl_yeast_display_via_display! {
+    i8, i16, i32, i64, i128, isize,
+    u8, u16, u32, u64, u128, usize,
+    f32, f64,
+    bool, char,
+    str, String,
+}
+
+impl<T: YeastDisplay + ?Sized> YeastDisplay for &T {
+    fn yeast_to_string(&self, ast: &Ast) -> String {
+        (**self).yeast_to_string(ast)
+    }
+}
+
 pub const CHILD_FIELD: u16 = u16::MAX;
 
 #[derive(Debug)]
@@ -160,6 +221,9 @@ pub struct Ast {
     root: Id,
     nodes: Vec<Node>,
     schema: schema::Schema,
+    /// Original source bytes the tree was parsed from. Used to resolve
+    /// `NodeContent::Range` to text for synthesized literal nodes.
+    source: Vec<u8>,
 }
 
 impl std::fmt::Debug for Ast {
@@ -182,11 +246,41 @@ impl Ast {
         schema: schema::Schema,
         tree: &tree_sitter::Tree,
         language: &tree_sitter::Language,
+    ) -> Self {
+        Self::from_tree_with_schema_and_source(schema, tree, language, Vec::new())
+    }
+
+    pub fn from_tree_with_schema_and_source(
+        schema: schema::Schema,
+        tree: &tree_sitter::Tree,
+        language: &tree_sitter::Language,
+        source: Vec<u8>,
     ) -> Self {
         let mut visitor = visitor::Visitor::new(language.clone());
         visitor.visit(tree);
 
-        visitor.build_with_schema(schema)
+        let mut ast = visitor.build_with_schema(schema);
+        ast.source = source;
+        ast
+    }
+
+    /// Returns the source text for `id`, resolving `NodeContent::Range`
+    /// against the stored source bytes when available.
+    pub fn source_text(&self, id: Id) -> String {
+        let Some(node) = self.get_node(id) else { return String::new(); };
+        match &node.content {
+            NodeContent::Range(range) => {
+                let start = range.start_byte;
+                let end = range.end_byte;
+                if end <= self.source.len() && start <= end {
+                    String::from_utf8_lossy(&self.source[start..end]).into_owned()
+                } else {
+                    String::new()
+                }
+            }
+            NodeContent::String(s) => s.to_string(),
+            NodeContent::DynamicString(s) => s.clone(),
+        }
     }
 
     pub fn walk(&self) -> AstCursor {
@@ -894,8 +988,17 @@ impl<'a> Runner<'a> {
         })
     }
 
-    pub fn run_from_tree(&self, tree: &tree_sitter::Tree) -> Result<Ast, String> {
-        let mut ast = Ast::from_tree_with_schema(self.schema.clone(), tree, &self.language);
+    pub fn run_from_tree(
+        &self,
+        tree: &tree_sitter::Tree,
+        source: &[u8],
+    ) -> Result<Ast, String> {
+        let mut ast = Ast::from_tree_with_schema_and_source(
+            self.schema.clone(),
+            tree,
+            &self.language,
+            source.to_vec(),
+        );
         self.run_phases(&mut ast)?;
         Ok(ast)
     }
@@ -908,7 +1011,12 @@ impl<'a> Runner<'a> {
         let tree = parser
             .parse(input, None)
             .ok_or_else(|| "Failed to parse input".to_string())?;
-        let mut ast = Ast::from_tree_with_schema(self.schema.clone(), &tree, &self.language);
+        let mut ast = Ast::from_tree_with_schema_and_source(
+            self.schema.clone(),
+            &tree,
+            &self.language,
+            input.as_bytes().to_vec(),
+        );
         self.run_phases(&mut ast)?;
         Ok(ast)
     }

shared/yeast/src/visitor.rs

@@ -52,6 +52,7 @@ impl Visitor {
             root: 0,
             schema,
             nodes: self.nodes.into_iter().map(|n| n.inner).collect(),
+            source: Vec::new(),
         }
     }
 

shared/yeast/tests/test.rs

@@ -1060,3 +1060,54 @@ fn test_desugar_for_with_multiple_assignment() {
     "#,
     );
 }
+
+/// Regression test: `#{capture}` in a template must render the *source text*
+/// of the captured node, not its arena `Id`. Previously, captures were bound
+/// as `usize`, so `#{cap}` printed the integer id (e.g. `"3"`) via `Display`.
+/// Captures are now bound as `NodeRef`, which has no `Display` impl and
+/// resolves to the captured node's source text via `YeastDisplay`.
+#[test]
+fn test_hash_brace_renders_capture_source_text() {
+    let rule = rule!(
+        (call
+            method: (identifier) @name
+            receiver: (identifier) @recv
+        )
+        =>
+        (call
+            method: (identifier #{name})
+            receiver: (identifier #{recv})
+            arguments: (argument_list)
+        )
+    );
+    let dump = run_and_dump("foo.bar()", vec![rule]);
+    assert_dump_eq(
+        &dump,
+        r#"
+        program
+          call
+            arguments: argument_list "foo.bar()"
+            method: identifier "bar"
+            receiver: identifier "foo"
+    "#,
+    );
+}
+
+/// Regression test: non-`NodeRef` values in `#{expr}` still render via their
+/// `Display` impl (covered by `YeastDisplay`'s blanket impls for primitives).
+#[test]
+fn test_hash_brace_renders_integer_expression() {
+    let rule = rule!(
+        (identifier) @_
+        =>
+        (identifier #{1 + 2})
+    );
+    let dump = run_and_dump("foo", vec![rule]);
+    assert_dump_eq(
+        &dump,
+        r#"
+        program
+          identifier "3"
+    "#,
+    );
+}

unified/extractor/tests/corpus/swift/desugar.txt

@@ -17,7 +17,7 @@ source_file
 top_level
   body:
     binary_expr
-      operator: operator "3"
+      operator: operator "+"
       left: unsupported_node "1"
       right: unsupported_node "2"
 
@@ -40,6 +40,6 @@ source_file
 top_level
   body:
     binary_expr
-      operator: operator "3"
+      operator: operator "+"
       left: name_expr "foo"
       right: name_expr "bar"