1use proc_macro2::TokenStream;
2use quote::quote;
3use syn::{parse_quote, Data, DeriveInput, Fields};
4
5use crate::helpers::{
6 non_enum_error, non_single_field_variant_error, HasStrumVariantProperties, HasTypeProperties,
7};
8
9fn get_arms<F>(ast: &DeriveInput, transparent_fn: F) -> syn::Result<Vec<TokenStream>>
10where
11 F: Fn(&TokenStream) -> TokenStream,
12{
13 let name = &ast.ident;
14 let mut arms = Vec::new();
15 let variants = match &ast.data {
16 Data::Enum(v) => &v.variants,
17 _ => return Err(non_enum_error()),
18 };
19
20 let type_properties = ast.get_type_properties()?;
21
22 for variant in variants {
23 let ident = &variant.ident;
24 let variant_properties = variant.get_variant_properties()?;
25
26 if variant_properties.disabled.is_some() {
27 continue;
28 }
29
30 if let Some(..) = variant_properties.transparent {
31 let arm = super::extract_single_field_variant_and_then(name, variant, |tok| {
32 transparent_fn(tok)
33 })
34 .map_err(|_| non_single_field_variant_error("transparent"))?;
35
36 arms.push(arm);
37 continue;
38 }
39
40 // Look at all the serialize attributes.
41 // Use `to_string` attribute (not `as_ref_str` or something) to keep things consistent
42 // (i.e. always `enum.as_ref().to_string() == enum.to_string()`).
43 let output = variant_properties
44 .get_preferred_name(type_properties.case_style, type_properties.prefix.as_ref());
45 let params = match variant.fields {
46 Fields::Unit => quote! {},
47 Fields::Unnamed(..) => quote! { (..) },
48 Fields::Named(..) => quote! { {..} },
49 };
50
51 arms.push(quote! { #name::#ident #params => #output });
52 }
53
54 if arms.len() < variants.len() {
55 arms.push(quote! {
56 _ => panic!(
57 "AsRef::<str>::as_ref() or AsStaticRef::<str>::as_static() \
58 called on disabled variant.",
59 )
60 });
61 }
62
63 Ok(arms)
64}
65
66pub fn as_ref_str_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
67 let name: &Ident = &ast.ident;
68 let (impl_generics: ImplGenerics<'_>, ty_generics: TypeGenerics<'_>, where_clause: Option<&WhereClause>) = ast.generics.split_for_impl();
69 let arms: Vec = get_arms(ast, |tok: &TokenStream| {
70 quote! { ::core::convert::AsRef::<str>::as_ref(#tok) }
71 })?;
72
73 Ok(quote! {
74 impl #impl_generics ::core::convert::AsRef<str> for #name #ty_generics #where_clause {
75 #[inline]
76 fn as_ref(&self) -> &str {
77 match *self {
78 #(#arms),*
79 }
80 }
81 }
82 })
83}
84
85pub enum GenerateTraitVariant {
86 AsStaticStr,
87 From,
88}
89
90pub fn as_static_str_inner(
91 ast: &DeriveInput,
92 trait_variant: &GenerateTraitVariant,
93) -> syn::Result<TokenStream> {
94 let name = &ast.ident;
95 let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
96 let arms = &get_arms(ast, |tok| {
97 quote! { ::core::convert::From::from(#tok) }
98 })?;
99
100 let type_properties = ast.get_type_properties()?;
101 let strum_module_path = type_properties.crate_module_path();
102
103 let mut generics = ast.generics.clone();
104 generics
105 .params
106 .push(syn::GenericParam::Lifetime(syn::LifetimeParam::new(
107 parse_quote!('_derivative_strum),
108 )));
109 let (impl_generics2, _, _) = generics.split_for_impl();
110
111 Ok(match trait_variant {
112 GenerateTraitVariant::AsStaticStr => quote! {
113 impl #impl_generics #strum_module_path::AsStaticRef<str> for #name #ty_generics #where_clause {
114 #[inline]
115 fn as_static(&self) -> &'static str {
116 match *self {
117 #(#arms),*
118 }
119 }
120 }
121 },
122 GenerateTraitVariant::From if !type_properties.const_into_str => quote! {
123 impl #impl_generics ::core::convert::From<#name #ty_generics> for &'static str #where_clause {
124 #[inline]
125 fn from(x: #name #ty_generics) -> &'static str {
126 match x {
127 #(#arms),*
128 }
129 }
130 }
131 impl #impl_generics2 ::core::convert::From<&'_derivative_strum #name #ty_generics> for &'static str #where_clause {
132 #[inline]
133 fn from(x: &'_derivative_strum #name #ty_generics) -> &'static str {
134 match *x {
135 #(#arms),*
136 }
137 }
138 }
139 },
140 GenerateTraitVariant::From => quote! {
141 impl #impl_generics #name #ty_generics #where_clause {
142 pub const fn into_str(&self) -> &'static str {
143 match self {
144 #(#arms),*
145 }
146 }
147 }
148
149 impl #impl_generics ::core::convert::From<#name #ty_generics> for &'static str #where_clause {
150 fn from(x: #name #ty_generics) -> &'static str {
151 match x {
152 #(#arms),*
153 }
154 }
155 }
156 impl #impl_generics2 ::core::convert::From<&'_derivative_strum #name #ty_generics> for &'static str #where_clause {
157 fn from(x: &'_derivative_strum #name #ty_generics) -> &'static str {
158 x.into_str()
159 }
160 }
161 },
162 })
163}
164