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, buf) = #field_type::#from_compact_ident(buf, buf.len());
70                            #ident::#variant_name(inner)
71                        }
72                    });
73                    self.fields_iterator.next();
74                }
75                FieldTypes::EnumVariant(_) => self.enum_lines.push(quote! {
76                    #current_variant_index => #ident::#variant_name,
77                }),
78                FieldTypes::StructField(_) => unreachable!(),
79            };
80        } else {
81            // This variant has no fields: Unit type
82            self.enum_lines.push(quote! {
83                #current_variant_index => #ident::#variant_name,
84            });
85        }
86        self.current_variant_index += 1;
87    }
88
89    /// Generates `to_compact` code for an enum variant.
90    ///
91    /// `fields_iterator` might look something like [`VariantUnit`, `VariantUnnamedField`, Field,
92    /// `VariantUnit`...].
93    pub fn to(&mut self, variant_name: &str, ident: &Ident) {
94        let variant_name = format_ident!("{variant_name}");
95        let current_variant_index = self.current_variant_index;
96
97        if let Some(next_field) = self.fields_iterator.peek() {
98            match next_field {
99                FieldTypes::EnumUnnamedField((_, use_alt_impl)) => {
100                    let to_compact_ident = if *use_alt_impl {
101                        format_ident!("specialized_to_compact")
102                    } else {
103                        format_ident!("to_compact")
104                    };
105
106                    // Unnamed type
107                    self.enum_lines.push(quote! {
108                        #ident::#variant_name(field) => {
109                            field.#to_compact_ident(&mut buffer);
110                            #current_variant_index
111                        },
112                    });
113                    self.fields_iterator.next();
114                }
115                FieldTypes::EnumVariant(_) => self.enum_lines.push(quote! {
116                    #ident::#variant_name => #current_variant_index,
117                }),
118                FieldTypes::StructField(_) => unreachable!(),
119            };
120        } else {
121            // This variant has no fields: Unit type
122            self.enum_lines.push(quote! {
123                #ident::#variant_name => #current_variant_index,
124            });
125        }
126        self.current_variant_index += 1;
127    }
128}
reth_codecs_derive/compact/enums.rs

reth_codecs_derive/compact/
enums.rs
