strum_macros/macros/
enum_discriminants.rs

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}