1#![doc = "Provides a procedural macro that exposes a compile-time static list of all variants of an enum."]
2
3use proc_macro::TokenStream;
4use quote::quote;
5use syn::spanned::Spanned;
6use syn::{Attribute, Data, DeriveInput, Fields, parse_macro_input};
7
8#[proc_macro_derive(Enumly)]
39pub fn derive_enumly(input: TokenStream) -> TokenStream {
40 let input = parse_macro_input!(input as DeriveInput);
41
42 if let Some(err) = non_exhaustive_error(&input.attrs) {
43 return err.to_compile_error().into();
44 }
45
46 let data_enum = match input.data {
47 Data::Enum(data_enum) => data_enum,
48 _ => {
49 return syn::Error::new(input.ident.span(), "Enumly can only be derived for enums")
50 .to_compile_error()
51 .into();
52 }
53 };
54
55 let mut variant_idents = Vec::with_capacity(data_enum.variants.len());
56
57 for variant in data_enum.variants {
58 if let Some(err) = non_exhaustive_error(&variant.attrs) {
59 return err.to_compile_error().into();
60 }
61
62 match variant.fields {
63 Fields::Unit => variant_idents.push(variant.ident),
64 _ => {
65 return syn::Error::new(
66 variant.ident.span(),
67 "Enumly only supports unit variants; tuple and struct variants are not allowed",
68 )
69 .to_compile_error()
70 .into();
71 }
72 }
73 }
74
75 let name = &input.ident;
76 let count = variant_idents.len();
77 let variant_exprs = variant_idents
78 .iter()
79 .map(|variant| quote! { Self::#variant });
80 let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
81
82 let expanded = quote! {
83 impl #impl_generics #name #ty_generics #where_clause {
84 pub const COUNT: usize = #count;
85 pub const VARIANTS: &'static [Self] = &[#(#variant_exprs),*];
86 }
87 };
88
89 TokenStream::from(expanded)
90}
91
92fn non_exhaustive_error(attrs: &[Attribute]) -> Option<syn::Error> {
93 attrs
94 .iter()
95 .find(|attr| attr.path().is_ident("non_exhaustive"))
96 .map(|attr| {
97 syn::Error::new(
98 attr.span(),
99 "Enumly does not support #[non_exhaustive] enums or variants",
100 )
101 })
102}