reth_codecs_derive/compact/
enums.rs

1use super::*;
2
3#[derive(Debug)]
4pub struct EnumHandler<'a> {
5    current_variant_index: u8,
6    fields_iterator: std::iter::Peekable<std::slice::Iter<'a, FieldTypes>>,
7    enum_lines: Vec<TokenStream2>,
8}
9
10impl<'a> EnumHandler<'a> {
11    pub fn new(fields: &'a FieldList) -> Self {
12        EnumHandler {
13            current_variant_index: 0u8,
14            enum_lines: vec![],
15            fields_iterator: fields.iter().peekable(),
16        }
17    }
18
19    pub fn next_field(&mut self) -> Option<&'a FieldTypes> {
20        self.fields_iterator.next()
21    }
22
23    pub fn generate_to(mut self, ident: &Ident) -> Vec<TokenStream2> {
24        while let Some(field) = self.next_field() {
25            match field {
26                //  The following method will advance the
27                // `fields_iterator` by itself and stop right before the next variant.
28                FieldTypes::EnumVariant(name) => self.to(name, ident),
29                FieldTypes::EnumUnnamedField(_) | FieldTypes::StructField(_) => unreachable!(),
30            }
31        }
32        self.enum_lines
33    }
34
35    pub fn generate_from(mut self, ident: &Ident) -> Vec<TokenStream2> {
36        while let Some(field) = self.next_field() {
37            match field {
38                //  The following method will advance the
39                // `fields_iterator` by itself and stop right before the next variant.
40                FieldTypes::EnumVariant(name) => self.from(name, ident),
41                FieldTypes::EnumUnnamedField(_) | FieldTypes::StructField(_) => unreachable!(),
42            }
43        }
44        self.enum_lines
45    }
46
47    /// Generates `from_compact` code for an enum variant.
48    ///
49    /// `fields_iterator` might look something like \[`VariantUnit`, `VariantUnnamedField`, Field,
50    /// `VariantUnit`...\].
51    pub fn from(&mut self, variant_name: &str, ident: &Ident) {
52        let variant_name = format_ident!("{variant_name}");
53        let current_variant_index = self.current_variant_index;
54
55        if let Some(next_field) = self.fields_iterator.peek() {
56            match next_field {
57                FieldTypes::EnumUnnamedField((next_ftype, use_alt_impl)) => {
58                    // This variant is of the type `EnumVariant(UnnamedField)`
59                    let field_type = format_ident!("{next_ftype}");
60                    let from_compact_ident = if *use_alt_impl {
61                        format_ident!("specialized_from_compact")
62                    } else {
63                        format_ident!("from_compact")
64                    };
65
66                    // Unnamed type
67                    self.enum_lines.push(quote! {
68                        #current_variant_index => {
69                            let (inner, new_buf) = #field_type::#from_compact_ident(buf, buf.len());
70                            buf = new_buf;
71                            #ident::#variant_name(inner)
72                        }
73                    });
74                    self.fields_iterator.next();
75                }
76                FieldTypes::EnumVariant(_) => self.enum_lines.push(quote! {
77                    #current_variant_index => #ident::#variant_name,
78                }),
79                FieldTypes::StructField(_) => unreachable!(),
80            };
81        } else {
82            // This variant has no fields: Unit type
83            self.enum_lines.push(quote! {
84                #current_variant_index => #ident::#variant_name,
85            });
86        }
87        self.current_variant_index += 1;
88    }
89
90    /// Generates `to_compact` code for an enum variant.
91    ///
92    /// `fields_iterator` might look something like [`VariantUnit`, `VariantUnnamedField`, Field,
93    /// `VariantUnit`...].
94    pub fn to(&mut self, variant_name: &str, ident: &Ident) {
95        let variant_name = format_ident!("{variant_name}");
96        let current_variant_index = self.current_variant_index;
97
98        if let Some(next_field) = self.fields_iterator.peek() {
99            match next_field {
100                FieldTypes::EnumUnnamedField((_, use_alt_impl)) => {
101                    let to_compact_ident = if *use_alt_impl {
102                        format_ident!("specialized_to_compact")
103                    } else {
104                        format_ident!("to_compact")
105                    };
106
107                    // Unnamed type
108                    self.enum_lines.push(quote! {
109                        #ident::#variant_name(field) => {
110                            field.#to_compact_ident(&mut buffer);
111                            #current_variant_index
112                        },
113                    });
114                    self.fields_iterator.next();
115                }
116                FieldTypes::EnumVariant(_) => self.enum_lines.push(quote! {
117                    #ident::#variant_name => #current_variant_index,
118                }),
119                FieldTypes::StructField(_) => unreachable!(),
120            };
121        } else {
122            // This variant has no fields: Unit type
123            self.enum_lines.push(quote! {
124                #ident::#variant_name => #current_variant_index,
125            });
126        }
127        self.current_variant_index += 1;
128    }
129}
reth_codecs_derive/compact/enums.rs

reth_codecs_derive/compact/
enums.rs
