atomic_enum_derive/
code_gen.rs

1use proc_macro::TokenStream;
2use proc_macro2::{TokenStream as TokenStream2, Ident, Span};
3
4use crate::parser::ParsedInput;
5
6pub struct CodeGenerator {
7    input: ParsedInput,
8}
9
10impl CodeGenerator {
11    pub fn generate(input: ParsedInput) -> TokenStream {
12        let generator = CodeGenerator { input };
13
14        let trait_u16_from_self = generator.gen_trait_u16_from_self();
15        let trait_self_from_u16 = generator.get_trait_self_from_u16();
16
17        let ident = generator.ident();
18        quote::quote!(
19            #[automatically_derived]
20            impl atomic_enum::Atomize for #ident {}
21
22            #[automatically_derived]
23            #trait_u16_from_self
24
25            #[automatically_derived]
26            #trait_self_from_u16
27        ).into()
28    }
29
30    fn ident(&self) -> Ident {
31        self.input.enum_ident.clone()
32    }
33
34    fn gen_trait_u16_from_self(&self) -> TokenStream2 {
35        let ident = self.ident();
36        let mut match_stmt = quote::quote!();
37        let mut is_first = true;
38
39        for e_val in self.input.enum_values.iter() {
40            let e_ident = e_val.ident.clone();
41            let discriminant = e_val.discriminant.clone();
42
43            if is_first {
44                is_first = false;
45                match_stmt = quote::quote!( #ident::#e_ident => #discriminant );
46            } else {
47                match_stmt = quote::quote!( #match_stmt, #ident::#e_ident => #discriminant );
48            }
49        }
50
51        quote::quote!(
52            impl From<#ident> for u16 {
53                fn from(value: #ident) -> Self {
54                    match value {
55                        #match_stmt
56                    }
57                }
58            }
59        )
60    }
61
62    fn get_trait_self_from_u16(&self) -> TokenStream2 {
63        let ident = self.ident();
64        let mut match_stmt = quote::quote!();
65        let mut is_first = true;
66        let unknown_ident = Ident::new("UnknownField", Span::call_site());
67
68        for e_val in self.input.enum_values.iter() {
69            // Skip the UnknownField, it is used as default handler
70            if &e_val.ident == &unknown_ident {
71                continue;
72            }
73
74            let ident = e_val.ident.clone();
75            let discriminant = e_val.discriminant.clone();
76
77            if is_first {
78                is_first = false;
79                match_stmt = quote::quote!( #discriminant => Self::#ident );
80            } else {
81                match_stmt = quote::quote!(#match_stmt, #discriminant => Self::#ident);
82            }
83        }
84
85        match_stmt = quote::quote!(#match_stmt, _ => Self::#unknown_ident);
86
87        quote::quote!(
88            impl From<u16> for #ident {
89                fn from(value: u16) -> Self {
90                    match value {
91                        #match_stmt
92                    }
93                }
94            }
95        )
96    }
97}
atomic_enum_derive/code_gen.rs

atomic_enum_derive/
code_gen.rs
