1use std::collections::HashMap;
2
3use proc_macro2::TokenStream;
4use quote::quote;
5use syn::{Data, DeriveInput, Fields, Lit};
6
7use crate::helpers::{non_enum_error, HasStrumVariantProperties, HasTypeProperties};
8
9#[derive(Hash, PartialEq, Eq)]
10enum PropertyType {
11 String,
12 Integer,
13 Bool,
14}
15
16const PROPERTY_TYPES: [PropertyType; 3] = [
17 PropertyType::String,
18 PropertyType::Integer,
19 PropertyType::Bool,
20];
21
22pub fn enum_properties_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
23 let name = &ast.ident;
24 let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
25 let variants = match &ast.data {
26 Data::Enum(v) => &v.variants,
27 _ => return Err(non_enum_error()),
28 };
29 let type_properties = ast.get_type_properties()?;
30 let strum_module_path = type_properties.crate_module_path();
31
32 let mut built_arms: HashMap<_, _> = PROPERTY_TYPES.iter().map(|p| (p, Vec::new())).collect();
33
34 for variant in variants {
35 let ident = &variant.ident;
36 let variant_properties = variant.get_variant_properties()?;
37 let mut arms: HashMap<_, _> = PROPERTY_TYPES.iter().map(|p| (p, Vec::new())).collect();
38 // But you can disable the messages.
39 if variant_properties.disabled.is_some() {
40 continue;
41 }
42
43 let params = match variant.fields {
44 Fields::Unit => quote! {},
45 Fields::Unnamed(..) => quote! { (..) },
46 Fields::Named(..) => quote! { {..} },
47 };
48
49 for (key, value) in variant_properties.props {
50 let property_type = match value {
51 Lit::Str(..) => PropertyType::String,
52 Lit::Bool(..) => PropertyType::Bool,
53 Lit::Int(..) => PropertyType::Integer,
54 _ => todo!("TODO"),
55 };
56
57 arms.get_mut(&property_type)
58 .unwrap()
59 .push(quote! { #key => ::core::option::Option::Some( #value )});
60 }
61
62 for property in &PROPERTY_TYPES {
63 arms.get_mut(&property)
64 .unwrap()
65 .push(quote! { _ => ::core::option::Option::None });
66 let arms_as_string = &arms[property];
67 built_arms.get_mut(&property).unwrap().push(quote! {
68 &#name::#ident #params => {
69 match prop {
70 #(#arms_as_string),*
71 }
72 }
73 });
74 }
75 }
76
77 for (_, arms) in built_arms.iter_mut() {
78 if arms.len() < variants.len() {
79 arms.push(quote! { _ => ::core::option::Option::None });
80 }
81 }
82
83 let (built_string_arms, built_int_arms, built_bool_arms) = (
84 &built_arms[&PropertyType::String],
85 &built_arms[&PropertyType::Integer],
86 &built_arms[&PropertyType::Bool],
87 );
88
89 Ok(quote! {
90 impl #impl_generics #strum_module_path::EnumProperty for #name #ty_generics #where_clause {
91 #[inline]
92 fn get_str(&self, prop: &str) -> ::core::option::Option<&'static str> {
93 match self {
94 #(#built_string_arms),*
95 }
96 }
97
98 #[inline]
99 fn get_int(&self, prop: &str) -> ::core::option::Option<i64> {
100 match self {
101 #(#built_int_arms),*
102 }
103 }
104
105 #[inline]
106 fn get_bool(&self, prop: &str) -> ::core::option::Option<bool> {
107 match self {
108 #(#built_bool_arms),*
109 }
110 }
111
112 }
113 })
114}
115