1use proc_macro2::{Span, TokenStream, TokenTree};
2use quote::{quote, ToTokens};
3use syn::parse_quote;
4use syn::{Data, DeriveInput, Fields};
5
6use crate::helpers::{non_enum_error, strum_discriminants_passthrough_error, HasTypeProperties};
7
8/// Attributes to copy from the main enum's variants to the discriminant enum's variants.
9///
10/// Attributes not in this list may be for other `proc_macro`s on the main enum, and may cause
11/// compilation problems when copied across.
12const ATTRIBUTES_TO_COPY: &[&str] = &["doc", "cfg", "allow", "deny", "strum_discriminants"];
13
14pub fn enum_discriminants_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
15 let name = &ast.ident;
16 let vis = &ast.vis;
17
18 let variants = match &ast.data {
19 Data::Enum(v) => &v.variants,
20 _ => return Err(non_enum_error()),
21 };
22
23 // Derives for the generated enum
24 let type_properties = ast.get_type_properties()?;
25
26 let derives = type_properties.discriminant_derives;
27
28 let derives = quote! {
29 #[derive(Clone, Copy, Debug, PartialEq, Eq, #(#derives),*)]
30 };
31
32 // Work out the name
33 let default_name = syn::Ident::new(&format!("{}Discriminants", name), Span::call_site());
34
35 let discriminants_name = type_properties.discriminant_name.unwrap_or(default_name);
36 let discriminants_vis = type_properties
37 .discriminant_vis
38 .unwrap_or_else(|| vis.clone());
39
40 // Pass through all other attributes
41 let pass_though_attributes = type_properties.discriminant_others;
42
43 // Add the variants without fields, but exclude the `strum` meta item
44 let mut discriminants = Vec::new();
45 for variant in variants {
46 let ident = &variant.ident;
47
48 // Don't copy across the "strum" meta attribute. Only passthrough the whitelisted
49 // attributes and proxy `#[strum_discriminants(...)]` attributes
50 let attrs = variant
51 .attrs
52 .iter()
53 .filter(|attr| {
54 ATTRIBUTES_TO_COPY
55 .iter()
56 .any(|attr_whitelisted| attr.path.is_ident(attr_whitelisted))
57 })
58 .map(|attr| {
59 if attr.path.is_ident("strum_discriminants") {
60 let passthrough_group = attr
61 .tokens
62 .clone()
63 .into_iter()
64 .next()
65 .ok_or_else(|| strum_discriminants_passthrough_error(attr))?;
66 let passthrough_attribute = match passthrough_group {
67 TokenTree::Group(ref group) => group.stream(),
68 _ => {
69 return Err(strum_discriminants_passthrough_error(&passthrough_group));
70 }
71 };
72 if passthrough_attribute.is_empty() {
73 return Err(strum_discriminants_passthrough_error(&passthrough_group));
74 }
75 Ok(quote! { #[#passthrough_attribute] })
76 } else {
77 Ok(attr.to_token_stream())
78 }
79 })
80 .collect::<Result<Vec<_>, _>>()?;
81
82 discriminants.push(quote! { #(#attrs)* #ident });
83 }
84
85 // Ideally:
86 //
87 // * For `Copy` types, we `impl From<TheEnum> for TheEnumDiscriminants`
88 // * For `!Copy` types, we `impl<'enum> From<&'enum TheEnum> for TheEnumDiscriminants`
89 //
90 // That way we ensure users are not able to pass a `Copy` type by reference. However, the
91 // `#[derive(..)]` attributes are not in the parsed tokens, so we are not able to check if a
92 // type is `Copy`, so we just implement both.
93 //
94 // See <https://github.com/dtolnay/syn/issues/433>
95 // ---
96 // let is_copy = unique_meta_list(type_meta.iter(), "derive")
97 // .map(extract_list_metas)
98 // .map(|metas| {
99 // metas
100 // .filter_map(get_meta_ident)
101 // .any(|derive| derive.to_string() == "Copy")
102 // }).unwrap_or(false);
103
104 let arms = variants
105 .iter()
106 .map(|variant| {
107 let ident = &variant.ident;
108 let params = match &variant.fields {
109 Fields::Unit => quote! {},
110 Fields::Unnamed(_fields) => {
111 quote! { (..) }
112 }
113 Fields::Named(_fields) => {
114 quote! { { .. } }
115 }
116 };
117
118 quote! { #name::#ident #params => #discriminants_name::#ident }
119 })
120 .collect::<Vec<_>>();
121
122 let from_fn_body = quote! { match val { #(#arms),* } };
123
124 let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
125 let impl_from = quote! {
126 impl #impl_generics ::core::convert::From< #name #ty_generics > for #discriminants_name #where_clause {
127 fn from(val: #name #ty_generics) -> #discriminants_name {
128 #from_fn_body
129 }
130 }
131 };
132 let impl_from_ref = {
133 let mut generics = ast.generics.clone();
134
135 let lifetime = parse_quote!('_enum);
136 let enum_life = quote! { & #lifetime };
137 generics.params.push(lifetime);
138
139 // Shadows the earlier `impl_generics`
140 let (impl_generics, _, _) = generics.split_for_impl();
141
142 quote! {
143 impl #impl_generics ::core::convert::From< #enum_life #name #ty_generics > for #discriminants_name #where_clause {
144 fn from(val: #enum_life #name #ty_generics) -> #discriminants_name {
145 #from_fn_body
146 }
147 }
148 }
149 };
150
151 Ok(quote! {
152 /// Auto-generated discriminant enum variants
153 #derives
154 #(#[ #pass_though_attributes ])*
155 #discriminants_vis enum #discriminants_name {
156 #(#discriminants),*
157 }
158
159 #impl_from
160 #impl_from_ref
161 })
162}
163