Merge pull request #22875 from github/tausbn/yeast-split-input-output-arenas

yeast: Split input/output ASTs

Author: tausbn

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

@@ -66,7 +66,7 @@ impl<'a> AstNode for Node<'a> {
     }
 }
 
-impl AstNode for yeast::Node {
+impl<I: yeast::TreeId> AstNode for yeast::Node<I> {
     fn kind(&self) -> &str {
         yeast::Node::kind_name(self)
     }
@@ -1043,7 +1043,7 @@ fn traverse(tree: &Tree, visitor: &mut Visitor) {
     }
 }
 
-fn traverse_yeast(tree: &yeast::Ast, visitor: &mut Visitor) {
+fn traverse_yeast(tree: &yeast::OutputAst, visitor: &mut Visitor) {
     let mut cursor = tree.walk();
     visitor.enter_node(cursor.node());
     let mut recurse = true;

shared/yeast-macros/src/lib.rs

@@ -50,17 +50,17 @@ pub fn query(input: TokenStream) -> TokenStream {
     }
 }
 
-/// Build a single AST node from a template, returning its `Id`.
+/// Build a single AST node from a template, returning its `OutputId`.
 ///
 /// # Template syntax
 ///
 /// ```text
 /// (kind "literal")             - leaf with static content
 /// (kind #{expr})               - leaf with computed content (expr.to_string())
 /// {expr}                       - embed a Rust expression, dispatched via
-///                                the `IntoFieldIds` trait: `Id` pushes a
-///                                single id; iterables (`Vec<Id>`,
-///                                `Option<Id>`, iterator chains) splice
+///                                the `IntoFieldIds` trait: `OutputId` pushes a
+///                                single id; iterables (`Vec<OutputId>`,
+///                                `Option<OutputId>`, iterator chains) splice
 ///                                their elements
 /// field: {expr}                - extend a named field with `{expr}`'s ids
 /// field: (kind ...)?           - set the field only if every `#{expr}`
@@ -108,9 +108,9 @@ pub fn tree(input: TokenStream) -> TokenStream {
     }
 }
 
-/// Build a list of AST nodes from a template, returning `Vec<Id>`.
+/// Build a list of AST nodes from a template, returning `Vec<OutputId>`.
 ///
-/// Like `tree!` but returns `Vec<Id>` and supports multiple top-level
+/// Like `tree!` but returns `Vec<OutputId>` and supports multiple top-level
 /// elements. All syntax from `tree!` is available.
 ///
 /// Uses the `BuildCtx` binding named `ctx` from the surrounding scope:
@@ -148,9 +148,11 @@ pub fn trees(input: TokenStream) -> TokenStream {
 /// ```
 ///
 /// Captures become Rust variables automatically:
-/// - `@name` (no quantifier) → `name: Id`
-/// - `@name` (after `*`/`+`) → `name: Vec<Id>`
-/// - `@name` (after `?`) → `name: Option<Id>`
+/// - `@name` (no quantifier) → `name: OutputId`
+/// - `@name` (after `*`/`+`) → `name: Vec<OutputId>`
+/// - `@name` (after `?`) → `name: Option<OutputId>`
+///
+/// Raw `@@name` captures use the corresponding input `Id` types instead.
 ///
 /// A guard is an optional Rust condition between the query and `=>`. It is
 /// evaluated after the query matches and before any `@` captures are
@@ -190,8 +192,8 @@ pub fn rule(input: TokenStream) -> TokenStream {
 ///     output: "ast_types.yml",
 ///     [
 ///         (source_file (_)* @cs) => (top_level body: {..cs}),
-///         (simple_identifier) @id => (name_expr identifier: (identifier #{id})),
-///         rule!((integer_literal) @lit => (int_literal #{lit})),
+///         (simple_identifier) @@id => (name_expr identifier: (identifier #{id})),
+///         rule!((integer_literal) @@lit => (int_literal #{lit})),
 ///         helper_fn(),
 ///     ]
 /// };

shared/yeast-macros/src/lower.rs

@@ -33,8 +33,20 @@ impl Rule {
             .iter()
             .map(|capture| capture.name.to_string())
             .collect::<Vec<_>>();
-        let raw_bindings = capture_bindings(raw_captures.into_iter());
-        let translated_bindings = capture_bindings(translated_captures.into_iter());
+        let captures_ident = Ident::new("__captures", Span::call_site());
+        let translated_captures_ident = Ident::new("__translated_captures", Span::call_site());
+        let guard_bindings =
+            capture_bindings(captures.iter(), quote! { yeast::Id }, &captures_ident);
+        let raw_bindings = capture_bindings(
+            raw_captures.into_iter(),
+            quote! { yeast::Id },
+            &captures_ident,
+        );
+        let translated_bindings = capture_bindings(
+            translated_captures.into_iter(),
+            quote! { yeast::OutputId },
+            &translated_captures_ident,
+        );
         let replacement = self.replacement.lower();
         let guard = self
             .guard
@@ -48,20 +60,20 @@ impl Rule {
                 yeast::Rule::guarded(
                     __query,
                     Box::new(|__ast: &yeast::Ast, __captures: &yeast::captures::Captures, __user_ctx: &mut _| {
-                        #(#raw_bindings)*
-                        #(#translated_bindings)*
+                        #(#guard_bindings)*
                         let ast = __ast;
                         let #ctx = __user_ctx;
                         Ok(#guard)
                     }),
-                    Box::new(|__ast: &mut yeast::Ast, mut __captures: yeast::captures::Captures, __source_range: Option<yeast::Range>, __user_ctx: &mut _, __translator: yeast::TranslatorHandle<'_, _>| {
+                    Box::new(|__ast: &yeast::Ast, __output: &mut yeast::OutputAst, __captures: yeast::captures::Captures, __source_range: Option<yeast::Range>, __user_ctx: &mut _, __translator: yeast::TranslatorHandle<'_, _>| {
                         let __skip: &[&str] = &[#(#raw_capture_names),*];
-                        __translator.auto_translate_captures(&mut __captures, __ast, __user_ctx, __skip)?;
+                        let __translated_captures =
+                            __translator.translate_captures(&__captures, __ast, __output, __user_ctx, __skip)?;
                         #(#raw_bindings)*
                         #(#translated_bindings)*
-                        let mut #ctx = yeast::build::BuildCtx::with_translator(__ast, &__captures, __source_range, __user_ctx, __translator);
-                        let __result: Vec<yeast::Id> = { #replacement };
-                        let __result = #ctx.finish_rule(__result);
+                        let mut #ctx = yeast::build::BuildCtx::with_translator(__ast, __output, __source_range, __user_ctx, __translator);
+                        let __result: Vec<yeast::OutputId> = { #replacement };
+                        let __result = #ctx.finish_rule(__result)?;
                         Ok(__result)
                     }),
                 )
@@ -77,9 +89,10 @@ impl Replacement {
             Replacement::Rust { body, .. } => {
                 let body = rust_block(body);
                 quote! {
-                    let __value = #body;
-                    let mut __ids: Vec<yeast::Id> = Vec::new();
-                    yeast::IntoFieldIds::extend_into(__value, &mut __ids);
+                    let __value = ::std::iter::IntoIterator::into_iter(#body)
+                        .collect::<Vec<_>>();
+                    let mut __ids: Vec<yeast::OutputId> = Vec::new();
+                    ctx.extend_field_ids(__value, &mut __ids);
                     __ids
                 }
             }
@@ -241,25 +254,29 @@ fn combine_cardinality(parent: Cardinality, child: Cardinality) -> Cardinality {
     }
 }
 
-fn capture_bindings<'a>(captures: impl Iterator<Item = &'a CaptureBinding>) -> Vec<TokenStream> {
+fn capture_bindings<'a>(
+    captures: impl Iterator<Item = &'a CaptureBinding>,
+    id_type: TokenStream,
+    captures_ident: &Ident,
+) -> Vec<TokenStream> {
     captures
         .map(|capture| {
             let name_string = capture.name.to_string();
             let name = Ident::new(&name_string, Span::call_site());
             match (capture.cardinality.multiple, capture.cardinality.required) {
                 (true, _) => {
                     quote! {
-                        let #name: Vec<yeast::Id> = __captures.get_all(#name_string);
+                        let #name: Vec<#id_type> = #captures_ident.get_all(#name_string);
                     }
                 }
                 (false, false) => {
                     quote! {
-                        let #name: Option<yeast::Id> = __captures.get_opt(#name_string);
+                        let #name: Option<#id_type> = #captures_ident.get_opt(#name_string);
                     }
                 }
                 (false, true) => {
                     quote! {
-                        let #name: yeast::Id = __captures.get_var(#name_string).unwrap();
+                        let #name: #id_type = #captures_ident.get_var(#name_string).unwrap();
                     }
                 }
             }
@@ -288,15 +305,17 @@ impl Template {
                 Template::Splice(block) => {
                     let block = rust_block(block);
                     quote! {
-                        yeast::IntoFieldIds::extend_into(#block, &mut __nodes);
+                        let __value = ::std::iter::IntoIterator::into_iter(#block)
+                            .collect::<Vec<_>>();
+                        ctx.extend_field_ids(__value, &mut __nodes);
                     }
                 }
             })
             .collect::<Vec<_>>();
 
         quote! {
             {
-                let mut __nodes: Vec<yeast::Id> = Vec::new();
+                let mut __nodes: Vec<yeast::OutputId> = Vec::new();
                 #(#items)*
                 __nodes
             }
@@ -308,7 +327,7 @@ impl Template {
             Template::Node(node) => node.lower(context),
             Template::Splice(block) => {
                 let block = rust_block(block);
-                quote! { ::std::convert::Into::<yeast::Id>::into(#block) }
+                quote! { ::std::convert::Into::<yeast::OutputId>::into(#block) }
             }
         }
     }
@@ -334,9 +353,9 @@ impl Node {
                                     break #label ::std::option::Option::None;
                                 };
                                 let __value =
-                                    yeast::YeastDisplay::yeast_to_string(__value_ref, &*#ctx.ast);
+                                    yeast::YeastDisplay::yeast_to_string(__value_ref, &#ctx);
                                 let __source_range =
-                                    yeast::YeastSourceRange::yeast_source_range(__value_ref, &*#ctx.ast);
+                                    yeast::YeastSourceRange::yeast_source_range(__value_ref, &#ctx);
                                 #ctx.literal_with_source_range(
                                     #kind_str,
                                     &__value,
@@ -349,9 +368,9 @@ impl Node {
                             {
                                 let __expr = #block;
                                 let __value =
-                                    yeast::YeastDisplay::yeast_to_string(&__expr, &*#ctx.ast);
+                                    yeast::YeastDisplay::yeast_to_string(&__expr, &#ctx);
                                 let __source_range =
-                                    yeast::YeastSourceRange::yeast_source_range(&__expr, &*#ctx.ast);
+                                    yeast::YeastSourceRange::yeast_source_range(&__expr, &#ctx);
                                 #ctx.literal_with_source_range(
                                     #kind_str,
                                     &__value,
@@ -362,7 +381,7 @@ impl Node {
                     }
                 }
             };
-            quote! { let __node: yeast::Id = #node; }
+            quote! { let __node: yeast::OutputId = #node; }
         } else {
             TokenStream::new()
         };
@@ -375,7 +394,9 @@ impl Node {
                 Field::Location { block } => {
                     let block = rust_block(block);
                     stmts.push(quote! {
-                        yeast::IntoFieldIds::extend_into(#block, &mut __location_ids);
+                        let __value = ::std::iter::IntoIterator::into_iter(#block)
+                            .collect::<Vec<_>>();
+                        #ctx.extend_location_ranges(__value, &mut __location_ranges);
                     });
                 }
                 Field::Named {
@@ -398,8 +419,10 @@ impl Node {
                             debug_assert!(!optional);
                             let block = rust_block(block);
                             stmts.push(quote! {
-                                let mut #temp: Vec<yeast::Id> = Vec::new();
-                                yeast::IntoFieldIds::extend_into(#block, &mut #temp);
+                                let mut #temp: Vec<yeast::OutputId> = Vec::new();
+                                let __value = ::std::iter::IntoIterator::into_iter(#block)
+                                    .collect::<Vec<_>>();
+                                #ctx.extend_field_ids(__value, &mut #temp);
                             });
                             field_args.push(quote! {
                                 if !#temp.is_empty() {
@@ -414,7 +437,7 @@ impl Node {
                             });
                             stmts.push(quote! {
                                 #[allow(unused_labels)]
-                                let #temp: ::std::option::Option<yeast::Id> = #label: {
+                                let #temp: ::std::option::Option<yeast::OutputId> = #label: {
                                     ::std::option::Option::Some(#value)
                                 };
                             });
@@ -426,7 +449,7 @@ impl Node {
                         }
                         Template::Node(child) => {
                             let value = child.lower(context);
-                            stmts.push(quote! { let #temp: yeast::Id = #value; });
+                            stmts.push(quote! { let #temp: yeast::OutputId = #value; });
                             field_args.push(quote! {
                                 __fields.push((#field_str, vec![#temp]));
                             });
@@ -438,28 +461,21 @@ impl Node {
 
         let build_node = if self.content.is_none() {
             quote! {
-                let mut __fields: Vec<(&str, Vec<yeast::Id>)> = Vec::new();
+                let mut __fields: Vec<(&str, Vec<yeast::OutputId>)> = Vec::new();
                 #(#field_args)*
-                let __node: yeast::Id = #ctx.node(#kind_str, __fields);
+                let __node: yeast::OutputId = #ctx.node(#kind_str, __fields);
             }
         } else {
             TokenStream::new()
         };
 
         quote! {
             {
-                let mut __location_ids: Vec<yeast::Id> = Vec::new();
+                let mut __location_ranges: Vec<yeast::Range> = Vec::new();
                 #node_init
                 #(#stmts)*
                 #build_node
-                let __source_range = __location_ids
-                    .into_iter()
-                    .filter_map(|__source| {
-                        #ctx.ast
-                            .get_node(__source)
-                            .and_then(|__node| __node.source_range())
-                    })
-                    .reduce(yeast::Range::union);
+                let __source_range = __location_ranges.into_iter().reduce(yeast::Range::union);
                 #ctx.set_node_source_range(__node, __source_range)
             }
         }

shared/yeast-macros/src/rule_parse.rs

@@ -59,8 +59,18 @@ fn validate_unique_capture_names(
 pub(crate) fn parse_rule(input: TokenStream) -> Result<Rule> {
     let (pattern_and_guard, replacement) = split_arrow(input)?;
     let (pattern, guard) = split_guard(pattern_and_guard)?;
+    let pattern = parse_pattern(pattern)?;
+    if let Pattern::Capture {
+        name, raw: false, ..
+    } = &pattern
+    {
+        return Err(syn::Error::new_spanned(
+            name,
+            "a capture of the rule root must use `@@`; the matched root cannot be eagerly translated",
+        ));
+    }
     Ok(Rule {
-        pattern: parse_pattern(pattern)?,
+        pattern,
         guard,
         replacement: parse_replacement(replacement)?,
     })
@@ -548,7 +558,7 @@ mod tests {
     #[test]
     fn parses_template_replacement() {
         let rule = parse_rule(quote! {
-            (identifier) @name
+            (identifier) @@name
             =>
             (identifier #{name})
         })
@@ -558,6 +568,22 @@ mod tests {
         assert!(matches!(rule.replacement, Replacement::Templates(_)));
     }
 
+    #[test]
+    fn rejects_eager_capture_of_rule_root() {
+        let result = parse_rule(quote! {
+            (identifier) @name
+            =>
+            (identifier #{name})
+        });
+        let Err(error) = result else {
+            panic!("expected eager rule-root capture to be rejected");
+        };
+        assert_eq!(
+            error.to_string(),
+            "a capture of the rule root must use `@@`; the matched root cannot be eagerly translated"
+        );
+    }
+
     #[test]
     fn repeated_capture_owns_its_cardinality() {
         let pattern = parse_pattern(quote!((array (identifier)* @items))).unwrap();

shared/yeast/src/bin/main.rs

@@ -20,7 +20,14 @@ fn main() {
     let args = Cli::parse();
     let language = get_language(&args.language);
     let source = std::fs::read_to_string(&args.file).unwrap();
-    let runner: yeast::Runner = yeast::Runner::new(language, &[]);
-    let ast = runner.run(&source).unwrap();
+    let mut parser = tree_sitter::Parser::new();
+    parser.set_language(&language).unwrap();
+    let tree = parser.parse(&source, None).unwrap();
+    let ast = yeast::Ast::from_tree_with_schema_and_source(
+        yeast::schema::from_language(&language),
+        &tree,
+        &language,
+        source.as_bytes().to_vec(),
+    );
     println!("{}", ast.print(&source, ast.get_root()));
 }