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 let strum_module_path = type_properties.crate_module_path();
26
27 let derives = type_properties.discriminant_derives;
28
29 let derives = quote! {
30 #[derive(Clone, Copy, Debug, PartialEq, Eq, #(#derives),*)]
31 };
32
33 // Work out the name
34 let default_name = syn::Ident::new(&format!("{}Discriminants", name), Span::call_site());
35
36 let discriminants_name = type_properties.discriminant_name.unwrap_or(default_name);
37 let discriminants_vis = type_properties
38 .discriminant_vis
39 .as_ref()
40 .unwrap_or_else(|| &vis);
41
42 // Pass through all other attributes
43 let pass_though_attributes = type_properties.discriminant_others;
44
45 let repr = type_properties.enum_repr.map(|repr| quote!(#[repr(#repr)]));
46
47 // Add the variants without fields, but exclude the `strum` meta item
48 let mut discriminants = Vec::new();
49 for variant in variants {
50 let ident = &variant.ident;
51 let discriminant = variant
52 .discriminant
53 .as_ref()
54 .map(|(_, expr)| quote!( = #expr));
55
56 // Don't copy across the "strum" meta attribute. Only passthrough the whitelisted
57 // attributes and proxy `#[strum_discriminants(...)]` attributes
58 let attrs = variant
59 .attrs
60 .iter()
61 .filter(|attr| {
62 ATTRIBUTES_TO_COPY
63 .iter()
64 .any(|attr_whitelisted| attr.path().is_ident(attr_whitelisted))
65 })
66 .map(|attr| {
67 if attr.path().is_ident("strum_discriminants") {
68 let mut ts = attr.meta.require_list()?.to_token_stream().into_iter();
69
70 // Discard strum_discriminants(...)
71 let _ = ts.next();
72
73 let passthrough_group = ts
74 .next()
75 .ok_or_else(|| strum_discriminants_passthrough_error(attr))?;
76 let passthrough_attribute = match passthrough_group {
77 TokenTree::Group(ref group) => group.stream(),
78 _ => {
79 return Err(strum_discriminants_passthrough_error(&passthrough_group));
80 }
81 };
82 if passthrough_attribute.is_empty() {
83 return Err(strum_discriminants_passthrough_error(&passthrough_group));
84 }
85 Ok(quote! { #[#passthrough_attribute] })
86 } else {
87 Ok(attr.to_token_stream())
88 }
89 })
90 .collect::<Result<Vec<_>, _>>()?;
91
92 discriminants.push(quote! { #(#attrs)* #ident #discriminant});
93 }
94
95 // Ideally:
96 //
97 // * For `Copy` types, we `impl From<TheEnum> for TheEnumDiscriminants`
98 // * For `!Copy` types, we `impl<'enum> From<&'enum TheEnum> for TheEnumDiscriminants`
99 //
100 // That way we ensure users are not able to pass a `Copy` type by reference. However, the
101 // `#[derive(..)]` attributes are not in the parsed tokens, so we are not able to check if a
102 // type is `Copy`, so we just implement both.
103 //
104 // See <https://github.com/dtolnay/syn/issues/433>
105 // ---
106 // let is_copy = unique_meta_list(type_meta.iter(), "derive")
107 // .map(extract_list_metas)
108 // .map(|metas| {
109 // metas
110 // .filter_map(get_meta_ident)
111 // .any(|derive| derive.to_string() == "Copy")
112 // }).unwrap_or(false);
113
114 let arms = variants
115 .iter()
116 .map(|variant| {
117 let ident = &variant.ident;
118 let params = match &variant.fields {
119 Fields::Unit => quote! {},
120 Fields::Unnamed(_fields) => {
121 quote! { (..) }
122 }
123 Fields::Named(_fields) => {
124 quote! { { .. } }
125 }
126 };
127
128 quote! { #name::#ident #params => #discriminants_name::#ident }
129 })
130 .collect::<Vec<_>>();
131
132 let from_fn_body = quote! { match val { #(#arms),* } };
133
134 let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
135 let impl_from = quote! {
136 impl #impl_generics ::core::convert::From< #name #ty_generics > for #discriminants_name #where_clause {
137 #[inline]
138 fn from(val: #name #ty_generics) -> #discriminants_name {
139 #from_fn_body
140 }
141 }
142 };
143 let impl_from_ref = {
144 let mut generics = ast.generics.clone();
145
146 let lifetime = parse_quote!('_enum);
147 let enum_life = quote! { & #lifetime };
148 generics.params.push(lifetime);
149
150 // Shadows the earlier `impl_generics`
151 let (impl_generics, _, _) = generics.split_for_impl();
152
153 quote! {
154 impl #impl_generics ::core::convert::From< #enum_life #name #ty_generics > for #discriminants_name #where_clause {
155 #[inline]
156 fn from(val: #enum_life #name #ty_generics) -> #discriminants_name {
157 #from_fn_body
158 }
159 }
160 }
161 };
162
163 // For now, only implement IntoDiscriminant if the user has not overriden the visibility.
164 let impl_into_discriminant = match type_properties.discriminant_vis {
165 // If the visibilty is unspecified or `pub` then we implement IntoDiscriminant
166 None | Some(syn::Visibility::Public(..)) => quote! {
167 impl #impl_generics #strum_module_path::IntoDiscriminant for #name #ty_generics #where_clause {
168 type Discriminant = #discriminants_name;
169
170 #[inline]
171 fn discriminant(&self) -> Self::Discriminant {
172 <Self::Discriminant as ::core::convert::From<&Self>>::from(self)
173 }
174 }
175 },
176 // If it's something restricted such as `pub(super)` then we skip implementing the
177 // trait for now. There are certainly scenarios where they could be equivalent, but
178 // as a heuristic, if someone is overriding the visibility, it's because they want
179 // the discriminant type to be less visible than the original type.
180 _ => quote! {},
181 };
182
183 Ok(quote! {
184 /// Auto-generated discriminant enum variants
185 #derives
186 #repr
187 #(#[ #pass_though_attributes ])*
188 #discriminants_vis enum #discriminants_name {
189 #(#discriminants),*
190 }
191
192 #impl_into_discriminant
193 #impl_from
194 #impl_from_ref
195 })
196}
197