cenum_utils_derive/
lib.rs1use std::{collections::HashSet};
28
29use proc_macro2::{Span, TokenStream};
30
31use quote::quote;
32use rustc_hash::{FxBuildHasher};
33use syn::{parse_macro_input, DeriveInput, Ident, Type};
34
35#[proc_macro_derive(ConstEnum)]
39pub fn derive_const_enum(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
40 impl_const_enum(parse_macro_input!(input as DeriveInput)).unwrap_or_else(|error| error.into_compile_error()).into()
41}
42
43thread_local! {
44 static PRIMITIVE_TYPES: HashSet<Ident, FxBuildHasher> = {
46 let type_tokens = [
47 quote! { u8 },
48 quote! { u16 },
49 quote! { u32 },
50 quote! { u64 },
51 quote! { u128 },
52 quote! { usize },
53 quote! { i8 },
54 quote! { i16 },
55 quote! { i32 },
56 quote! { i64 },
57 quote! { i128 },
58 quote! { isize },
59 ];
60
61 HashSet::from_iter(type_tokens.into_iter().map(|tokens| syn::parse2(tokens).unwrap()))
62 };
63}
64
65fn impl_const_enum(item: DeriveInput) -> Result<TokenStream, syn::Error> {
66 let syn::Data::Enum(enum_data) = item.data else { return Err(syn::Error::new_spanned(item, "ConstEnum: expected an enum")); };
67
68 let representation: Option<Type> = item.attrs.iter().find_map(|attribute| {
70 attribute.meta
71 .path()
72 .get_ident()
73 .and_then(|ident| {
74 (ident == "repr")
75 .then(|| attribute.parse_args().ok())
76 .flatten()
77 })
78 });
79
80 let enum_name: Ident = item.ident;
81 let variant_count: usize = enum_data.variants.len();
82
83 let variant_idents = enum_data.variants.iter().map(|variant| &variant.ident);
84 let variant_names = enum_data.variants.iter().map(|variant| variant.ident.to_string());
85
86 let enum_discriminant: Option<Type> = representation.and_then(|ty| {
88 let Type::Path(type_path) = &ty else { return None; };
89
90 if !type_path.qself.is_none() { return None; };
91 let ident: &Ident = type_path.path.get_ident()?;
92
93 PRIMITIVE_TYPES.with(|primitives| primitives.contains(ident)).then_some(ty)
94 });
95
96 let crate_name: Ident = Ident::new("cenum_utils", Span::call_site());
97
98 let count_impl: TokenStream = quote! {
99 impl ::#crate_name::EnumCount for #enum_name {
100 const COUNT: usize = #variant_count;
101 }
102 };
103
104 let names_impl: TokenStream = quote! {
105 impl ::#crate_name::EnumNames for #enum_name {
106 const NAMES: &[&str] = &[#(#variant_names),*];
107 }
108 };
109
110 let discriminants_impl: Option<TokenStream> = enum_discriminant.map(|discriminant| {
111 quote! {
112 impl ::#crate_name::EnumDiscriminants for #enum_name {
113 type Discriminant = #discriminant;
114
115 const DISCRIMINANTS: &[Self::Discriminant] = &[#(Self::#variant_idents as #discriminant),*];
116 }
117 }
118 });
119
120 Ok(quote! { #count_impl #names_impl #discriminants_impl })
121}