From 760fb6301997bfdd9f1159a5ac52d4ff6262f270 Mon Sep 17 00:00:00 2001
From: rzvxa <rzvxa@protonmail.com>
Date: Wed, 7 Aug 2024 15:29:00 +0330
Subject: [PATCH] feat(ast_codegen): support for `generate_derive` marker.

---
 crates/oxc_ast_macros/src/lib.rs     |  2 +-
 tasks/ast_codegen/src/schema/defs.rs | 11 +++++++++++
 tasks/ast_codegen/src/schema/mod.rs  | 21 +++++++++++++++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/crates/oxc_ast_macros/src/lib.rs b/crates/oxc_ast_macros/src/lib.rs
index 7701e27244d8cb..eff30de0943b97 100644
--- a/crates/oxc_ast_macros/src/lib.rs
+++ b/crates/oxc_ast_macros/src/lib.rs
@@ -49,7 +49,7 @@ pub fn ast(_args: TokenStream, input: TokenStream) -> TokenStream {
 /// Does not generate any code.
 /// Only purpose is to allow using `#[scope]`, `#[visit]`, and other attrs in the AST node type defs.
 /// These "marker" attributes are used in codegen.
-#[proc_macro_derive(Ast, attributes(scope, visit, span, serde, tsify))]
+#[proc_macro_derive(Ast, attributes(scope, visit, span, serde, tsify, generate_derive))]
 pub fn ast_derive(_item: TokenStream) -> TokenStream {
     TokenStream::new()
 }
diff --git a/tasks/ast_codegen/src/schema/defs.rs b/tasks/ast_codegen/src/schema/defs.rs
index b4d965fb3371ff..46bb2926614a12 100644
--- a/tasks/ast_codegen/src/schema/defs.rs
+++ b/tasks/ast_codegen/src/schema/defs.rs
@@ -26,6 +26,15 @@ impl TypeDef {
     pub fn visitable(&self) -> bool {
         with_either!(self, it => it.visitable)
     }
+
+    pub fn generated_derives(&self) -> &Vec<String> {
+        with_either!(self, it => &it.generated_derives)
+    }
+
+    pub fn generates_derive(&self, derive: &str) -> bool {
+        let generated_derives = self.generated_derives();
+        generated_derives.iter().any(|it| it == derive)
+    }
 }
 
 #[derive(Debug, Serialize)]
@@ -41,6 +50,7 @@ pub struct StructDef {
     pub size_32: usize,
     pub align_32: usize,
     pub offsets_32: Option<Vec<usize>>,
+    pub generated_derives: Vec<String>,
     #[serde(skip)]
     pub markers: OuterMarkers,
 }
@@ -60,6 +70,7 @@ pub struct EnumDef {
     pub size_32: usize,
     pub align_32: usize,
     pub offsets_32: Option<Vec<usize>>,
+    pub generated_derives: Vec<String>,
 }
 
 impl EnumDef {
diff --git a/tasks/ast_codegen/src/schema/mod.rs b/tasks/ast_codegen/src/schema/mod.rs
index 57a9ad50c8608d..d588e9a8fb62e7 100644
--- a/tasks/ast_codegen/src/schema/mod.rs
+++ b/tasks/ast_codegen/src/schema/mod.rs
@@ -148,6 +148,8 @@ fn lower_ast_enum(it @ rust::Enum { item, meta }: &rust::Enum, ctx: &crate::Earl
         size_32,
         align_32,
         offsets_32,
+
+        generated_derives: parse_generate_derive(&item.attrs),
     }
 }
 
@@ -178,6 +180,8 @@ fn lower_ast_struct(
         align_32,
         offsets_32,
         markers: parse_outer_markers(&item.attrs).unwrap(),
+
+        generated_derives: parse_generate_derive(&item.attrs),
     }
 }
 
@@ -264,6 +268,23 @@ fn get_docs(attrs: &[syn::Attribute]) -> Vec<String> {
         .collect()
 }
 
+fn parse_generate_derive(attrs: &[syn::Attribute]) -> Vec<String> {
+    let mut derives = std::collections::HashSet::new();
+    for attr in attrs {
+        if !attr.path().is_ident("generate_derive") {
+            continue;
+        }
+
+        let args: syn::punctuated::Punctuated<syn::Ident, syn::Token![,]> =
+            attr.parse_args_with(syn::punctuated::Punctuated::parse_terminated).unwrap();
+
+        for arg in args {
+            derives.insert(arg.to_string());
+        }
+    }
+    Vec::from_iter(derives)
+}
+
 macro_rules! with_either {
     ($def:expr, $it:ident => $body:expr) => {
         match $def {