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

#![allow(incomplete_features)]
#![feature(generic_const_exprs)]
extern crate proc_macro;
extern crate const_enum_tools;

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, DeriveInput};

#[proc_macro_derive(VariantIterable)]
pub fn derive_answer_fn(enum_item: TokenStream) -> TokenStream {
    let ast: syn::DeriveInput = parse_macro_input!(enum_item as DeriveInput);

    match ast.data {
        syn::Data::Union(union_data) => {
            let err = syn::Error::new_spanned(union_data.union_token, "Unexpected union declaration: VariantIterable can only be derived for enums.");
            err.into_compile_error().into()
        },
        syn::Data::Struct(struct_data) => {
            let err = syn::Error::new_spanned(struct_data.struct_token, "Unexpected union declaration: VariantIterable can only be derived for enums.");
            err.into_compile_error().into()
        },
        syn::Data::Enum(enum_field_data) => {
            let variants = enum_field_data.variants;
            let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
            let name = ast.ident;
            let variant_count = variants.len();

            let mut variant_index_match_arms = Vec::new();
            let mut variant_names = Vec::new();
            let mut all_unit_no_discriminant = true;

            for (index, variant) in variants.iter().enumerate() {
                let variant_name = &variant.ident;

                variant_index_match_arms.push(
                    match &variant.fields {
                        syn::Fields::Named(fields) => {
                            all_unit_no_discriminant = false;
                            let mapped = fields.named.iter().map(|_| { quote!(_) });
                            quote!(
                                #name::#variant_name(#(#mapped),*) => {
                                    #index
                                }
                            )
                        },
                        syn::Fields::Unnamed(fields) => {
                            all_unit_no_discriminant = false;
                            let mapped = fields.unnamed.iter().map(|_| { quote!(_) });
                            quote!(
                                #name::#variant_name(#(#mapped),*) => {
                                    #index
                                }
                            )
                        },
                        syn::Fields::Unit => {
                            // If there is an explicit discriminant, we might not be able to perform the bitwise copy
                            // optimization.
                            if let Some(discriminant) = &variant.discriminant {
                                match discriminant.1.clone() {
                                    // If the discriminant expression is a literal, we can check if it is equal to the default value.
                                    syn::Expr::Lit(lit) => {
                                        match lit.lit {
                                            syn::Lit::Int(int_lit) => {
                                                // If the first part of the literal before the type is the same as what it would be
                                                // because of the position in the enum, we're good. Otherwise, no optimization.
                                                if int_lit.base10_digits() != index.to_string().as_str() {
                                                    all_unit_no_discriminant = false;
                                                }
                                            },
                                            _ => {
                                                all_unit_no_discriminant = false;
                                            }
                                        }
                                    },
                                    // Otherwise, since we cannot evaluate arbitrary const expressions, we will not be able to optimize.
                                    // This involves using the long match arms list.
                                    _ => {
                                        all_unit_no_discriminant = false;
                                    },
                                }
                            }
                            quote!(
                                #name::#variant_name => {
                                    #index
                                }
                            )
                        },
                    }
                );

                variant_names.push({
                    let variant_name_string = variant_name.to_string();
                    quote!(
                        #variant_name_string
                    )
                });

            }

            let variant_countable_impl = quote!(
                impl #impl_generics const_enum_tools::VariantCountable for #name #ty_generics #where_clause {
                    const VARIANT_COUNT: usize = #variant_count;
                }
            );

            // If there are no explicit discriminants
            let variant_index_body = if all_unit_no_discriminant {
                quote!(
                    unsafe {
                        (self as *const Self).read() as usize
                    }
                )
            }
            else {
                quote!(
                    match self {
                        #(
                            #variant_index_match_arms
                        ),*
                    }
                )
            };

            quote!(
                #variant_countable_impl

                impl #impl_generics const_enum_tools::VariantIterable for #name #ty_generics #where_clause {
                    #[inline]
                    fn variant_index (&self) -> usize {
                        #variant_index_body
                    }

                    const VARIANTS: [&'static str; #variant_count] = [#(#variant_names),*];
                }
            ).into()
        }
    }

}