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 {
     }
 }
 
+/// Implemented by `define_derive!` and `define_generator!` macros
+pub trait CodegenBase {
+    fn file_path() -> &'static str;
+}
+
 /// Creates a generated file warning + required information for a generated file.
-macro_rules! generated_header {
+pub fn generate_header(file_path: &str) -> TokenStream {
-    () => {{
+    let file_path = file_path.replace('\\', "/");
-        let file = file!().replace("\\", "/");
+
-        // TODO add generation date, AST source hash, etc here.
+    // TODO: Add generation date, AST source hash, etc here.
-        let edit_comment = format!("@ To edit this generated file you have to edit `{file}`");
+    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
     }
-    }};
 }
-
-pub(crate) use generated_header;

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;
-}
+
+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;
-}
+
+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;
-}
+
+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;
-}
+
+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;
-}
+
+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;
-}
+
+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,23 +40,11 @@ pub trait Derive {
     fn prelude() -> TokenStream {
         TokenStream::default()
     }
-}
 
-pub trait DeriveTemplate: Derive {
+    // Standard methods
-    fn template(module_path: Vec<&str>, impls: TokenStream) -> TokenStream;
-}
 
-#[derive(Debug, Clone)]
-pub struct DeriveOutput(pub Vec<(PathBuf, TokenStream)>);
-
-macro_rules! define_derive {
-    ($vis:vis struct $ident:ident $($lifetime:lifetime)? $($rest:tt)*) => {
-        $vis struct $ident $($lifetime)? $($rest)*
-
-        impl $($lifetime)? $crate::derives::DeriveTemplate for $ident $($lifetime)? {
     fn template(module_paths: Vec<&str>, impls: TokenStream) -> TokenStream {
-                use itertools::Itertools;
+        let header = generate_header(Self::file_path());
-                let header = $crate::codegen::generated_header!();
         let prelude = Self::prelude();
 
         // from `x::y::z` to `crate::y::z::*`
@@ -74,23 +73,8 @@ macro_rules! define_derive {
             #impls
         }
     }
-        }
-
-        impl $($lifetime)? $crate::codegen::Runner for $ident $($lifetime)? {
-            type Context = $crate::codegen::LateCtx;
-            type Output = $crate::derives::DeriveOutput;
-
-            fn name(&self) -> &'static str {
-                stringify!($ident)
-            }
-
-            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;
 
+    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
@@ -98,7 +82,9 @@ macro_rules! define_derive {
             .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)| {
+            .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) {
@@ -108,36 +94,64 @@ macro_rules! define_derive {
                     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();
+                modules.sort_unstable();
 
                 acc.push((
-                            $crate::output(
+                    crate::output(
                         format!("crates/{}", path.split("::").next().unwrap()).as_str(),
                         &filename,
                     ),
                     Self::template(
                         modules,
-                                streams
+                        streams.into_iter().fold(TokenStream::new(), |mut acc, it| {
-                                .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 {
+    ($ident:ident $($lifetime:lifetime)?) => {
+        const _: () = {
+            use $crate::{
+                codegen::{CodegenBase, LateCtx, Runner},
+                derives::DeriveOutput,
+                Result,
+            };
+
+            impl $($lifetime)? CodegenBase for $ident $($lifetime)? {
+                fn file_path() -> &'static str {
+                    file!()
+                }
+            }
+
+            impl $($lifetime)? Runner for $ident $($lifetime)? {
+                type Context = LateCtx;
+                type Output = DeriveOutput;
+
+                fn name(&self) -> &'static str {
+                    stringify!($ident)
+                }
+
+                fn run(&mut self, ctx: &LateCtx) -> Result<DeriveOutput> {
+                    self.output(ctx)
+                }
+            }
+        };
     };
 }
 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;
-}
+
+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;
-}
+
+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;
-}
+
+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 },
 }
 
+pub trait Generator: CodegenBase {
+    // Methods defined by implementer
+
+    fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput;
+
+    // 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 {
-    ($vis:vis struct $ident:ident $($lifetime:lifetime)? $($rest:tt)*) => {
+    ($ident:ident $($lifetime:lifetime)?) => {
-        $vis struct $ident $($lifetime)? $($rest)*
+        const _: () = {
-        impl $($lifetime)? $crate::codegen::Runner for $ident $($lifetime)? {
+            use $crate::{
-            type Context = $crate::codegen::LateCtx;
+                codegen::{CodegenBase, LateCtx, Runner},
-            type Output = $crate::GeneratorOutput;
+                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: &$crate::codegen::LateCtx) -> $crate::Result<Self::Output> {
+                fn run(&mut self, ctx: &LateCtx) -> Result<GeneratorOutput> {
-                Ok(self.generate(ctx))
+                    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;
-}
+
+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;
-}
 
-define_generator! {
+define_generator!(VisitGenerator);
-    pub struct VisitMutGenerator;
-}
 
 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: