feat(ast_tools): generate mapping of trait name to crate in oxc_ast_macros (#8870)

`oxc_ast_macros` contains a function which takes the name of a trait, and returns the crate that trait is defined in.

Instead of having that info hard-coded in 2 places, codify it in generator for that trait in `oxc_ast_tools`, and codegen this function in `oxc_ast_macros` instead.

.github/.generated_ast_watch_list.yml

@@ -20,6 +20,7 @@ src:
   - 'crates/oxc_ast/src/generated/get_id.rs'
   - 'crates/oxc_ast/src/generated/visit.rs'
   - 'crates/oxc_ast/src/generated/visit_mut.rs'
+  - 'crates/oxc_ast_macros/src/generated/mod.rs'
   - 'crates/oxc_ast_macros/src/lib.rs'
   - 'crates/oxc_regular_expression/src/ast.rs'
   - 'crates/oxc_regular_expression/src/generated/derive_clone_in.rs'

crates/oxc_ast_macros/src/ast.rs

@@ -2,6 +2,8 @@ use proc_macro2::TokenStream;
 use quote::quote;
 use syn::{punctuated::Punctuated, token::Comma, Attribute, Fields, Ident, Item, ItemEnum};
 
+use crate::generated::get_trait_crate_and_generics;
+
 pub fn ast(input: &Item) -> TokenStream {
     let (head, tail) = match input {
         Item::Enum(enum_) => (enum_repr(enum_), assert_generated_derives(&enum_.attrs)),
@@ -44,19 +46,20 @@ fn enum_repr(enum_: &ItemEnum) -> TokenStream {
 fn assert_generated_derives(attrs: &[Attribute]) -> TokenStream {
     // NOTE: At this level we don't care if a trait is derived multiple times, It is the
     // responsibility of the `ast_tools` to raise errors for those.
-    let assertion = attrs
+    let assertions = attrs
         .iter()
         .filter(|attr| attr.path().is_ident("generate_derive"))
         .flat_map(parse_attr)
-        .map(|derive| {
-            let (abs_derive, generics) = abs_trait(&derive);
+        .map(|trait_ident| {
+            let trait_name = trait_ident.to_string();
+            let (trait_path, generics) = get_trait_crate_and_generics(&trait_name);
             quote! {{
                 // NOTE: these are wrapped in a scope to avoid the need for unique identifiers.
-                trait AssertionTrait: #abs_derive #generics {}
-                impl<T: #derive #generics> AssertionTrait for T {}
+                trait AssertionTrait: #trait_path #generics {}
+                impl<T: #trait_ident #generics> AssertionTrait for T {}
             }}
         });
-    quote!(const _: () = { #(#assertion)* };)
+    quote!( const _: () = { #(#assertions)* }; )
 }
 
 #[inline]
@@ -65,34 +68,3 @@ fn parse_attr(attr: &Attribute) -> impl Iterator<Item = Ident> {
         .expect("`#[generate_derive]` only accepts traits as single segment paths. Found an invalid argument.")
         .into_iter()
 }
-
-// TODO: benchmark this to see if a lazy static cell containing `HashMap` would perform better.
-#[inline]
-fn abs_trait(
-    ident: &Ident,
-) -> (/* absolute type path */ TokenStream, /* possible generics */ TokenStream) {
-    if ident == "CloneIn" {
-        (quote!(::oxc_allocator::CloneIn), quote!(<'static>))
-    } else if ident == "GetSpan" {
-        (quote!(::oxc_span::GetSpan), TokenStream::default())
-    } else if ident == "GetSpanMut" {
-        (quote!(::oxc_span::GetSpanMut), TokenStream::default())
-    } else if ident == "GetAddress" {
-        (quote!(::oxc_allocator::GetAddress), TokenStream::default())
-    } else if ident == "ContentEq" {
-        (quote!(::oxc_span::ContentEq), TokenStream::default())
-    } else if ident == "ESTree" {
-        (quote!(::oxc_estree::ESTree), TokenStream::default())
-    } else {
-        invalid_derive(ident)
-    }
-}
-
-#[cold]
-fn invalid_derive(ident: &Ident) -> ! {
-    panic!(
-        "Invalid derive trait(generate_derive): {ident}.\n\
-        Help: If you are trying to implement a new `generate_derive` trait, \
-        make sure to add it to the list in `abs_trait` function."
-    )
-}

crates/oxc_ast_macros/src/generated/mod.rs

@@ -0,0 +1,17 @@
+// Auto-generated code, DO NOT EDIT DIRECTLY!
+// To edit this generated file you have to edit `tasks/ast_tools/src/main.rs`
+
+use proc_macro2::TokenStream;
+use quote::quote;
+
+pub fn get_trait_crate_and_generics(trait_name: &str) -> (TokenStream, TokenStream) {
+    match trait_name {
+        "CloneIn" => (quote!(::oxc_allocator::CloneIn), quote!(< 'static >)),
+        "GetAddress" => (quote!(::oxc_allocator::GetAddress), TokenStream::new()),
+        "GetSpan" => (quote!(::oxc_span::GetSpan), TokenStream::new()),
+        "GetSpanMut" => (quote!(::oxc_span::GetSpanMut), TokenStream::new()),
+        "ContentEq" => (quote!(::oxc_span::ContentEq), TokenStream::new()),
+        "ESTree" => (quote!(::oxc_estree::ESTree), TokenStream::new()),
+        _ => panic!("Invalid derive trait(generate_derive): {trait_name}"),
+    }
+}

crates/oxc_ast_macros/src/lib.rs

@@ -2,6 +2,7 @@ use proc_macro::TokenStream;
 use syn::{parse_macro_input, Item};
 
 mod ast;
+mod generated;
 
 /// This attribute serves two purposes:
 ///

tasks/ast_tools/src/derives/clone_in.rs

@@ -21,6 +21,14 @@ impl Derive for DeriveCloneIn {
         "CloneIn"
     }
 
+    fn trait_has_lifetime(&self) -> bool {
+        true
+    }
+
+    fn crate_name(&self) -> &'static str {
+        "oxc_allocator"
+    }
+
     /// Register that accept `#[clone_in]` attr on struct fields.
     fn attrs(&self) -> &[(&'static str, AttrPositions)] {
         &[("clone_in", AttrPositions::StructField)]

tasks/ast_tools/src/derives/content_eq.rs

@@ -19,6 +19,10 @@ impl Derive for DeriveContentEq {
         "ContentEq"
     }
 
+    fn crate_name(&self) -> &'static str {
+        "oxc_span"
+    }
+
     fn prelude(&self) -> TokenStream {
         quote! {
             #![allow(clippy::match_same_arms)]

tasks/ast_tools/src/derives/estree.rs

@@ -26,6 +26,10 @@ impl Derive for DeriveESTree {
         "ESTree"
     }
 
+    fn crate_name(&self) -> &'static str {
+        "oxc_estree"
+    }
+
     fn snake_name(&self) -> String {
         "estree".to_string()
     }

tasks/ast_tools/src/derives/get_address.rs

@@ -17,6 +17,10 @@ impl Derive for DeriveGetAddress {
         "GetAddress"
     }
 
+    fn crate_name(&self) -> &'static str {
+        "oxc_allocator"
+    }
+
     fn prelude(&self) -> TokenStream {
         quote! {
             #![allow(clippy::match_same_arms)]

tasks/ast_tools/src/derives/get_span.rs

@@ -21,6 +21,10 @@ impl Derive for DeriveGetSpan {
         "GetSpan"
     }
 
+    fn crate_name(&self) -> &'static str {
+        "oxc_span"
+    }
+
     /// Register that accept `#[span]` attr on struct fields.
     fn attrs(&self) -> &[(&'static str, AttrPositions)] {
         &[("span", AttrPositions::StructField)]
@@ -79,6 +83,11 @@ impl Derive for DeriveGetSpanMut {
         "GetSpanMut"
     }
 
+    /// Get crate trait is defined in.
+    fn crate_name(&self) -> &'static str {
+        "oxc_span"
+    }
+
     fn prelude(&self) -> TokenStream {
         quote! {
             #![allow(clippy::match_same_arms)]

tasks/ast_tools/src/derives/mod.rs

@@ -30,6 +30,16 @@ pub trait Derive: Runner {
     /// Get trait name.
     fn trait_name(&self) -> &'static str;
 
+    /// Get if trait has lifetime.
+    ///
+    /// Default to `false`, but can be overridden.
+    fn trait_has_lifetime(&self) -> bool {
+        false
+    }
+
+    /// Get crate trait is defined in.
+    fn crate_name(&self) -> &'static str;
+
     /// Get snake case trait name.
     ///
     /// Defaults to `trait_name()` converted to snake case.

tasks/ast_tools/src/main.rs

@@ -172,6 +172,7 @@
 use std::{fmt::Write, fs};
 
 use bpaf::{Bpaf, Parser};
+use quote::quote;
 use rayon::prelude::*;
 
 mod codegen;
@@ -187,9 +188,10 @@ use codegen::{get_runners, Codegen, Runner};
 use derives::Derive;
 use generators::Generator;
 use logger::{log, log_failed, log_result, log_success, logln};
-use output::{Output, RawOutput};
+use output::{output_path, Output, RawOutput};
 use parse::parse_files;
 use schema::Schema;
+use utils::create_ident;
 
 /// Paths to source files containing AST types
 static SOURCE_PATHS: &[&str] = &[
@@ -295,6 +297,7 @@ fn main() {
     logln!("All Derives and Generators... Done!");
 
     // Edit `lib.rs` in `oxc_ast_macros` crate
+    outputs.push(generate_proc_macro());
     outputs.push(generate_updated_proc_macro(&codegen));
 
     // Add CI filter file to outputs
@@ -336,6 +339,38 @@ fn generate_ci_filter(outputs: &[RawOutput]) -> RawOutput {
     Output::Yaml { path: GITHUB_WATCH_LIST_PATH.to_string(), code }.into_raw(file!())
 }
 
+/// Generate function for proc macro in `oxc_ast_macros` crate.
+///
+/// This function translates trait name to path to the trait and any generic params.
+fn generate_proc_macro() -> RawOutput {
+    let match_arms = DERIVES.iter().map(|derive| {
+        let trait_name = derive.trait_name();
+        let trait_ident = create_ident(trait_name);
+        let crate_ident = create_ident(derive.crate_name());
+        if derive.trait_has_lifetime() {
+            quote!( #trait_name => (quote!(::#crate_ident::#trait_ident), quote!( <'static> )) )
+        } else {
+            quote!( #trait_name => (quote!(::#crate_ident::#trait_ident), TokenStream::new()) )
+        }
+    });
+
+    let output = quote! {
+        use proc_macro2::TokenStream;
+        use quote::quote;
+
+        ///@@line_break
+        pub fn get_trait_crate_and_generics(trait_name: &str) -> (TokenStream, TokenStream) {
+            match trait_name {
+                #(#match_arms,)*
+                _ => panic!("Invalid derive trait(generate_derive): {trait_name}"),
+            }
+        }
+    };
+
+    Output::Rust { path: output_path(AST_MACROS_CRATE_PATH, "mod.rs"), tokens: output }
+        .into_raw(file!())
+}
+
 /// Update the list of helper attributes for `Ast` derive proc macro in `oxc_ast_macros` crate
 /// to include all attrs which generators/derives utilize.
 ///