enumly_derive/
lib.rs

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/// Derive macro that exposes compile-time constants for the full set of enum variants.
9///
10/// ---
11/// # Examples
12/// ```ignore
13/// use enumly::Enumly;
14///
15/// #[derive(Enumly, Debug, PartialEq)]
16/// enum Color {
17///     Red,
18///     Green,
19///     Blue,
20/// }
21///
22/// assert_eq!(Color::COUNT, 3);
23/// assert_eq!(Color::VARIANTS, &[Color::Red, Color::Green, Color::Blue]);
24/// ```
25///
26/// ---
27/// Fails to compile when any variant is not unit:
28/// ```compile_fail
29/// use enumly::Enumly;
30///
31/// #[derive(Enumly)]
32/// enum Bad {
33///     Tuple(u8),
34///     Struct { value: u8 },
35/// }
36/// ```
37///
38#[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}