Add support for #[enumcapsulate(discriminant(…))] enum/variant-level helper attributes

regexident · regexident · commit 26b95df12fb4 · 2024-12-01T00:08:21.000+01:00
diff --git a/macros/src/config/for_enum.rs b/macros/src/config/for_enum.rs
@@ -1,9 +1,12 @@
+mod discriminant;
 mod encapsulate;
 mod standard;
+mod variant_discriminant;
 
 pub(crate) use self::encapsulate::EncapsulateDeriveEnumConfig;
+pub(crate) use self::variant_discriminant::VariantDiscriminantDeriveEnumConfig;
 
-use self::standard::EnumConfig;
+use self::{discriminant::DiscriminantConfig, standard::EnumConfig};
 
 pub(crate) type FromDeriveEnumConfig = EnumConfig;
 pub(crate) type TryIntoDeriveEnumConfig = EnumConfig;
diff --git a/macros/src/config/for_enum/discriminant.rs b/macros/src/config/for_enum/discriminant.rs
@@ -0,0 +1,45 @@
+use syn::meta::ParseNestedMeta;
+
+use crate::attr::{NAME, REPR};
+
+#[derive(Clone, Default)]
+pub(crate) struct DiscriminantConfig {
+    repr: Option<syn::Type>,
+    ident: Option<syn::Ident>,
+}
+
+impl DiscriminantConfig {
+    pub(crate) fn parse(
+        &mut self,
+        meta: &ParseNestedMeta,
+        _item_enum: &syn::ItemEnum,
+    ) -> Result<(), syn::Error> {
+        meta.parse_nested_meta(|meta| {
+            if meta.path.is_ident(REPR) {
+                if self.repr.is_some() {
+                    return Err(meta.error("repr already specified"));
+                }
+
+                self.repr = Some(meta.value()?.parse()?);
+            } else if meta.path.is_ident(NAME) {
+                if self.ident.is_some() {
+                    return Err(meta.error("name already specified"));
+                }
+
+                self.ident = Some(meta.value()?.parse()?);
+            } else {
+                return Err(meta.error("unsupported discriminant attribute"));
+            }
+
+            Ok(())
+        })
+    }
+
+    pub(crate) fn repr(&self) -> Option<&syn::Type> {
+        self.repr.as_ref()
+    }
+
+    pub(crate) fn ident(&self) -> Option<&syn::Ident> {
+        self.ident.as_ref()
+    }
+}
diff --git a/macros/src/config/for_enum/variant_discriminant.rs b/macros/src/config/for_enum/variant_discriminant.rs
@@ -0,0 +1,40 @@
+use crate::{attr::DISCRIMINANT, parse_enumcapsulate_attrs};
+
+use super::DiscriminantConfig;
+
+#[derive(Clone, Default)]
+pub(crate) struct VariantDiscriminantDeriveEnumConfig {
+    discriminant: Option<DiscriminantConfig>,
+}
+
+impl VariantDiscriminantDeriveEnumConfig {
+    pub(crate) fn from_enum(item_enum: &syn::ItemEnum) -> Result<Self, syn::Error> {
+        let mut this = Self::default();
+
+        parse_enumcapsulate_attrs(&item_enum.attrs, |meta| {
+            if meta.path.is_ident(DISCRIMINANT) {
+                let mut discriminant = this.discriminant.take().unwrap_or_default();
+
+                discriminant.parse(&meta, item_enum)?;
+
+                this.discriminant = Some(discriminant);
+            }
+
+            Ok(())
+        })?;
+
+        Ok(this)
+    }
+
+    pub fn repr(&self) -> Option<&syn::Type> {
+        self.discriminant
+            .as_ref()
+            .and_then(|discriminant| discriminant.repr())
+    }
+
+    pub fn ident(&self) -> Option<&syn::Ident> {
+        self.discriminant
+            .as_ref()
+            .and_then(|discriminant| discriminant.ident())
+    }
+}
diff --git a/macros/src/config/for_variant.rs b/macros/src/config/for_variant.rs
@@ -1,6 +1,10 @@
+mod discriminant;
 mod standard;
+mod variant_discriminant;
 
-use self::standard::VariantConfig;
+pub(crate) use self::variant_discriminant::VariantDiscriminantDeriveVariantConfig;
+
+use self::{discriminant::DiscriminantConfig, standard::VariantConfig};
 
 pub(crate) type FromDeriveVariantConfig = VariantConfig;
 pub(crate) type TryIntoDeriveVariantConfig = VariantConfig;
