reth_codecs_derive/compact/
mod.rs

1use proc_macro::TokenStream;
2use proc_macro2::{Ident, TokenStream as TokenStream2};
3use quote::{format_ident, quote};
4use syn::{Data, DeriveInput, Generics};
5
6mod generator;
7use generator::*;
8
9mod enums;
10use enums::*;
11
12mod flags;
13use flags::*;
14
15mod structs;
16use structs::*;
17
18use crate::ZstdConfig;
19
20// Helper Alias type
21type FieldType = String;
22/// `Compact` has alternative functions that can be used as a workaround for type
23/// specialization of fixed sized types.
24///
25/// Example: `Vec<B256>` vs `Vec<U256>`. The first does not
26/// require the len of the element, while the latter one does.
27type UseAlternative = bool;
28// Helper Alias type
29#[derive(Debug, Clone, Eq, PartialEq)]
30pub struct StructFieldDescriptor {
31    name: String,
32    ftype: String,
33    is_compact: bool,
34    use_alt_impl: bool,
35    is_reference: bool,
36}
37// Helper Alias type
38type FieldList = Vec<FieldTypes>;
39
40#[derive(Debug, Clone, Eq, PartialEq)]
41pub enum FieldTypes {
42    StructField(StructFieldDescriptor),
43    EnumVariant(String),
44    EnumUnnamedField((FieldType, UseAlternative)),
45}
46
47/// Derives the `Compact` trait and its from/to implementations.
48pub fn derive(input: DeriveInput, zstd: Option<ZstdConfig>) -> TokenStream {
49    let mut output = quote! {};
50
51    let DeriveInput { ident, data, generics, attrs, .. } = input;
52
53    let has_lifetime = has_lifetime(&generics);
54
55    let fields = get_fields(&data);
56    output.extend(generate_flag_struct(&ident, &attrs, has_lifetime, &fields, zstd.is_some()));
57    output.extend(generate_from_to(&ident, &attrs, has_lifetime, &fields, zstd));
58    output.into()
59}
60
61pub fn has_lifetime(generics: &Generics) -> bool {
62    generics.lifetimes().next().is_some()
63}
64
65/// Given a list of fields on a struct, extract their fields and types.
66pub fn get_fields(data: &Data) -> FieldList {
67    let mut fields = vec![];
68
69    match data {
70        Data::Struct(data) => match data.fields {
71            syn::Fields::Named(ref data_fields) => {
72                for field in &data_fields.named {
73                    load_field(field, &mut fields, false);
74                }
75                assert_eq!(fields.len(), data_fields.named.len(), "get_fields");
76            }
77            syn::Fields::Unnamed(ref data_fields) => {
78                assert_eq!(
79                    data_fields.unnamed.len(),
80                    1,
81                    "Compact only allows one unnamed field. Consider making it a struct."
82                );
83                load_field(&data_fields.unnamed[0], &mut fields, false);
84            }
85            syn::Fields::Unit => todo!(),
86        },
87        Data::Enum(data) => {
88            for variant in &data.variants {
89                fields.push(FieldTypes::EnumVariant(variant.ident.to_string()));
90
91                match &variant.fields {
92                    syn::Fields::Named(_) => {
93                        panic!(
94                            "Not allowed to have Enum Variants with multiple named fields. Make it a struct instead."
95                        )
96                    }
97                    syn::Fields::Unnamed(data_fields) => {
98                        assert_eq!(
99                            data_fields.unnamed.len(),
100                            1,
101                            "Compact only allows one unnamed field. Consider making it a struct."
102                        );
103                        load_field(&data_fields.unnamed[0], &mut fields, true);
104                    }
105                    syn::Fields::Unit => (),
106                }
107            }
108        }
109        Data::Union(_) => todo!(),
110    }
111
112    fields
113}
114
115fn load_field(field: &syn::Field, fields: &mut FieldList, is_enum: bool) {
116    match field.ty {
117        syn::Type::Reference(ref reference) => match &*reference.elem {
118            syn::Type::Path(path) => {
119                load_field_from_segments(&path.path.segments, is_enum, fields, field)
120            }
121            _ => unimplemented!("{:?}", &field.ident),
122        },
123        syn::Type::Path(ref path) => {
124            load_field_from_segments(&path.path.segments, is_enum, fields, field)
125        }
126        _ => unimplemented!("{:?}", &field.ident),
127    }
128}
129
130fn load_field_from_segments(
131    segments: &syn::punctuated::Punctuated<syn::PathSegment, syn::token::PathSep>,
132    is_enum: bool,
133    fields: &mut Vec<FieldTypes>,
134    field: &syn::Field,
135) {
136    if !segments.is_empty() {
137        let mut ftype = String::new();
138
139        let mut use_alt_impl: UseAlternative = false;
140
141        for (index, segment) in segments.iter().enumerate() {
142            ftype.push_str(&segment.ident.to_string());
143            if index < segments.len() - 1 {
144                ftype.push_str("::");
145            }
146
147            use_alt_impl = should_use_alt_impl(&ftype, segment);
148        }
149
150        if is_enum {
151            fields.push(FieldTypes::EnumUnnamedField((ftype, use_alt_impl)));
152        } else {
153            let should_compact = is_flag_type(&ftype) ||
154                field.attrs.iter().any(|attr| {
155                    attr.path().segments.iter().any(|path| path.ident == "maybe_zero")
156                });
157
158            fields.push(FieldTypes::StructField(StructFieldDescriptor {
159                name: field.ident.as_ref().map(|i| i.to_string()).unwrap_or_default(),
160                ftype,
161                is_compact: should_compact,
162                use_alt_impl,
163                is_reference: matches!(field.ty, syn::Type::Reference(_)),
164            }));
165        }
166    }
167}
168
169/// Since there's no impl specialization in rust stable atm, once we find we have a
170/// Vec/Option we try to find out if it's a Vec/Option of a fixed size data type, e.g. `Vec<B256>`.
171///
172/// If so, we use another impl to code/decode its data.
173fn should_use_alt_impl(ftype: &str, segment: &syn::PathSegment) -> bool {
174    if ftype == "Vec" || ftype == "Option" {
175        if let syn::PathArguments::AngleBracketed(ref args) = segment.arguments {
176            if let Some(syn::GenericArgument::Type(syn::Type::Path(arg_path))) = args.args.last() {
177                if let (Some(path), 1) =
178                    (arg_path.path.segments.first(), arg_path.path.segments.len())
179                {
180                    if [
181                        "B256",
182                        "Address",
183                        "Address",
184                        "Bloom",
185                        "TxHash",
186                        "BlockHash",
187                        "CompactPlaceholder",
188                    ]
189                    .contains(&path.ident.to_string().as_str())
190                    {
191                        return true
192                    }
193                }
194            }
195        }
196    }
197    false
198}
199
200/// Given the field type in a string format, return the amount of bits necessary to save its maximum
201/// length.
202pub fn get_bit_size(ftype: &str) -> u8 {
203    match ftype {
204        "TransactionKind" | "TxKind" | "bool" | "Option" | "Signature" => 1,
205        "TxType" | "OpTxType" | "SeismicTxType" => 2,
206        "u64" | "BlockNumber" | "TxNumber" | "ChainId" | "NumTransactions" => 4,
207        "u128" => 5,
208        "U256" => 6,
209        "u8" => 1,
210        _ => 0,
211    }
212}
213
214/// Given the field type in a string format, checks if its type should be added to the
215/// `StructFlags`.
216pub fn is_flag_type(ftype: &str) -> bool {
217    get_bit_size(ftype) > 0
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223    use similar_asserts::assert_eq;
224    use syn::parse2;
225
226    #[test]
227    fn compact_codec() {
228        let f_struct = quote! {
229             #[derive(Debug, PartialEq, Clone)]
230             pub struct TestStruct {
231                 f_u64: u64,
232                 f_u256: U256,
233                 f_bool_t: bool,
234                 f_bool_f: bool,
235                 f_option_none: Option<U256>,
236                 f_option_some: Option<B256>,
237                 f_option_some_u64: Option<u64>,
238                 f_vec_empty: Vec<U256>,
239                 f_vec_some: Vec<Address>,
240             }
241        };
242
243        // Generate code that will impl the `Compact` trait.
244        let mut output = quote! {};
245        let DeriveInput { ident, data, attrs, .. } = parse2(f_struct).unwrap();
246        let fields = get_fields(&data);
247        output.extend(generate_flag_struct(&ident, &attrs, false, &fields, false));
248        output.extend(generate_from_to(&ident, &attrs, false, &fields, None));
249
250        // Expected output in a TokenStream format. Commas matter!
251        let should_output = quote! {
252            impl TestStruct {
253                #[doc = "Used bytes by [`TestStructFlags`]"]
254                pub const fn bitflag_encoded_bytes() -> usize {
255                    2u8 as usize
256                }
257                #[doc = "Unused bits for new fields by [`TestStructFlags`]"]
258                pub const fn bitflag_unused_bits() -> usize {
259                    1u8 as usize
260                }
261            }
262
263            pub use TestStruct_flags::TestStructFlags;
264
265            #[expect(non_snake_case)]
266            mod TestStruct_flags {
267                use reth_codecs::__private::Buf;
268                use reth_codecs::__private::modular_bitfield;
269                use reth_codecs::__private::modular_bitfield::prelude::*;
270                #[doc = "Fieldset that facilitates compacting the parent type. Used bytes: 2 | Unused bits: 1"]
271                #[bitfield]
272                #[derive(Clone, Copy, Debug, Default)]
273                pub struct TestStructFlags {
274                    pub f_u64_len: B4,
275                    pub f_u256_len: B6,
276                    pub f_bool_t_len: B1,
277                    pub f_bool_f_len: B1,
278                    pub f_option_none_len: B1,
279                    pub f_option_some_len: B1,
280                    pub f_option_some_u64_len: B1,
281                    #[skip]
282                    unused: B1,
283                }
284                impl TestStructFlags {
285                    #[doc = r" Deserializes this fieldset and returns it, alongside the original slice in an advanced position."]
286                    pub fn from(mut buf: &[u8]) -> (Self, &[u8]) {
287                        (
288                            TestStructFlags::from_bytes([buf.get_u8(), buf.get_u8(),]),
289                            buf
290                        )
291                    }
292                }
293            }
294            #[cfg(test)]
295            #[expect(dead_code)]
296            #[test_fuzz::test_fuzz]
297            fn fuzz_test_test_struct(obj: TestStruct) {
298                use reth_codecs::Compact;
299                let mut buf = vec![];
300                let len = obj.clone().to_compact(&mut buf);
301                let (same_obj, buf) = TestStruct::from_compact(buf.as_ref(), len);
302                assert_eq!(obj, same_obj);
303            }
304            #[test]
305            #[expect(missing_docs)]
306            pub fn fuzz_test_struct() {
307                fuzz_test_test_struct(TestStruct::default())
308            }
309            impl reth_codecs::Compact for TestStruct {
310                fn to_compact<B>(&self, buf: &mut B) -> usize where B: reth_codecs::__private::bytes::BufMut + AsMut<[u8]> {
311                    let mut flags = TestStructFlags::default();
312                    let mut total_length = 0;
313                    let mut buffer = reth_codecs::__private::bytes::BytesMut::new();
314                    let f_u64_len = self.f_u64.to_compact(&mut buffer);
315                    flags.set_f_u64_len(f_u64_len as u8);
316                    let f_u256_len = self.f_u256.to_compact(&mut buffer);
317                    flags.set_f_u256_len(f_u256_len as u8);
318                    let f_bool_t_len = self.f_bool_t.to_compact(&mut buffer);
319                    flags.set_f_bool_t_len(f_bool_t_len as u8);
320                    let f_bool_f_len = self.f_bool_f.to_compact(&mut buffer);
321                    flags.set_f_bool_f_len(f_bool_f_len as u8);
322                    let f_option_none_len = self.f_option_none.to_compact(&mut buffer);
323                    flags.set_f_option_none_len(f_option_none_len as u8);
324                    let f_option_some_len = self.f_option_some.specialized_to_compact(&mut buffer);
325                    flags.set_f_option_some_len(f_option_some_len as u8);
326                    let f_option_some_u64_len = self.f_option_some_u64.to_compact(&mut buffer);
327                    flags.set_f_option_some_u64_len(f_option_some_u64_len as u8);
328                    let f_vec_empty_len = self.f_vec_empty.to_compact(&mut buffer);
329                    let f_vec_some_len = self.f_vec_some.specialized_to_compact(&mut buffer);
330                    let flags = flags.into_bytes();
331                    total_length += flags.len() + buffer.len();
332                    buf.put_slice(&flags);
333                    buf.put(buffer);
334                    total_length
335                }
336                fn from_compact(mut buf: &[u8], len: usize) -> (Self, &[u8]) {
337                    let (flags, mut buf) = TestStructFlags::from(buf);
338                    let (f_u64, new_buf) = u64::from_compact(buf, flags.f_u64_len() as usize);
339                    buf = new_buf;
340                    let (f_u256, new_buf) = U256::from_compact(buf, flags.f_u256_len() as usize);
341                    buf = new_buf;
342                    let (f_bool_t, new_buf) = bool::from_compact(buf, flags.f_bool_t_len() as usize);
343                    buf = new_buf;
344                    let (f_bool_f, new_buf) = bool::from_compact(buf, flags.f_bool_f_len() as usize);
345                    buf = new_buf;
346                    let (f_option_none, new_buf) = Option::from_compact(buf, flags.f_option_none_len() as usize);
347                    buf = new_buf;
348                    let (f_option_some, new_buf) = Option::specialized_from_compact(buf, flags.f_option_some_len() as usize);
349                    buf = new_buf;
350                    let (f_option_some_u64, new_buf) = Option::from_compact(buf, flags.f_option_some_u64_len() as usize);
351                    buf = new_buf;
352                    let (f_vec_empty, new_buf) = Vec::from_compact(buf, buf.len());
353                    buf = new_buf;
354                    let (f_vec_some, new_buf) = Vec::specialized_from_compact(buf, buf.len());
355                    buf = new_buf;
356                    let obj = TestStruct {
357                        f_u64: f_u64,
358                        f_u256: f_u256,
359                        f_bool_t: f_bool_t,
360                        f_bool_f: f_bool_f,
361                        f_option_none: f_option_none,
362                        f_option_some: f_option_some,
363                        f_option_some_u64: f_option_some_u64,
364                        f_vec_empty: f_vec_empty,
365                        f_vec_some: f_vec_some,
366                    };
367                    (obj, buf)
368                }
369            }
370        };
371
372        assert_eq!(
373            syn::parse2::<syn::File>(output).unwrap(),
374            syn::parse2::<syn::File>(should_output).unwrap()
375        );
376    }
377}
reth_codecs_derive/compact/mod.rs

reth_codecs_derive/compact/
mod.rs
