1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81
use quote::quote; use syn::{DeriveInput, Data, Fields}; use quote::format_ident; #[proc_macro_derive(EnumConstValue)] pub fn enum_const_value(input: proc_macro::TokenStream) -> proc_macro::TokenStream { let derive_input: DeriveInput = syn::parse(input).unwrap(); let original_enum_type_name = &derive_input.ident; let const_enum_type_ident = format_ident!("{}ConstValue", &original_enum_type_name); let mut generate_const_enum = false; let (idents_and_matcher, original_matcher): (Vec<_>, Vec<_>) = match derive_input.data { Data::Enum(e) => { e .variants .into_iter() .enumerate() .map( |(index, variant) | { let index = index as i32; let variant_name = variant.ident; let tokens = match variant.fields { Fields::Unit => quote! {}, Fields::Unnamed(..) => { generate_const_enum = true; quote! { (..) } } Fields::Named(..) => { generate_const_enum = true; quote! { {..} } } }; let original_enum_matcher = quote! { &#original_enum_type_name::#variant_name#tokens => #index }; let const_enum_matcher = quote! { &#const_enum_type_ident::#variant_name => #index }; ((variant_name, const_enum_matcher), original_enum_matcher) }) .unzip() } _ => panic!("Only enums are supported") }; let (variant_names, const_matcher): (Vec<_>, Vec<_>) = idents_and_matcher.into_iter().unzip(); let mut tokens = quote! { impl #original_enum_type_name { pub fn const_value(&self) -> i32 { match self { #(#original_matcher),* } } } }; if generate_const_enum { tokens.extend(quote! { pub enum #const_enum_type_ident { #(#variant_names),* } impl #const_enum_type_ident { pub fn value_for_variant(variant: &#const_enum_type_ident) -> i32 { match variant { #(#const_matcher),* } } pub fn const_value(&self) -> i32 { #const_enum_type_ident::value_for_variant(self) } } }); } tokens.into() }