1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167

use proc_macro2::{Span, TokenStream, TokenTree};
use quote::{quote, ToTokens};
use syn::parse_quote;
use syn::{Data, DeriveInput};

use crate::helpers::{non_enum_error, strum_discriminants_passthrough_error, HasTypeProperties};

/// Attributes to copy from the main enum's variants to the discriminant enum's variants.
///
/// Attributes not in this list may be for other `proc_macro`s on the main enum, and may cause
/// compilation problems when copied across.
const ATTRIBUTES_TO_COPY: &[&str] = &["doc", "cfg", "allow", "deny", "strum_discriminants"];

pub fn enum_discriminants_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
    let name = &ast.ident;
    let vis = &ast.vis;

    let variants = match &ast.data {
        Data::Enum(v) => &v.variants,
        _ => return Err(non_enum_error()),
    };

    // Derives for the generated enum
    let type_properties = ast.get_type_properties()?;

    let derives = type_properties.discriminant_derives;

    let derives = quote! {
        #[derive(Clone, Copy, Debug, PartialEq, Eq, #(#derives),*)]
    };

    // Work out the name
    let default_name = syn::Ident::new(
        &format!("{}Discriminants", name.to_string()),
        Span::call_site(),
    );

    let discriminants_name = type_properties.discriminant_name.unwrap_or(default_name);
    let discriminants_vis = type_properties
        .discriminant_vis
        .unwrap_or_else(|| vis.clone());

    // Pass through all other attributes
    let pass_though_attributes = type_properties.discriminant_others;

    // Add the variants without fields, but exclude the `strum` meta item
    let mut discriminants = Vec::new();
    for variant in variants {
        let ident = &variant.ident;

        // Don't copy across the "strum" meta attribute. Only passthrough the whitelisted
        // attributes and proxy `#[strum_discriminants(...)]` attributes
        let attrs = variant
            .attrs
            .iter()
            .filter(|attr| {
                ATTRIBUTES_TO_COPY
                    .iter()
                    .any(|attr_whitelisted| attr.path.is_ident(attr_whitelisted))
            })
            .map(|attr| {
                if attr.path.is_ident("strum_discriminants") {
                    let passthrough_group = attr
                        .tokens
                        .clone()
                        .into_iter()
                        .next()
                        .ok_or_else(|| strum_discriminants_passthrough_error(attr))?;
                    let passthrough_attribute = match passthrough_group {
                        TokenTree::Group(ref group) => group.stream(),
                        _ => {
                            return Err(strum_discriminants_passthrough_error(passthrough_group));
                        }
                    };
                    if passthrough_attribute.is_empty() {
                        return Err(strum_discriminants_passthrough_error(passthrough_group));
                    }
                    Ok(quote! { #[#passthrough_attribute] })
                } else {
                    Ok(attr.to_token_stream())
                }
            })
            .collect::<Result<Vec<_>, _>>()?;

        discriminants.push(quote! { #(#attrs)* #ident });
    }

    // Ideally:
    //
    // * For `Copy` types, we `impl From<TheEnum> for TheEnumDiscriminants`
    // * For `!Copy` types, we `impl<'enum> From<&'enum TheEnum> for TheEnumDiscriminants`
    //
    // That way we ensure users are not able to pass a `Copy` type by reference. However, the
    // `#[derive(..)]` attributes are not in the parsed tokens, so we are not able to check if a
    // type is `Copy`, so we just implement both.
    //
    // See <https://github.com/dtolnay/syn/issues/433>
    // ---
    // let is_copy = unique_meta_list(type_meta.iter(), "derive")
    //     .map(extract_list_metas)
    //     .map(|metas| {
    //         metas
    //             .filter_map(get_meta_ident)
    //             .any(|derive| derive.to_string() == "Copy")
    //     }).unwrap_or(false);

    let arms = variants
        .iter()
        .map(|variant| {
            let ident = &variant.ident;

            use syn::Fields::*;
            let params = match &variant.fields {
                Unit => quote! {},
                Unnamed(_fields) => {
                    quote! { (..) }
                }
                Named(_fields) => {
                    quote! { { .. } }
                }
            };

            quote! { #name::#ident #params => #discriminants_name::#ident }
        })
        .collect::<Vec<_>>();

    let from_fn_body = quote! { match val { #(#arms),* } };

    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
    let impl_from = quote! {
        impl #impl_generics ::core::convert::From< #name #ty_generics > for #discriminants_name #where_clause {
            fn from(val: #name #ty_generics) -> #discriminants_name {
                #from_fn_body
            }
        }
    };
    let impl_from_ref = {
        let mut generics = ast.generics.clone();

        let lifetime = parse_quote!('_enum);
        let enum_life = quote! { & #lifetime };
        generics.params.push(lifetime);

        // Shadows the earlier `impl_generics`
        let (impl_generics, _, _) = generics.split_for_impl();

        quote! {
            impl #impl_generics ::core::convert::From< #enum_life #name #ty_generics > for #discriminants_name #where_clause {
                fn from(val: #enum_life #name #ty_generics) -> #discriminants_name {
                    #from_fn_body
                }
            }
        }
    };

    Ok(quote! {
        /// Auto-generated discriminant enum variants
        #derives
        #(#[ #pass_though_attributes ])*
        #discriminants_vis enum #discriminants_name {
            #(#discriminants),*
        }

        #impl_from
        #impl_from_ref
    })
}