refactor(ast_tools): reduce macro usage (#6895)

1. Reduce the amount of code in `define_derive!` and `define_generator!` macros. This makes the code easier to read, and gives type hints in IDE.
2. Remove `generated_header!` macro and insert header as a blanket action, instead of repeated code in every generator.

crates/oxc_ast/src/generated/ast_builder.rs

@@ -1,7 +1,8 @@
-//! AST node factories
 // Auto-generated code, DO NOT EDIT DIRECTLY!
 // To edit this generated file you have to edit `tasks/ast_tools/src/generators/ast_builder.rs`
 
+//! AST node factories
+
 #![allow(
     clippy::default_trait_access,
     clippy::too_many_arguments,

tasks/ast_tools/src/codegen.rs

@@ -238,18 +238,20 @@ impl AstCodegen {
     }
 }
 
-/// Creates a generated file warning + required information for a generated file.
-macro_rules! generated_header {
-    () => {{
-        let file = file!().replace("\\", "/");
-        // TODO add generation date, AST source hash, etc here.
-        let edit_comment = format!("@ To edit this generated file you have to edit `{file}`");
-        quote::quote! {
-            //!@ Auto-generated code, DO NOT EDIT DIRECTLY!
-            #![doc = #edit_comment]
-            //!@@line_break
-        }
-    }};
+/// Implemented by `define_derive!` and `define_generator!` macros
+pub trait CodegenBase {
+    fn file_path() -> &'static str;
 }
 
-pub(crate) use generated_header;
+/// Creates a generated file warning + required information for a generated file.
+pub fn generate_header(file_path: &str) -> TokenStream {
+    let file_path = file_path.replace('\\', "/");
+
+    // TODO: Add generation date, AST source hash, etc here.
+    let edit_comment = format!("@ To edit this generated file you have to edit `{file_path}`");
+    quote::quote! {
+        //!@ Auto-generated code, DO NOT EDIT DIRECTLY!
+        #![doc = #edit_comment]
+        //!@@line_break
+    }
+}

tasks/ast_tools/src/derives/clone_in.rs

@@ -3,16 +3,16 @@ use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 use syn::Ident;
 
-use super::{define_derive, Derive, DeriveOutput};
+use super::{define_derive, Derive};
 use crate::{
     codegen::LateCtx,
     markers::CloneInAttribute,
     schema::{EnumDef, GetIdent, StructDef, TypeDef},
 };
 
-define_derive! {
-    pub struct DeriveCloneIn;
-}
+pub struct DeriveCloneIn;
+
+define_derive!(DeriveCloneIn);
 
 impl Derive for DeriveCloneIn {
     fn trait_name() -> &'static str {

tasks/ast_tools/src/derives/content_eq.rs

@@ -2,16 +2,16 @@ use itertools::Itertools;
 use proc_macro2::TokenStream;
 use quote::quote;
 
-use super::{define_derive, Derive, DeriveOutput};
+use super::{define_derive, Derive};
 use crate::{
     codegen::LateCtx,
     schema::{EnumDef, GetGenerics, StructDef, ToType, TypeDef},
     util::ToIdent,
 };
 
-define_derive! {
-    pub struct DeriveContentEq;
-}
+pub struct DeriveContentEq;
+
+define_derive!(DeriveContentEq);
 
 const IGNORE_FIELDS: [(/* field name */ &str, /* field type */ &str); 6] = [
     ("span", "Span"),

tasks/ast_tools/src/derives/content_hash.rs

@@ -2,16 +2,16 @@ use itertools::Itertools;
 use proc_macro2::TokenStream;
 use quote::quote;
 
-use super::{define_derive, Derive, DeriveOutput};
+use super::{define_derive, Derive};
 use crate::{
     codegen::LateCtx,
     schema::{EnumDef, GetGenerics, StructDef, ToType, TypeDef},
     util::ToIdent,
 };
 
-define_derive! {
-    pub struct DeriveContentHash;
-}
+pub struct DeriveContentHash;
+
+define_derive!(DeriveContentHash);
 
 const IGNORE_FIELD_TYPES: [/* type name */ &str; 4] = [
     "Span",

tasks/ast_tools/src/derives/estree.rs

@@ -2,7 +2,7 @@ use convert_case::{Case, Casing};
 use proc_macro2::TokenStream;
 use quote::quote;
 
-use super::{define_derive, Derive, DeriveOutput};
+use super::{define_derive, Derive};
 use crate::{
     codegen::LateCtx,
     markers::ESTreeStructAttribute,
@@ -12,9 +12,9 @@ use crate::{
     },
 };
 
-define_derive! {
-    pub struct DeriveESTree;
-}
+pub struct DeriveESTree;
+
+define_derive!(DeriveESTree);
 
 impl Derive for DeriveESTree {
     fn trait_name() -> &'static str {

tasks/ast_tools/src/derives/get_span.rs

@@ -2,16 +2,16 @@ use proc_macro2::TokenStream;
 use quote::quote;
 use syn::Ident;
 
-use super::{define_derive, Derive, DeriveOutput};
+use super::{define_derive, Derive};
 use crate::{
     codegen::LateCtx,
     schema::{EnumDef, GetGenerics, StructDef, ToType, TypeDef},
     util::{ToIdent, TypeWrapper},
 };
 
-define_derive! {
-    pub struct DeriveGetSpan;
-}
+pub struct DeriveGetSpan;
+
+define_derive!(DeriveGetSpan);
 
 impl Derive for DeriveGetSpan {
     fn trait_name() -> &'static str {
@@ -47,9 +47,9 @@ impl Derive for DeriveGetSpan {
     }
 }
 
-define_derive! {
-    pub struct DeriveGetSpanMut;
-}
+pub struct DeriveGetSpanMut;
+
+define_derive!(DeriveGetSpanMut);
 
 impl Derive for DeriveGetSpanMut {
     fn trait_name() -> &'static str {

tasks/ast_tools/src/derives/mod.rs

@@ -1,9 +1,15 @@
 use std::path::PathBuf;
 
 use convert_case::{Case, Casing};
+use itertools::Itertools;
 use proc_macro2::TokenStream;
+use rustc_hash::{FxHashMap, FxHashSet};
 
-use crate::{codegen::LateCtx, schema::TypeDef};
+use crate::{
+    codegen::{generate_header, CodegenBase, LateCtx},
+    schema::TypeDef,
+    Result,
+};
 
 mod clone_in;
 mod content_eq;
@@ -17,7 +23,12 @@ pub use content_hash::DeriveContentHash;
 pub use estree::DeriveESTree;
 pub use get_span::{DeriveGetSpan, DeriveGetSpanMut};
 
-pub trait Derive {
+#[derive(Debug, Clone)]
+pub struct DeriveOutput(pub Vec<(PathBuf, TokenStream)>);
+
+pub trait Derive: CodegenBase {
+    // Methods defined by implementer
+
     fn trait_name() -> &'static str;
 
     fn snake_name() -> String {
@@ -29,115 +40,118 @@ pub trait Derive {
     fn prelude() -> TokenStream {
         TokenStream::default()
     }
-}
 
-pub trait DeriveTemplate: Derive {
-    fn template(module_path: Vec<&str>, impls: TokenStream) -> TokenStream;
-}
+    // Standard methods
 
-#[derive(Debug, Clone)]
-pub struct DeriveOutput(pub Vec<(PathBuf, TokenStream)>);
+    fn template(module_paths: Vec<&str>, impls: TokenStream) -> TokenStream {
+        let header = generate_header(Self::file_path());
+        let prelude = Self::prelude();
+
+        // from `x::y::z` to `crate::y::z::*`
+        let use_modules = module_paths.into_iter().map(|it| {
+            let local_path = ["crate"]
+                .into_iter()
+                .chain(it.strip_suffix("::mod").unwrap_or(it).split("::").skip(1))
+                .chain(["*"])
+                .join("::");
+            let use_module: syn::ItemUse =
+                syn::parse_str(format!("use {local_path};").as_str()).unwrap();
+            quote::quote! {
+                ///@@line_break
+                #use_module
+            }
+        });
+
+        quote::quote! {
+            #header
+
+            #prelude
+
+            #(#use_modules)*
+
+            ///@@line_break
+            #impls
+        }
+    }
+
+    fn output(&mut self, ctx: &LateCtx) -> Result<DeriveOutput> {
+        let trait_name = Self::trait_name();
+        let filename = format!("derive_{}.rs", Self::snake_name());
+        let output = ctx
+            .schema()
+            .into_iter()
+            .filter(|def| def.generates_derive(trait_name))
+            .map(|def| (def, self.derive(def, ctx)))
+            .fold(
+                FxHashMap::<&str, (FxHashSet<&str>, Vec<TokenStream>)>::default(),
+                |mut acc, (def, stream)| {
+                    let module_path = def.module_path();
+                    let krate = module_path.split("::").next().unwrap();
+                    if !acc.contains_key(krate) {
+                        acc.insert(krate, Default::default());
+                    }
+                    let streams = acc.get_mut(krate).expect("We checked this right above!");
+                    streams.0.insert(module_path);
+                    streams.1.push(stream);
+                    acc
+                },
+            )
+            .into_iter()
+            .sorted_by(|lhs, rhs| lhs.0.cmp(rhs.0))
+            .fold(Vec::new(), |mut acc, (path, (modules, streams))| {
+                let mut modules = Vec::from_iter(modules);
+                modules.sort_unstable();
+
+                acc.push((
+                    crate::output(
+                        format!("crates/{}", path.split("::").next().unwrap()).as_str(),
+                        &filename,
+                    ),
+                    Self::template(
+                        modules,
+                        streams.into_iter().fold(TokenStream::new(), |mut acc, it| {
+                            acc.extend(quote::quote! {
+                                ///@@line_break
+                            });
+                            acc.extend(it);
+                            acc
+                        }),
+                    ),
+                ));
+                acc
+            });
+        Ok(DeriveOutput(output))
+    }
+}
 
 macro_rules! define_derive {
-    ($vis:vis struct $ident:ident $($lifetime:lifetime)? $($rest:tt)*) => {
-        $vis struct $ident $($lifetime)? $($rest)*
+    ($ident:ident $($lifetime:lifetime)?) => {
+        const _: () = {
+            use $crate::{
+                codegen::{CodegenBase, LateCtx, Runner},
+                derives::DeriveOutput,
+                Result,
+            };
 
-        impl $($lifetime)? $crate::derives::DeriveTemplate for $ident $($lifetime)? {
-            fn template(module_paths: Vec<&str>, impls: TokenStream) -> TokenStream {
-                use itertools::Itertools;
-                let header = $crate::codegen::generated_header!();
-                let prelude = Self::prelude();
-
-                // from `x::y::z` to `crate::y::z::*`
-                let use_modules = module_paths.into_iter().map(|it| {
-                    let local_path = ["crate"]
-                        .into_iter()
-                        .chain(it.strip_suffix("::mod").unwrap_or(it).split("::").skip(1))
-                        .chain(["*"])
-                        .join("::");
-                    let use_module: syn::ItemUse =
-                        syn::parse_str(format!("use {local_path};").as_str()).unwrap();
-                    quote::quote! {
-                        ///@@line_break
-                        #use_module
-                    }
-                });
-
-                quote::quote! {
-                    #header
-
-                    #prelude
-
-                    #(#use_modules)*
-
-                    ///@@line_break
-                    #impls
+            impl $($lifetime)? CodegenBase for $ident $($lifetime)? {
+                fn file_path() -> &'static str {
+                    file!()
                 }
             }
-        }
 
-        impl $($lifetime)? $crate::codegen::Runner for $ident $($lifetime)? {
-            type Context = $crate::codegen::LateCtx;
-            type Output = $crate::derives::DeriveOutput;
+            impl $($lifetime)? Runner for $ident $($lifetime)? {
+                type Context = LateCtx;
+                type Output = DeriveOutput;
 
-            fn name(&self) -> &'static str {
-                stringify!($ident)
+                fn name(&self) -> &'static str {
+                    stringify!($ident)
+                }
+
+                fn run(&mut self, ctx: &LateCtx) -> Result<DeriveOutput> {
+                    self.output(ctx)
+                }
             }
-
-            fn run(&mut self, ctx: &$crate::codegen::LateCtx) -> $crate::Result<Self::Output> {
-                use std::vec::Vec;
-                use itertools::Itertools;
-                use rustc_hash::{FxHashMap, FxHashSet};
-
-                use $crate::derives::DeriveTemplate;
-
-                let trait_name = Self::trait_name();
-                let filename = format!("derive_{}.rs", Self::snake_name());
-                let output = ctx
-                    .schema()
-                    .into_iter()
-                    .filter(|def| def.generates_derive(trait_name))
-                    .map(|def| (def, self.derive(def, ctx)))
-                    .fold(FxHashMap::<&str, (FxHashSet<&str>, Vec<TokenStream>)>::default(), |mut acc, (def, stream)| {
-                        let module_path = def.module_path();
-                        let krate = module_path.split("::").next().unwrap();
-                        if !acc.contains_key(krate) {
-                            acc.insert(krate, Default::default());
-                        }
-                        let streams = acc.get_mut(krate).expect("We checked this right above!");
-                        streams.0.insert(module_path);
-                        streams.1.push(stream);
-                        acc
-                    })
-                    .into_iter()
-                    .sorted_by(|lhs, rhs| lhs.0.cmp(rhs.0))
-                    .fold(Vec::new(), |mut acc, (path, (modules, streams))| {
-                        let mut modules = Vec::from_iter(modules);
-                        modules.sort();
-
-                        acc.push((
-                            $crate::output(
-                                format!("crates/{}", path.split("::").next().unwrap()).as_str(),
-                                &filename,
-                            ),
-                            Self::template(
-                                modules,
-                                streams
-                                .into_iter()
-                                .fold(TokenStream::new(), |mut acc, it| {
-                                    acc.extend(quote::quote!{
-                                        ///@@line_break
-                                    });
-                                    acc.extend(it);
-                                    acc
-                                })
-                            )
-                        ));
-                        acc
-                    });
-                Ok(DeriveOutput(output))
-            }
-        }
+        };
     };
 }
 pub(crate) use define_derive;

tasks/ast_tools/src/generators/assert_layouts.rs

@@ -4,16 +4,16 @@ use syn::Type;
 
 use super::define_generator;
 use crate::{
-    codegen::{generated_header, LateCtx},
+    codegen::LateCtx,
     output,
     schema::{FieldDef, ToType, TypeDef},
     util::ToIdent,
     Generator, GeneratorOutput,
 };
 
-define_generator! {
-    pub struct AssertLayouts;
-}
+pub struct AssertLayouts;
+
+define_generator!(AssertLayouts);
 
 impl Generator for AssertLayouts {
     fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
@@ -26,13 +26,9 @@ impl Generator for AssertLayouts {
             })
             .collect::<(Vec<TokenStream>, Vec<TokenStream>)>();
 
-        let header = generated_header!();
-
         GeneratorOutput::Rust {
             path: output(crate::AST_CRATE, "assert_layouts.rs"),
             tokens: quote! {
-                #header
-
                 use std::mem::{align_of, offset_of, size_of};
 
                 ///@@line_break

tasks/ast_tools/src/generators/ast_builder.rs

@@ -10,7 +10,7 @@ use syn::{parse_quote, Ident, Type};
 
 use super::define_generator;
 use crate::{
-    codegen::{generated_header, LateCtx},
+    codegen::LateCtx,
     output,
     schema::{
         EnumDef, FieldDef, GetIdent, InheritDef, StructDef, ToType, TypeDef, TypeName, VariantDef,
@@ -19,9 +19,9 @@ use crate::{
     Generator, GeneratorOutput,
 };
 
-define_generator! {
-    pub struct AstBuilderGenerator;
-}
+pub struct AstBuilderGenerator;
+
+define_generator!(AstBuilderGenerator);
 
 impl Generator for AstBuilderGenerator {
     fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
@@ -32,14 +32,12 @@ impl Generator for AstBuilderGenerator {
             .map(|it| generate_builder_fn(it, ctx))
             .collect_vec();
 
-        let header = generated_header!();
-
         GeneratorOutput::Rust {
             path: output(crate::AST_CRATE, "ast_builder.rs"),
             tokens: quote! {
                 //! AST node factories
-                #header
 
+                //!@@line_break
                 #![allow(
                     clippy::default_trait_access,
                     clippy::too_many_arguments,

tasks/ast_tools/src/generators/ast_kind.rs

@@ -5,15 +5,15 @@ use syn::{parse_quote, Arm, ImplItemFn, Variant};
 
 use super::define_generator;
 use crate::{
-    codegen::{generated_header, LateCtx},
+    codegen::LateCtx,
     output,
     schema::{GetIdent, ToType},
     Generator, GeneratorOutput,
 };
 
-define_generator! {
-    pub struct AstKindGenerator;
-}
+pub struct AstKindGenerator;
+
+define_generator!(AstKindGenerator);
 
 pub const BLACK_LIST: [&str; 61] = [
     "Expression",
@@ -125,15 +125,12 @@ impl Generator for AstKindGenerator {
             })
             .collect_vec();
 
-        let header = generated_header!();
-
         GeneratorOutput::Rust {
             path: output(crate::AST_CRATE, "ast_kind.rs"),
             tokens: quote! {
-                #header
                 #![allow(missing_docs)] ///@ FIXME (in ast_tools/src/generators/ast_kind.rs)
-                ///@@line_break
 
+                ///@@line_break
                 use oxc_span::{GetSpan, Span};
 
                 ///@@line_break

tasks/ast_tools/src/generators/mod.rs

@@ -1,8 +1,12 @@
 use std::path::PathBuf;
 
 use proc_macro2::TokenStream;
+use quote::quote;
 
-use crate::codegen::LateCtx;
+use crate::{
+    codegen::{generate_header, CodegenBase, LateCtx},
+    Result,
+};
 
 mod assert_layouts;
 mod ast_builder;
@@ -16,31 +20,62 @@ pub use ast_kind::AstKindGenerator;
 pub use typescript::TypescriptGenerator;
 pub use visit::{VisitGenerator, VisitMutGenerator};
 
-pub trait Generator {
-    fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput;
-}
-
 #[derive(Debug, Clone)]
 pub enum GeneratorOutput {
     Rust { path: PathBuf, tokens: TokenStream },
     Text { path: PathBuf, content: String },
 }
 
-macro_rules! define_generator {
-    ($vis:vis struct $ident:ident $($lifetime:lifetime)? $($rest:tt)*) => {
-        $vis struct $ident $($lifetime)? $($rest)*
-        impl $($lifetime)? $crate::codegen::Runner for $ident $($lifetime)? {
-            type Context = $crate::codegen::LateCtx;
-            type Output = $crate::GeneratorOutput;
+pub trait Generator: CodegenBase {
+    // Methods defined by implementer
 
-            fn name(&self) -> &'static str {
-                stringify!($ident)
-            }
+    fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput;
 
-            fn run(&mut self, ctx: &$crate::codegen::LateCtx) -> $crate::Result<Self::Output> {
-                Ok(self.generate(ctx))
-            }
+    // Standard methods
+
+    fn output(&mut self, ctx: &LateCtx) -> Result<GeneratorOutput> {
+        let mut output = self.generate(ctx);
+
+        if let GeneratorOutput::Rust { tokens, .. } = &mut output {
+            let header = generate_header(Self::file_path());
+            *tokens = quote! {
+                #header
+                #tokens
+            };
         }
+
+        Ok(output)
+    }
+}
+
+macro_rules! define_generator {
+    ($ident:ident $($lifetime:lifetime)?) => {
+        const _: () = {
+            use $crate::{
+                codegen::{CodegenBase, LateCtx, Runner},
+                generators::GeneratorOutput,
+                Result,
+            };
+
+            impl $($lifetime)? CodegenBase for $ident $($lifetime)? {
+                fn file_path() -> &'static str {
+                    file!()
+                }
+            }
+
+            impl $($lifetime)? Runner for $ident $($lifetime)? {
+                type Context = LateCtx;
+                type Output = GeneratorOutput;
+
+                fn name(&self) -> &'static str {
+                    stringify!($ident)
+                }
+
+                fn run(&mut self, ctx: &LateCtx) -> Result<GeneratorOutput> {
+                    self.output(ctx)
+                }
+            }
+        };
     };
 }
 pub(crate) use define_generator;

tasks/ast_tools/src/generators/typescript.rs

@@ -18,9 +18,9 @@ use crate::{
 
 const CUSTOM_TYPESCRIPT: &str = include_str!("../../../../crates/oxc_ast/src/ast/types.d.ts");
 
-define_generator! {
-    pub struct TypescriptGenerator;
-}
+pub struct TypescriptGenerator;
+
+define_generator!(TypescriptGenerator);
 
 impl Generator for TypescriptGenerator {
     fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {

tasks/ast_tools/src/generators/visit.rs

@@ -9,7 +9,7 @@ use syn::{parse_quote, Ident};
 
 use super::define_generator;
 use crate::{
-    codegen::{generated_header, LateCtx},
+    codegen::LateCtx,
     generators::ast_kind::BLACK_LIST as KIND_BLACK_LIST,
     markers::VisitArg,
     output,
@@ -18,13 +18,9 @@ use crate::{
     Generator, GeneratorOutput,
 };
 
-define_generator! {
-    pub struct VisitGenerator;
-}
+pub struct VisitGenerator;
 
-define_generator! {
-    pub struct VisitMutGenerator;
-}
+define_generator!(VisitGenerator);
 
 impl Generator for VisitGenerator {
     fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
@@ -35,6 +31,10 @@ impl Generator for VisitGenerator {
     }
 }
 
+pub struct VisitMutGenerator;
+
+define_generator!(VisitMutGenerator);
+
 impl Generator for VisitMutGenerator {
     fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
         GeneratorOutput::Rust {
@@ -45,8 +45,6 @@ impl Generator for VisitMutGenerator {
 }
 
 fn generate_visit<const MUT: bool>(ctx: &LateCtx) -> TokenStream {
-    let header = generated_header!();
-
     let (visits, walks) = VisitBuilder::new(ctx, MUT).build();
 
     let walk_mod = if MUT { quote!(walk_mut) } else { quote!(walk) };
@@ -72,8 +70,6 @@ fn generate_visit<const MUT: bool>(ctx: &LateCtx) -> TokenStream {
     };
 
     quote! {
-        #header
-
         //! Visitor Pattern
         //!
         //! See: