refactor(ast_codegen): simplify derive_get_span generator

tasks/ast_codegen/src/generators/derive_get_span.rs

@@ -1,6 +1,6 @@
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
-use syn::{Generics, Ident, Type};
+use syn::Ident;
 
 use crate::{
     output,
@@ -15,86 +15,62 @@ define_generator! {
     pub struct DeriveGetSpan;
 }
 
-define_generator! {
-    pub struct DeriveGetSpanMut;
-}
-
 impl Generator for DeriveGetSpan {
     fn name(&self) -> &'static str {
         stringify!(DeriveGetSpan)
     }
 
     fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
-        GeneratorOutput::Stream((
-            output(crate::AST_CRATE, "derive_get_span.rs"),
-            derive::<false>(ctx),
-        ))
+        let trait_name = format_ident!("GetSpan");
+        let method_name = format_ident!("span");
+        let self_type = quote!(&self);
+        let result_type = quote!(Span);
+        let result_expr = quote!(self.span);
+        let out = derive(&trait_name, &method_name, &self_type, &result_type, &result_expr, ctx);
+
+        GeneratorOutput::Stream((output(crate::AST_CRATE, "derive_get_span.rs"), out))
     }
 }
 
+define_generator! {
+    pub struct DeriveGetSpanMut;
+}
+
 impl Generator for DeriveGetSpanMut {
     fn name(&self) -> &'static str {
         stringify!(DeriveGetSpanMut)
     }
 
     fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
-        GeneratorOutput::Stream((
-            output(crate::AST_CRATE, "derive_get_span_mut.rs"),
-            derive::<true>(ctx),
-        ))
+        let trait_name = format_ident!("GetSpanMut");
+        let method_name = format_ident!("span_mut");
+        let self_type = quote!(&mut self);
+        let result_type = quote!(&mut Span);
+        let result_expr = quote!(&mut self.span);
+        let out = derive(&trait_name, &method_name, &self_type, &result_type, &result_expr, ctx);
+
+        GeneratorOutput::Stream((output(crate::AST_CRATE, "derive_get_span_mut.rs"), out))
     }
 }
 
-fn derive<const MUT: bool>(ctx: &LateCtx) -> TokenStream {
-    let (self_type, trait_ident, method_ident, result_type) = if MUT {
-        (
-            quote!(&mut self),
-            format_ident!("GetSpanMut"),
-            format_ident!("span_mut"),
-            quote!(&mut Span),
-        )
-    } else {
-        (quote!(&self), format_ident!("GetSpan"), format_ident!("span"), quote!(Span))
-    };
-    let inline = quote!(#[inline]);
-
-    let derive_enum = |it: &EnumDef| {
-        let generics = it.generics();
-        let typ = it.to_type();
-        impl_trait(
-            &trait_ident,
-            &method_ident,
-            &generics,
-            &typ,
-            &self_type,
-            &result_type,
-            None,
-            &derive_enum(it, &method_ident),
-        )
-    };
-
-    let derive_struct = |it: &StructDef| {
-        let generics = it.generics();
-        let typ = it.to_type();
-        impl_trait(
-            &trait_ident,
-            &method_ident,
-            &generics,
-            &typ,
-            &self_type,
-            &result_type,
-            Some(&inline),
-            &derive_struct::<MUT>(it, &method_ident),
-        )
-    };
+fn derive(
+    trait_name: &Ident,
+    method_name: &Ident,
+    self_type: &TokenStream,
+    result_type: &TokenStream,
+    result_expr: &TokenStream,
+    ctx: &LateCtx,
+) -> TokenStream {
     let impls: Vec<TokenStream> = ctx
         .schema
         .definitions
         .iter()
         .filter(|def| def.visitable())
         .map(|def| match &def {
-            TypeDef::Enum(it) => derive_enum(it),
-            TypeDef::Struct(it) => derive_struct(it),
+            TypeDef::Enum(def) => derive_enum(def, trait_name, method_name, self_type, result_type),
+            TypeDef::Struct(def) => {
+                derive_struct(def, trait_name, method_name, self_type, result_type, result_expr)
+            }
         })
         .collect();
 
@@ -105,57 +81,68 @@ fn derive<const MUT: bool>(ctx: &LateCtx) -> TokenStream {
         insert!("#![allow(clippy::match_same_arms)]");
         endl!();
 
-        use oxc_span::{#trait_ident, Span};
+        use oxc_span::{#trait_name, Span};
         endl!();
         use crate::ast::*;
+        endl!();
 
         #(#impls)*
     }
 }
 
-fn derive_enum(def: &EnumDef, method: &Ident) -> TokenStream {
+fn derive_enum(
+    def: &EnumDef,
+    trait_name: &Ident,
+    method_name: &Ident,
+    self_type: &TokenStream,
+    result_type: &TokenStream,
+) -> TokenStream {
+    let target_type = def.to_type();
+    let generics = def.generics();
+
     let matches = def.all_variants().map(|var| {
         let ident = var.ident();
-        quote!(Self :: #ident(it) => it.#method())
+        quote!(Self :: #ident(it) => it.#method_name())
     });
 
     quote! {
-        match self {
-            #(#matches),*
+        impl #generics #trait_name for #target_type {
+            fn #method_name(#self_type) -> #result_type {
+                match self {
+                    #(#matches),*
+                }
+            }
         }
+        endl!();
     }
 }
 
-fn derive_struct<const MUT: bool>(def: &StructDef, method: &Ident) -> TokenStream {
-    let inner_span_hint = def.fields.iter().find(|it| it.markers.span);
-    if let Some(span_field) = inner_span_hint {
+fn derive_struct(
+    def: &StructDef,
+    trait_name: &Ident,
+    method_name: &Ident,
+    self_type: &TokenStream,
+    result_type: &TokenStream,
+    result_expr: &TokenStream,
+) -> TokenStream {
+    let target_type = def.to_type();
+    let generics = def.generics();
+
+    let span_field = def.fields.iter().find(|field| field.markers.span);
+    let result_expr = if let Some(span_field) = span_field {
         let ident = span_field.name.as_ref().map(ToIdent::to_ident).unwrap();
-        quote!(self.#ident.#method())
-    } else if MUT {
-        quote!(&mut self.span)
+        quote!(self.#ident.#method_name())
     } else {
-        quote!(self.span)
-    }
-}
+        result_expr.clone()
+    };
 
-#[allow(clippy::too_many_arguments)]
-fn impl_trait(
-    trait_ident: &Ident,
-    method_ident: &Ident,
-    generics: &Option<Generics>,
-    target_type: &Type,
-    self_: &TokenStream,
-    result_type: &TokenStream,
-    inline: Option<&TokenStream>,
-    body: &TokenStream,
-) -> TokenStream {
     quote! {
-        endl!();
-        impl #generics #trait_ident for #target_type {
-            #inline
-            fn #method_ident(#self_) -> #result_type {
-                #body
+        impl #generics #trait_name for #target_type {
+            #[inline]
+            fn #method_name(#self_type) -> #result_type {
+                #result_expr
             }
         }
+        endl!();
     }
 }