enumkit_derive/
lib.rs

1use proc_macro::TokenStream;
2use quote::{format_ident, quote};
3use syn::{Data, DeriveInput, Fields};
4
5#[proc_macro_derive(EnumValues)]
6pub fn values_macro_derive(input: TokenStream) -> TokenStream {
7    let ast: DeriveInput = syn::parse(input).unwrap();
8    let name = &ast.ident;
9    if let Data::Enum(enum_data) = &ast.data {
10        let mut variants = vec![];
11        for v in enum_data.variants.iter() {
12            if let Fields::Unit = v.fields {
13                let var_name = &v.ident;
14                variants.push(quote! { #name::#var_name });
15            } else {
16                panic!(
17                    "Values macro can only be applied to \
18                        enums with Unit variants (no fields). '{}' is not a Unit variant",
19                    v.ident
20                );
21            }
22        }
23        let variant_count = variants.len();
24        let values_impl = quote! {
25            impl #name {
26
27                /// Produces an iterator over owned variants of this enum.
28                pub fn values() -> impl Iterator<Item=#name> {
29                    vec![
30                        #(#variants),*
31                    ].into_iter()
32                }
33
34                /// Gives the number of variants in this enum.
35                pub fn len() -> usize {
36                    #variant_count
37                }
38            }
39        };
40        values_impl.into()
41    } else {
42        panic!("Values macro can only be applied to enums.");
43    }
44}
45
46#[proc_macro_derive(EnumMapping)]
47pub fn mapping_macro_derive(input: TokenStream) -> TokenStream {
48    let ast: DeriveInput = syn::parse(input).unwrap();
49    let name = &ast.ident;
50    let visibility = &ast.vis;
51    if let Data::Enum(enum_data) = &ast.data {
52        let mut variants = vec![];
53        for v in enum_data.variants.iter() {
54            if let Fields::Unit = v.fields {
55                let var_name = &v.ident;
56                variants.push(quote! { #name::#var_name });
57            } else {
58                panic!(
59                    "Mapping macro can only be applied to \
60                        enums with Unit variants (no fields). '{}' is not a Unit variant",
61                    v.ident
62                );
63            }
64        }
65        let variant_count = variants.len();
66        let map_name = format_ident!("{}Mapping", name);
67        let into_iter_name = format_ident!("{}MappingIntoIter", name);
68        let iter_name = format_ident!("{}MappingIter", name);
69        let cases: Vec<_> = variants
70            .iter()
71            .enumerate()
72            .map(|(i, var)| {
73                quote! { #var => #i }
74            })
75            .collect();
76
77        let puts_construct: Vec<_> = variants
78            .iter()
79            .enumerate()
80            .map(|(_, var)| {
81                quote! { f(#var) }
82            })
83            .collect();
84
85        // cases of a match that map from enum index to enum value
86        let rcases: Vec<_> = variants
87            .iter()
88            .enumerate()
89            .map(|(i, var)| {
90                if i == variant_count - 1 {
91                    return quote! { _ => #var };
92                } else {
93                    quote! { #i => #var }
94                }
95            })
96            .collect();
97
98        let map_derives = if cfg!(feature="serde") {
99            quote! { #[derive(Copy, Clone, Serialize, Deserialize)] }
100        } else {
101            quote! { #[derive(Copy, Clone)] }
102        };
103
104        let values_impl = quote! {
105            #map_derives
106            #visibility struct #map_name <T>([T; #variant_count]);
107            impl<T> EnumMapping<T> for #map_name<T> {
108                type Enum = #name;
109                type EnumIter<'a> = #iter_name<'a, T>
110                    where T: 'a, Self: 'a;
111                fn get(&self, var: #name) -> &T {
112                    let index = match var {
113                        #(#cases),*
114                    };
115                    &self.0[index]
116                }
117                fn get_mut(&mut self, var: #name) -> &mut T {
118                    let index = match var {
119                        #(#cases),*
120                    };
121                    &mut self.0[index]
122                }
123                fn put(&mut self, var: #name, val: T) {
124                    let index = match var {
125                        #(#cases),*
126                    };
127                    self.0[index] = val;
128                }
129                fn new<F: FnMut(#name) -> T>(mut f: F) -> Self {
130                    let arr = [#(#puts_construct),*,];
131                    #map_name(arr)
132                }
133                fn iter(&self) -> Self::EnumIter<'_> {
134                    self.into_iter()
135                }
136            }
137
138
139            #visibility struct #into_iter_name<T>(Vec<T>, usize);
140            #visibility struct #iter_name<'a, T>(&'a #map_name<T>, usize);
141            impl<T> IntoIterator for #map_name<T> {
142                type Item = (#name, T);
143                type IntoIter = #into_iter_name<T>;
144
145                fn into_iter(self) -> Self::IntoIter {
146                    #into_iter_name(self.0.into_iter().rev().collect(), 0)
147                }
148            }
149
150            impl<'a, T> IntoIterator for &'a #map_name<T> {
151                type Item = (#name, &'a T);
152                type IntoIter = #iter_name<'a, T>;
153
154                fn into_iter(self) -> Self::IntoIter {
155                    #iter_name(self, 0)
156                }
157            }
158
159            impl<T> Iterator for #into_iter_name<T> {
160                type Item = (#name, T);
161
162                fn next(&mut self) -> Option<Self::Item> {
163                    self.0.pop().map(|t| {
164                        let i = self.1;
165                        self.1 += 1;
166                        (
167                            match i {
168                                #(#rcases),*
169                            },
170                            t
171                        )
172                    })
173                }
174            }
175
176            impl<'a, T> Iterator for #iter_name<'a, T> {
177                type Item = (#name, &'a T);
178
179                fn next(&mut self) -> Option<Self::Item> {
180                    if self.1 < #variant_count {
181                        let i = self.1;
182                        self.1 += 1;
183                        Some((
184                            match i {
185                                #(#rcases),*
186                            },
187                            &self.0.0[i]
188                        ))
189                    } else {
190                        None
191                    }
192                }
193            }
194
195        };
196        values_impl.into()
197    } else {
198        panic!("Mapping macro can only be applied to enums.");
199    }
200}
enumkit_derive/lib.rs

enumkit_derive/
lib.rs