diff --git a/macros/src/config/for_variant/discriminant.rs b/macros/src/config/for_variant/discriminant.rs
@@ -0,0 +1,45 @@
+use syn::meta::ParseNestedMeta;
+
+use crate::attr::{NAME, VALUE};
+
+#[derive(Clone, Default)]
+pub(crate) struct DiscriminantConfig {
+    expr: Option<syn::Expr>,
+    ident: Option<syn::Ident>,
+}
+
+impl DiscriminantConfig {
+    pub(crate) fn parse(
+        &mut self,
+        meta: &ParseNestedMeta,
+        _variant: &syn::Variant,
+    ) -> Result<(), syn::Error> {
+        meta.parse_nested_meta(|meta| {
+            if meta.path.is_ident(VALUE) {
+                if self.expr.is_some() {
+                    return Err(meta.error("value already specified"));
+                }
+
+                self.expr = Some(meta.value()?.parse()?);
+            } else if meta.path.is_ident(NAME) {
+                if self.ident.is_some() {
+                    return Err(meta.error("name already specified"));
+                }
+
+                self.ident = Some(meta.value()?.parse()?);
+            } else {
+                return Err(meta.error("unsupported discriminant attribute"));
+            }
+
+            Ok(())
+        })
+    }
+
+    pub(crate) fn expr(&self) -> Option<&syn::Expr> {
+        self.expr.as_ref()
+    }
+
+    pub(crate) fn ident(&self) -> Option<&syn::Ident> {
+        self.ident.as_ref()
+    }
+}
diff --git a/macros/src/config/for_variant/variant_discriminant.rs b/macros/src/config/for_variant/variant_discriminant.rs
@@ -0,0 +1,40 @@
+use crate::{attr::DISCRIMINANT, parse_enumcapsulate_attrs};
+
+use super::DiscriminantConfig;
+
+#[derive(Clone, Default)]
+pub(crate) struct VariantDiscriminantDeriveVariantConfig {
+    discriminant: Option<DiscriminantConfig>,
+}
+
+impl VariantDiscriminantDeriveVariantConfig {
+    pub(crate) fn from_variant(variant: &syn::Variant) -> Result<Self, syn::Error> {
+        let mut this = Self::default();
+
+        parse_enumcapsulate_attrs(&variant.attrs, |meta| {
+            if meta.path.is_ident(DISCRIMINANT) {
+                let mut discriminant = this.discriminant.take().unwrap_or_default();
+
+                discriminant.parse(&meta, variant)?;
+
+                this.discriminant = Some(discriminant);
+            }
+
+            Ok(())
+        })?;
+
+        Ok(this)
+    }
+
+    pub fn expr(&self) -> Option<&syn::Expr> {
+        self.discriminant
+            .as_ref()
+            .and_then(|discriminant| discriminant.expr())
+    }
+
+    pub fn ident(&self) -> Option<&syn::Ident> {
+        self.discriminant
+            .as_ref()
+            .and_then(|discriminant| discriminant.ident())
+    }
+}
diff --git a/macros/src/enum_deriver.rs b/macros/src/enum_deriver.rs
@@ -552,60 +552,83 @@ impl EnumDeriver {
     pub fn derive_variant_discriminant(&self) -> Result<TokenStream2, syn::Error> {
         let enum_ident = &self.item.ident;
 
-        let outer = enum_ident;
-        let outer_ty: Type = parse_quote_spanned! { outer.span() => #outer };
+        let enum_config = VariantDiscriminantDeriveEnumConfig::from_enum(&self.item)?;
+
+        let mut discriminant_enum_ident = quote::format_ident!("{enum_ident}Discriminant");
+        let mut repr_attr: Option<TokenStream2> = None;
+
+        if let Some(ident) = enum_config.ident() {
+            discriminant_enum_ident = ident.clone();
+        }
+
+        if let Some(ty) = enum_config.repr() {
+            repr_attr = Some(quote! {
+                #[repr(#ty)]
+            });
+        }
 
         let (impl_generics, type_generics, where_clause) = self.item.generics.split_for_impl();
 
         let variants = self.variants();
 
-        let discriminant_ident = quote::format_ident!("{outer}Discriminant");
-
         let mut discriminant_variants: Vec<TokenStream2> = vec![];
+        let mut match_arms: Vec<TokenStream2> = vec![];
 
         for variant in &variants {
-            let variant_ident = &variant.ident;
+            let variant_config = VariantDiscriminantDeriveVariantConfig::from_variant(variant)?;
 
-            discriminant_variants.push(quote! {
-                #variant_ident,
-            });
-        }
+            let variant_ident: &syn::Ident = &variant.ident;
 
-        let discriminant_enum = quote! {
-            #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
-            pub enum #discriminant_ident {
-                #(#discriminant_variants)*
+            let mut discriminant_variant_ident: syn::Ident = variant_ident.clone();
+            let mut discriminant_variant_expr: Option<&syn::Expr> =
+                variant.discriminant.as_ref().map(|(_, expr)| expr);
+
+            if let Some(ident) = variant_config.ident() {
+                discriminant_variant_ident = ident.clone();
             }
-        };
 
-        let mut match_arms: Vec<TokenStream2> = vec![];
+            if let Some(expr) = variant_config.expr() {
+                discriminant_variant_expr = Some(expr);
+            }
 
-        for variant in variants {
-            let variant_ident = &variant.ident;
-            let inner = variant_ident;
+            let variant_discriminant = discriminant_variant_expr.map(|expr| {
+                quote! { = #expr }
+            });
+
+            discriminant_variants.push(quote! {
+                #discriminant_variant_ident #variant_discriminant,
+            });
 
             let pattern = match &variant.fields {
                 Fields::Named(_) => quote! {
-                    #outer_ty::#inner { .. }
+                    #enum_ident::#variant_ident { .. }
                 },
                 Fields::Unnamed(_) => quote! {
-                    #outer_ty::#inner(..)
+                    #enum_ident::#variant_ident(..)
                 },
                 Fields::Unit => quote! {
-                    #outer_ty::#inner
+                    #enum_ident::#variant_ident
                 },
             };
 
             match_arms.push(quote! {
-                #pattern => #discriminant_ident::#inner,
+                #pattern => #discriminant_enum_ident::#discriminant_variant_ident,
             });
         }
 
+        let discriminant_enum = quote! {
+            #repr_attr
+            #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
+            pub enum #discriminant_enum_ident {
+                #(#discriminant_variants)*
+            }
+        };
+
         Ok(quote! {
             #discriminant_enum
 
-            impl #impl_generics ::enumcapsulate::VariantDiscriminant for #outer_ty #type_generics #where_clause {
-                type Discriminant = #discriminant_ident;
+            impl #impl_generics ::enumcapsulate::VariantDiscriminant for #enum_ident #type_generics #where_clause {
+                type Discriminant = #discriminant_enum_ident;
 
                 fn variant_discriminant(&self) -> Self::Discriminant {
                     match self {
diff --git a/macros/src/lib.rs b/macros/src/lib.rs
@@ -348,7 +348,7 @@ pub fn derive_variant_downcast(input: TokenStream) -> TokenStream {
 /// }
 /// ```
 ///
-#[proc_macro_derive(VariantDiscriminant)]
+#[proc_macro_derive(VariantDiscriminant, attributes(enumcapsulate))]
 pub fn derive_variant_discriminant(input: TokenStream) -> TokenStream {
     let item = parse_macro_input!(input as syn::ItemEnum);
 
diff --git a/macros/src/utils.rs b/macros/src/utils.rs
@@ -2,9 +2,15 @@ use proc_macro::TokenStream;
 use proc_macro2::TokenStream as TokenStream2;
 
 pub(crate) mod attr {
+    #![allow(dead_code)]
+
+    pub(crate) const DISCRIMINANT: &str = "discriminant";
     pub(crate) const EXCLUDE: &str = "exclude";
     pub(crate) const FIELD: &str = "field";
+    pub(crate) const NAME: &str = "name";
     pub(crate) const NAMESPACE: &str = "enumcapsulate";
+    pub(crate) const REPR: &str = "repr";
+    pub(crate) const VALUE: &str = "value";
 }
 
 pub(crate) mod macro_name {
diff --git a/tests/derive-tests/variant_discriminant/pass/enum/name.out.rs b/tests/derive-tests/variant_discriminant/pass/enum/name.out.rs
diff --git a/tests/derive-tests/variant_discriminant/pass/enum/name.rs b/tests/derive-tests/variant_discriminant/pass/enum/name.rs
diff --git a/tests/derive-tests/variant_discriminant/pass/enum/repr.out.rs b/tests/derive-tests/variant_discriminant/pass/enum/repr.out.rs
diff --git a/tests/derive-tests/variant_discriminant/pass/enum/repr.rs b/tests/derive-tests/variant_discriminant/pass/enum/repr.rs
