feat(db): Refactor storage libraries (#371)

crates/storage/codecs/derive/src/compact/mod.rs

@@ -0,0 +1,316 @@
+extern crate proc_macro2;
+use proc_macro::{self, TokenStream};
+use proc_macro2::{Ident, TokenStream as TokenStream2};
+use quote::{format_ident, quote};
+use syn::{parse_macro_input, Data, DeriveInput};
+
+mod generator;
+use generator::*;
+
+mod enums;
+use enums::*;
+
+mod flags;
+use flags::*;
+
+mod structs;
+use structs::*;
+
+// Helper Alias type
+type IsCompact = bool;
+// Helper Alias type
+type FieldName = String;
+// Helper Alias type
+type FieldType = String;
+/// `Compact` has alternative functions that can be used as a workaround for type
+/// specialization of fixed sized types.
+///
+/// Example: `Vec<H256>` vs `Vec<U256>`. The first does not
+/// require the len of the element, while the latter one does.
+type UseAlternative = bool;
+// Helper Alias type
+type StructFieldDescriptor = (FieldName, FieldType, IsCompact, UseAlternative);
+// Helper Alias type
+type FieldList = Vec<FieldTypes>;
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum FieldTypes {
+    StructField(StructFieldDescriptor),
+    EnumVariant(String),
+    EnumUnnamedField((FieldType, UseAlternative)),
+}
+
+/// Derives the [`Compact`] trait and its from/to implementations.
+pub fn derive(input: TokenStream) -> TokenStream {
+    let mut output = quote! {};
+
+    let DeriveInput { ident, data, .. } = parse_macro_input!(input);
+    let fields = get_fields(&data);
+    output.extend(generate_flag_struct(&ident, &fields));
+    output.extend(generate_from_to(&ident, &fields));
+    output.into()
+}
+
+/// Given a list of fields on a struct, extract their fields and types.
+pub fn get_fields(data: &Data) -> FieldList {
+    let mut fields = vec![];
+
+    match data {
+        Data::Struct(data) => match data.fields {
+            syn::Fields::Named(ref data_fields) => {
+                for field in &data_fields.named {
+                    load_field(field, &mut fields, false);
+                }
+                assert_eq!(fields.len(), data_fields.named.len());
+            }
+            syn::Fields::Unnamed(ref data_fields) => {
+                assert!(
+                    data_fields.unnamed.len() == 1,
+                    "Compact only allows one unnamed field. Consider making it a struct."
+                );
+                load_field(&data_fields.unnamed[0], &mut fields, false);
+            }
+            syn::Fields::Unit => todo!(),
+        },
+        Data::Enum(data) => {
+            for variant in &data.variants {
+                fields.push(FieldTypes::EnumVariant(variant.ident.to_string()));
+
+                match &variant.fields {
+                    syn::Fields::Named(_) => {
+                        panic!("Not allowed to have Enum Variants with multiple named fields. Make it a struct instead.")
+                    }
+                    syn::Fields::Unnamed(data_fields) => {
+                        assert!(
+                            data_fields.unnamed.len() == 1,
+                            "Compact only allows one unnamed field. Consider making it a struct."
+                        );
+                        load_field(&data_fields.unnamed[0], &mut fields, true);
+                    }
+                    syn::Fields::Unit => (),
+                }
+            }
+        }
+        Data::Union(_) => todo!(),
+    }
+
+    fields
+}
+
+fn load_field(field: &syn::Field, fields: &mut FieldList, is_enum: bool) {
+    if let syn::Type::Path(ref path) = field.ty {
+        let segments = &path.path.segments;
+        if !segments.is_empty() {
+            let mut ftype = String::new();
+
+            let mut use_alt_impl: UseAlternative = false;
+
+            for (index, segment) in segments.iter().enumerate() {
+                ftype.push_str(&segment.ident.to_string());
+                if index < segments.len() - 1 {
+                    ftype.push_str("::");
+                }
+
+                use_alt_impl = should_use_alt_impl(&ftype, segment);
+            }
+
+            if is_enum {
+                fields.push(FieldTypes::EnumUnnamedField((ftype.to_string(), use_alt_impl)));
+            } else {
+                let should_compact = is_flag_type(&ftype) ||
+                    field.attrs.iter().any(|attr| {
+                        attr.path.segments.iter().any(|path| path.ident == "maybe_zero")
+                    });
+
+                fields.push(FieldTypes::StructField((
+                    field.ident.as_ref().map(|i| i.to_string()).unwrap_or_default(),
+                    ftype,
+                    should_compact,
+                    use_alt_impl,
+                )));
+            }
+        }
+    }
+}
+
+/// Since there's no impl specialization in rust stable atm, once we find we have a
+/// Vec/Option we try to find out if it's a Vec/Option of a fixed size data type.
+/// eg, Vec<H256>. If so, we use another impl to code/decode its data.
+fn should_use_alt_impl(ftype: &String, segment: &syn::PathSegment) -> bool {
+    if *ftype == "Vec" || *ftype == "Option" {
+        if let syn::PathArguments::AngleBracketed(ref args) = segment.arguments {
+            if let Some(syn::GenericArgument::Type(syn::Type::Path(arg_path))) = args.args.last() {
+                if let (Some(path), 1) =
+                    (arg_path.path.segments.first(), arg_path.path.segments.len())
+                {
+                    if ["H256", "H160", "Address", "Bloom"]
+                        .contains(&path.ident.to_string().as_str())
+                    {
+                        return true
+                    }
+                }
+            }
+        }
+    }
+    false
+}
+
+/// Given the field type in a string format, return the amount of bits necessary to save its maximum
+/// length.
+pub fn get_bit_size(ftype: &str) -> u8 {
+    match ftype {
+        "bool" | "Option" => 1,
+        "TxType" => 2,
+        "u64" | "BlockNumber" | "TxNumber" | "ChainId" => 4,
+        "u128" => 5,
+        "U256" | "TxHash" => 6,
+        _ => 0,
+    }
+}
+
+/// Given the field type in a string format, checks if its type should be added to the
+/// StructFlags.
+pub fn is_flag_type(ftype: &str) -> bool {
+    get_bit_size(ftype) > 0
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use syn::parse2;
+
+    #[test]
+    fn gen() {
+        let f_struct = quote! {
+            #[derive(Debug, PartialEq, Clone)]
+            pub struct TestStruct {
+                f_u64: u64,
+                f_u256: U256,
+                f_bool_t: bool,
+                f_bool_f: bool,
+                f_option_none: Option<U256>,
+                f_option_some: Option<H256>,
+                f_option_some_u64: Option<u64>,
+                f_vec_empty: Vec<U256>,
+                f_vec_some: Vec<H160>,
+            }
+        };
+
+        // Generate code that will impl the `Compact` trait.
+        let mut output = quote! {};
+        let DeriveInput { ident, data, .. } = parse2(f_struct).unwrap();
+        let fields = get_fields(&data);
+        output.extend(generate_flag_struct(&ident, &fields));
+        output.extend(generate_from_to(&ident, &fields));
+
+        // Expected output in a TokenStream format. Commas matter!
+        let should_output = quote! {
+            pub use TestStruct_flags::TestStructFlags;
+            mod TestStruct_flags {
+                use bytes::Buf;
+                use modular_bitfield::prelude::*;
+
+                #[doc=r" Fieldset that facilitates compacting the parent type."]
+                #[bitfield]
+                #[derive(Clone, Copy, Debug, Default)]
+                pub struct TestStructFlags {
+                    pub f_u64_len: B4,
+                    pub f_u256_len: B6,
+                    pub f_bool_t_len: B1,
+                    pub f_bool_f_len: B1,
+                    pub f_option_none_len: B1,
+                    pub f_option_some_len: B1,
+                    pub f_option_some_u64_len: B1,
+                    #[skip]
+                    unused: B1,
+                }
+                impl TestStructFlags {
+                    #[doc=r" Deserializes this fieldset and returns it, alongside the original slice in an advanced position."]
+                    pub fn from(mut buf: &[u8]) -> (Self, &[u8]) {
+                        (
+                            TestStructFlags::from_bytes([buf.get_u8(), buf.get_u8(),]),
+                            buf
+                        )
+                    }
+                }
+            }
+            #[cfg(test)]
+            #[allow(dead_code)]
+            #[test_fuzz::test_fuzz]
+            fn fuzz_test_TestStruct(obj: TestStruct) {
+                let mut buf = vec![];
+                let len = obj.clone().to_compact(&mut buf);
+                let (same_obj, buf) = TestStruct::from_compact(buf.as_ref(), len);
+                assert_eq!(obj, same_obj);
+            }
+            #[test]
+            pub fn fuzz_TestStruct() {
+                fuzz_test_TestStruct(TestStruct::default())
+            }
+            impl Compact for TestStruct {
+                fn to_compact(self, buf: &mut impl bytes::BufMut) -> usize {
+                    let mut flags = TestStructFlags::default();
+                    let mut total_len = 0;
+                    let mut buffer = bytes::BytesMut::new();
+                    let f_u64_len = self.f_u64.to_compact(&mut buffer);
+                    flags.set_f_u64_len(f_u64_len as u8);
+                    let f_u256_len = self.f_u256.to_compact(&mut buffer);
+                    flags.set_f_u256_len(f_u256_len as u8);
+                    let f_bool_t_len = self.f_bool_t.to_compact(&mut buffer);
+                    flags.set_f_bool_t_len(f_bool_t_len as u8);
+                    let f_bool_f_len = self.f_bool_f.to_compact(&mut buffer);
+                    flags.set_f_bool_f_len(f_bool_f_len as u8);
+                    let f_option_none_len = self.f_option_none.to_compact(&mut buffer);
+                    flags.set_f_option_none_len(f_option_none_len as u8);
+                    let f_option_some_len = self.f_option_some.specialized_to_compact(&mut buffer);
+                    flags.set_f_option_some_len(f_option_some_len as u8);
+                    let f_option_some_u64_len = self.f_option_some_u64.to_compact(&mut buffer);
+                    flags.set_f_option_some_u64_len(f_option_some_u64_len as u8);
+                    let f_vec_empty_len = self.f_vec_empty.to_compact(&mut buffer);
+                    let f_vec_some_len = self.f_vec_some.specialized_to_compact(&mut buffer);
+                    let flags = flags.into_bytes();
+                    total_len += flags.len() + buffer.len();
+                    buf.put_slice(&flags);
+                    buf.put(buffer);
+                    total_len
+                }
+                fn from_compact(mut buf: &[u8], len: usize) -> (Self, &[u8]) {
+                    let (flags, mut buf) = TestStructFlags::from(buf);
+                    let mut f_u64 = u64::default();
+                    (f_u64, buf) = u64::from_compact(buf, flags.f_u64_len() as usize);
+                    let mut f_u256 = U256::default();
+                    (f_u256, buf) = U256::from_compact(buf, flags.f_u256_len() as usize);
+                    let mut f_bool_t = bool::default();
+                    (f_bool_t, buf) = bool::from_compact(buf, flags.f_bool_t_len() as usize);
+                    let mut f_bool_f = bool::default();
+                    (f_bool_f, buf) = bool::from_compact(buf, flags.f_bool_f_len() as usize);
+                    let mut f_option_none = Option::default();
+                    (f_option_none, buf) = Option::from_compact(buf, flags.f_option_none_len() as usize);
+                    let mut f_option_some = Option::default();
+                    (f_option_some, buf) = Option::specialized_from_compact(buf, flags.f_option_some_len() as usize);
+                    let mut f_option_some_u64 = Option::default();
+                    (f_option_some_u64, buf) =
+                        Option::from_compact(buf, flags.f_option_some_u64_len() as usize);
+                    let mut f_vec_empty = Vec::default();
+                    (f_vec_empty, buf) = Vec::from_compact(buf, buf.len());
+                    let mut f_vec_some = Vec::default();
+                    (f_vec_some, buf) = Vec::specialized_from_compact(buf, buf.len());
+                    let obj = TestStruct {
+                        f_u64: f_u64,
+                        f_u256: f_u256,
+                        f_bool_t: f_bool_t,
+                        f_bool_f: f_bool_f,
+                        f_option_none: f_option_none,
+                        f_option_some: f_option_some,
+                        f_option_some_u64: f_option_some_u64,
+                        f_vec_empty: f_vec_empty,
+                        f_vec_some: f_vec_some,
+                    };
+                    (obj, buf)
+                }
+            }
+        };
+
+        assert_eq!(output.to_string(), should_output.to_string());
+    }
+}