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

use proc_macro::TokenStream;

use darling::FromDeriveInput;
use darling::FromVariant;
use proc_macro_error::abort;
use syn::parse_macro_input;

use crate::util::get_unit_ident;

#[derive(FromVariant)]
/// Type for gathering each variants ident and fields.
struct VarOpts {
    ident: syn::Ident,
    fields: darling::ast::Fields<darling::util::Ignored>,
}

#[derive(FromDeriveInput)]
#[darling(attributes(unit_attrs), forward_attrs(unit_enum))]
/// Type for parsing the input and extracting the
/// unit_name attribute like: `#[unit_enum(UnitFoo)]`.
struct Opts {
    ident: syn::Ident,
    attrs: Vec<syn::Attribute>,
    data: darling::ast::Data<VarOpts, darling::util::Ignored>,
    #[darling(default)]
    forward: Option<syn::Meta>,
}

pub fn derive(input: TokenStream) -> TokenStream {
    // Parse the input.
    let input = parse_macro_input!(input);
    let opts = match Opts::from_derive_input(&input) {
        Ok(o) => o,
        Err(e) => abort!(e.span(), e),
    };
    let Opts {
        ident,
        attrs,
        data,
        forward,
    } = opts;

    // Extract the variants.
    let variants = match data {
        darling::ast::Data::Enum(variants) => variants,
        _ => abort!(ident, "UnitEnum can only be derived on Enums"),
    };

    // Parse the `unit_name` attribute.
    let unit_ident = get_unit_ident(&attrs);

    // Generate the unit variants.
    let units: proc_macro2::TokenStream = variants
        .iter()
        .map(|VarOpts { ident, .. }| quote::quote! {#ident,})
        .collect();

    // Generate the match arms for `match &Self` to the unit variant.
    let units_match: proc_macro2::TokenStream = variants
        .iter()
        .map(
            |VarOpts {
                 ident: v_ident,
                 fields,
                 ..
             }| {
                let enum_style = match fields.style {
                    darling::ast::Style::Struct => quote::quote! {{..}},
                    darling::ast::Style::Unit => quote::quote! {},
                    darling::ast::Style::Tuple => quote::quote! {(_)},
                };
                quote::quote! {#ident::#v_ident #enum_style => #unit_ident::#v_ident,}
            },
        )
        .collect();

    // Forward any attributes that are meant for the unit enum.
    let unit_attrs: proc_macro2::TokenStream = match forward {
        Some(syn::Meta::List(syn::MetaList { nested, .. })) => {
            nested.iter().map(|a| quote::quote! {#[#a]}).collect()
        }
        _ => quote::quote! {},
    };

    let output = quote::quote! {
        // Impl the UnitEnum for Self
        impl UnitEnum for #ident {
            type Unit = #unit_ident;

            fn to_unit(&self) -> Self::Unit {
                match self {
                    #units_match
                }
            }

            fn unit_iter() -> Box<dyn Iterator<Item = Self::Unit>>
            {
                Box::new(#unit_ident::iter())
            }
        }

        // Add the forwarded attributes and
        // declare the unit enum.
        #unit_attrs
        #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
        pub enum #unit_ident {
            #units
        }

        // Add a iter function that creates
        // an iterator for each variant.
        impl #unit_ident {
            pub fn iter() -> impl Iterator<Item = Self> {
                use #unit_ident::*;
                [#units].into_iter()
            }
        }
    };
    output.into()
}