enum_dict_derive/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::spanned::Spanned;
4use syn::{Data, DeriveInput, Error, Fields, LitStr, parse_macro_input};
5
6#[proc_macro_derive(DictKey)]
7pub fn derive_dict_key(input: TokenStream) -> TokenStream {
8    let input: DeriveInput = parse_macro_input!(input);
9    let ident = &input.ident;
10
11    let mut ident_names = vec![];
12
13    match &input.data {
14        Data::Enum(data) => {
15            for variant in &data.variants {
16                let ident = &variant.ident;
17                let name = LitStr::new(&ident.to_string(), ident.span());
18                ident_names.push(quote! { #name });
19
20                match &variant.fields {
21                    Fields::Unit => {}
22                    _ => {
23                        return Error::new(variant.span(), "DictKey can only be derived for unit variants")
24                            .to_compile_error()
25                            .into();
26                    }
27                }
28            }
29        }
30        _ => {
31            return Error::new(input.span(), "DictKey can only be derived for enums")
32                .to_compile_error()
33                .into();
34        }
35    }
36
37    quote! {
38        impl DictKey for #ident {
39            const FIELDS: &'static [&'static str] = &[#(#ident_names),*];
40
41            fn into_usize(self) -> usize {
42                self as usize
43            }
44        }
45    }
46    .into()
47}