enum_map_derive/
derive_enum.rs

Help
1// SPDX-FileCopyrightText: 2021 - 2023 Kamila Borowska <kamila@borowska.pw>
2// SPDX-FileCopyrightText: 2021 Bruno Corrêa Zimmermann <brunoczim@gmail.com>
3//
4// SPDX-License-Identifier: MIT OR Apache-2.0
5
6use crate::type_length;
7use proc_macro2::TokenStream;
8use quote::{format_ident, quote, ToTokens};
9use syn::{DataEnum, Fields, FieldsNamed, FieldsUnnamed, Ident, Variant};
10
11pub fn generate(name: Ident, data_enum: DataEnum) -> TokenStream {
12    let mut generator = EnumGenerator::empty();
13    for variant in &data_enum.variants {
14        generator.handle_variant(variant);
15    }
16    generator.finish(&name)
17}
18
19#[derive(Debug)]
20struct Length {
21    units: usize,
22    opaque: TokenStream,
23}
24
25impl ToTokens for Length {
26    fn to_tokens(&self, tokens: &mut TokenStream) {
27        let Self { units, opaque } = self;
28        tokens.extend(quote! { (#units + #opaque) });
29    }
30}
31
32/// Total length is the sum of each variant's length. To represent a variant, its number is added to
33/// the sum of previous variant lengths.
34#[derive(Debug)]
35struct EnumGenerator {
36    length: Length,
37    from_usize_arms: TokenStream,
38    into_usize_arms: TokenStream,
39}
40
41impl EnumGenerator {
42    fn empty() -> Self {
43        Self {
44            length: Length {
45                units: 0,
46                opaque: quote! { 0 },
47            },
48            from_usize_arms: quote! {},
49            into_usize_arms: quote! {},
50        }
51    }
52
53    fn finish(&self, name: &Ident) -> TokenStream {
54        let Self {
55            length,
56            from_usize_arms,
57            into_usize_arms,
58        } = self;
59
60        quote! {
61            #[automatically_derived]
62            impl ::enum_map::Enum for #name {
63                const LENGTH: ::enum_map::usize = #length;
64
65                #[inline]
66                fn from_usize(value: ::enum_map::usize) -> Self {
67                    #from_usize_arms {
68                        ::enum_map::out_of_bounds()
69                    }
70                }
71
72                #[inline]
73                fn into_usize(self) -> ::enum_map::usize {
74                    match self {
75                        #into_usize_arms
76                    }
77                }
78            }
79
80            #[automatically_derived]
81            impl<V> ::enum_map::EnumArray<V> for #name {
82                type Array = [V; #length];
83            }
84        }
85    }
86
87    fn handle_variant(&mut self, variant: &Variant) {
88        match &variant.fields {
89            Fields::Unit => self.handle_unit_variant(&variant.ident),
90            Fields::Unnamed(fields) => self.handle_unnamed_variant(&variant.ident, fields),
91            Fields::Named(fields) => self.handle_named_variant(&variant.ident, fields),
92        }
93    }
94
95    /// Becomes simply `1` in counting, since this is the size of the unit.
96    fn handle_unit_variant(&mut self, variant: &Ident) {
97        let Self {
98            length,
99            from_usize_arms,
100            into_usize_arms,
101        } = self;
102        *into_usize_arms = quote! { #into_usize_arms Self::#variant => #length, };
103        *from_usize_arms = quote! {
104            #from_usize_arms if value == #length {
105                Self::#variant
106            } else
107        };
108        self.length.units += 1;
109    }
110
111    /// Its size is the product of the sizes of its members. To represent this variant, one can
112    /// think of this as representing a little-endian number. First member is simply added, but
113    /// next members are multiplied before being added.
114    fn handle_unnamed_variant(&mut self, variant: &Ident, fields: &FieldsUnnamed) {
115        let Self {
116            length,
117            from_usize_arms,
118            into_usize_arms,
119        } = self;
120        let mut expr_into = quote! { #length };
121        let mut fields_length = quote! { 1usize };
122        let mut params_from = quote! {};
123        for (i, field) in fields.unnamed.iter().enumerate() {
124            let ident = format_ident!("p{}", i);
125            let ty = &field.ty;
126            let field_length = type_length(ty);
127
128            expr_into = quote! {
129                (#expr_into + #fields_length * ::enum_map::Enum::into_usize(#ident))
130            };
131
132            params_from = quote! {
133                #params_from <#ty as ::enum_map::Enum>::from_usize(
134                    (value - #length) / #fields_length % #field_length
135                ),
136            };
137
138            fields_length = quote! { (#fields_length * #field_length) };
139        }
140
141        *length = Length {
142            units: 0,
143            opaque: quote! { (#length + #fields_length) },
144        };
145
146        let from_arms = &from_usize_arms;
147        *from_usize_arms = quote! {
148            #from_arms if value < #length {
149                Self::#variant(#params_from)
150            } else
151        };
152
153        let mut params_into = quote! {};
154        for i in 0..fields.unnamed.len() {
155            let ident = format_ident!("p{}", i);
156            params_into = quote! { #params_into #ident, };
157        }
158
159        *into_usize_arms = quote! {
160            #into_usize_arms Self::#variant(#params_into) => #expr_into,
161        };
162    }
163
164    /// Its size is the product of the sizes of its members. To represent this variant, one can
165    /// think of this as representing a little-endian number. First member is simply added, but
166    /// next members are multiplied before being added.
167    fn handle_named_variant(&mut self, variant: &Ident, fields: &FieldsNamed) {
168        let Self {
169            length,
170            from_usize_arms,
171            into_usize_arms,
172        } = self;
173        let mut expr_into = quote! { #length };
174        let mut fields_length = quote! { 1usize };
175        let mut params_from = quote! {};
176
177        for field in fields.named.iter() {
178            let ident = field.ident.as_ref().unwrap();
179            let ty = &field.ty;
180            let field_length = type_length(ty);
181
182            expr_into = quote! {
183                (#expr_into + #fields_length * ::enum_map::Enum::into_usize(#ident))
184            };
185
186            params_from = quote! {
187                #params_from #ident: <#ty as ::enum_map::Enum>::from_usize(
188                    (value - #length) / #fields_length % #field_length
189                ),
190            };
191
192            fields_length = quote! { (#fields_length * #field_length) };
193        }
194
195        *length = Length {
196            units: 0,
197            opaque: quote! { (#length + #fields_length) },
198        };
199
200        *from_usize_arms = quote! {
201            #from_usize_arms if value < #length {
202                Self::#variant { #params_from }
203            } else
204        };
205
206        let mut params_into = quote! {};
207        for field in fields.named.iter() {
208            let ident = field.ident.as_ref().unwrap();
209            params_into = quote! { #params_into #ident, };
210        }
211
212        *into_usize_arms = quote! {
213            #into_usize_arms Self::#variant { #params_into } => #expr_into,
214        };
215    }
216}
enum_map_derive/derive_enum.rs

enum_map_derive/
derive_enum.rs
