From 76f2d33fa3bee100da80970be16895a229180a8c Mon Sep 17 00:00:00 2001 From: Arsenii Kulikov Date: Wed, 11 Dec 2024 01:01:50 +0400 Subject: [PATCH 1/3] wip --- Cargo.lock | 1 - crates/primitives/src/receipt.rs | 6 +- crates/prune/types/Cargo.toml | 1 - .../codecs/derive/src/compact/generator.rs | 77 ++++++++++--------- .../storage/codecs/derive/src/compact/mod.rs | 10 ++- crates/storage/codecs/derive/src/lib.rs | 50 ++++++++++-- crates/storage/codecs/src/private.rs | 2 +- 7 files changed, 98 insertions(+), 49 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 92b089a0eafd..218556ee9137 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8873,7 +8873,6 @@ dependencies = [ "alloy-primitives", "arbitrary", "assert_matches", - "bytes", "derive_more", "modular-bitfield", "proptest", diff --git a/crates/primitives/src/receipt.rs b/crates/primitives/src/receipt.rs index 62c664e22a46..549c032cee0d 100644 --- a/crates/primitives/src/receipt.rs +++ b/crates/primitives/src/receipt.rs @@ -13,8 +13,6 @@ use reth_primitives_traits::receipt::ReceiptExt; use serde::{Deserialize, Serialize}; use crate::TxType; -#[cfg(feature = "reth-codec")] -use reth_zstd_compressors::{RECEIPT_COMPRESSOR, RECEIPT_DECOMPRESSOR}; /// Retrieves gas spent by transactions as a vector of tuples (transaction index, gas used). pub use reth_primitives_traits::receipt::gas_spent_by_transactions; @@ -25,6 +23,10 @@ pub use reth_primitives_traits::receipt::gas_spent_by_transactions; )] #[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::CompactZstd))] #[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests)] +#[cfg_attr(any(test, feature = "reth-codec"), reth_zstd( + compressor = reth_zstd_compressors::RECEIPT_COMPRESSOR, + decompressor = reth_zstd_compressors::RECEIPT_DECOMPRESSOR +))] #[rlp(trailing)] pub struct Receipt { /// Receipt type. diff --git a/crates/prune/types/Cargo.toml b/crates/prune/types/Cargo.toml index 5446d6f76ff5..75093511310d 100644 --- a/crates/prune/types/Cargo.toml +++ b/crates/prune/types/Cargo.toml @@ -15,7 +15,6 @@ workspace = true reth-codecs.workspace = true alloy-primitives.workspace = true -bytes.workspace = true derive_more.workspace = true modular-bitfield.workspace = true serde.workspace = true diff --git a/crates/storage/codecs/derive/src/compact/generator.rs b/crates/storage/codecs/derive/src/compact/generator.rs index 63fef05ad705..a84913f59e81 100644 --- a/crates/storage/codecs/derive/src/compact/generator.rs +++ b/crates/storage/codecs/derive/src/compact/generator.rs @@ -1,6 +1,7 @@ //! Code generator for the `Compact` trait. use super::*; +use crate::ZstdConfig; use convert_case::{Case, Casing}; use syn::{Attribute, LitStr}; @@ -10,20 +11,20 @@ pub fn generate_from_to( attrs: &[Attribute], has_lifetime: bool, fields: &FieldList, - is_zstd: bool, + zstd: Option, ) -> TokenStream2 { let flags = format_ident!("{ident}Flags"); - let to_compact = generate_to_compact(fields, ident, is_zstd); - let from_compact = generate_from_compact(fields, ident, is_zstd); + let reth_codecs = parse_reth_codecs_path(attrs).unwrap(); + + let to_compact = generate_to_compact(fields, ident, zstd.clone(), &reth_codecs); + let from_compact = generate_from_compact(fields, ident, zstd); let snake_case_ident = ident.to_string().to_case(Case::Snake); let fuzz = format_ident!("fuzz_test_{snake_case_ident}"); let test = format_ident!("fuzz_{snake_case_ident}"); - let reth_codecs = parse_reth_codecs_path(attrs).unwrap(); - let lifetime = if has_lifetime { quote! { 'a } } else { @@ -77,7 +78,7 @@ pub fn generate_from_to( #fuzz_tests #impl_compact { - fn to_compact(&self, buf: &mut B) -> usize where B: bytes::BufMut + AsMut<[u8]> { + fn to_compact(&self, buf: &mut B) -> usize where B: #reth_codecs::__private::bytes::BufMut + AsMut<[u8]> { let mut flags = #flags::default(); let mut total_length = 0; #(#to_compact)* @@ -92,7 +93,11 @@ pub fn generate_from_to( } /// Generates code to implement the `Compact` trait method `to_compact`. -fn generate_from_compact(fields: &FieldList, ident: &Ident, is_zstd: bool) -> TokenStream2 { +fn generate_from_compact( + fields: &FieldList, + ident: &Ident, + zstd: Option, +) -> TokenStream2 { let mut lines = vec![]; let mut known_types = vec!["B256", "Address", "Bloom", "Vec", "TxHash", "BlockHash", "FixedBytes"]; @@ -147,38 +152,41 @@ fn generate_from_compact(fields: &FieldList, ident: &Ident, is_zstd: bool) -> To // If the type has compression support, then check the `__zstd` flag. Otherwise, use the default // code branch. However, even if it's a type with compression support, not all values are // to be compressed (thus the zstd flag). Ideally only the bigger ones. - is_zstd - .then(|| { - let decompressor = format_ident!("{}_DECOMPRESSOR", ident.to_string().to_uppercase()); - quote! { - if flags.__zstd() != 0 { - #decompressor.with(|decompressor| { - let decompressor = &mut decompressor.borrow_mut(); - let decompressed = decompressor.decompress(buf); - let mut original_buf = buf; - - let mut buf: &[u8] = decompressed; - #(#lines)* - (obj, original_buf) - }) - } else { + if let Some(zstd) = zstd { + let decompressor = zstd.decompressor; + quote! { + if flags.__zstd() != 0 { + #decompressor.with(|decompressor| { + let decompressor = &mut decompressor.borrow_mut(); + let decompressed = decompressor.decompress(buf); + let mut original_buf = buf; + + let mut buf: &[u8] = decompressed; #(#lines)* - (obj, buf) - } - } - }) - .unwrap_or_else(|| { - quote! { + (obj, original_buf) + }) + } else { #(#lines)* (obj, buf) } - }) + } + } else { + quote! { + #(#lines)* + (obj, buf) + } + } } /// Generates code to implement the `Compact` trait method `from_compact`. -fn generate_to_compact(fields: &FieldList, ident: &Ident, is_zstd: bool) -> Vec { +fn generate_to_compact( + fields: &FieldList, + ident: &Ident, + zstd: Option, + reth_codecs: &syn::Path, +) -> Vec { let mut lines = vec![quote! { - let mut buffer = bytes::BytesMut::new(); + let mut buffer = #reth_codecs::__private::bytes::BytesMut::new(); }]; let is_enum = fields.iter().any(|field| matches!(field, FieldTypes::EnumVariant(_))); @@ -198,7 +206,7 @@ fn generate_to_compact(fields: &FieldList, ident: &Ident, is_zstd: bool) -> Vec< // Just because a type supports compression, doesn't mean all its values are to be compressed. // We skip the smaller ones, and thus require a flag` __zstd` to specify if this value is // compressed or not. - if is_zstd { + if zstd.is_some() { lines.push(quote! { let mut zstd = buffer.len() > 7; if zstd { @@ -214,9 +222,8 @@ fn generate_to_compact(fields: &FieldList, ident: &Ident, is_zstd: bool) -> Vec< buf.put_slice(&flags); }); - if is_zstd { - let compressor = format_ident!("{}_COMPRESSOR", ident.to_string().to_uppercase()); - + if let Some(zstd) = zstd { + let compressor = zstd.compressor; lines.push(quote! { if zstd { #compressor.with(|compressor| { diff --git a/crates/storage/codecs/derive/src/compact/mod.rs b/crates/storage/codecs/derive/src/compact/mod.rs index 1c1723d2ec94..1525809dd1fa 100644 --- a/crates/storage/codecs/derive/src/compact/mod.rs +++ b/crates/storage/codecs/derive/src/compact/mod.rs @@ -15,6 +15,8 @@ use flags::*; mod structs; use structs::*; +use crate::ZstdConfig; + // Helper Alias type type IsCompact = bool; // Helper Alias type @@ -40,7 +42,7 @@ pub enum FieldTypes { } /// Derives the `Compact` trait and its from/to implementations. -pub fn derive(input: TokenStream, is_zstd: bool) -> TokenStream { +pub fn derive(input: TokenStream, zstd: Option) -> TokenStream { let mut output = quote! {}; let DeriveInput { ident, data, generics, attrs, .. } = parse_macro_input!(input); @@ -48,8 +50,8 @@ pub fn derive(input: TokenStream, is_zstd: bool) -> TokenStream { let has_lifetime = has_lifetime(&generics); let fields = get_fields(&data); - output.extend(generate_flag_struct(&ident, &attrs, has_lifetime, &fields, is_zstd)); - output.extend(generate_from_to(&ident, &attrs, has_lifetime, &fields, is_zstd)); + output.extend(generate_flag_struct(&ident, &attrs, has_lifetime, &fields, zstd.is_some())); + output.extend(generate_from_to(&ident, &attrs, has_lifetime, &fields, zstd)); output.into() } @@ -236,7 +238,7 @@ mod tests { let DeriveInput { ident, data, attrs, .. } = parse2(f_struct).unwrap(); let fields = get_fields(&data); output.extend(generate_flag_struct(&ident, &attrs, false, &fields, false)); - output.extend(generate_from_to(&ident, &attrs, false, &fields, false)); + output.extend(generate_from_to(&ident, &attrs, false, &fields, None)); // Expected output in a TokenStream format. Commas matter! let should_output = quote! { diff --git a/crates/storage/codecs/derive/src/lib.rs b/crates/storage/codecs/derive/src/lib.rs index 0b4015830f5d..a2799ecda8c9 100644 --- a/crates/storage/codecs/derive/src/lib.rs +++ b/crates/storage/codecs/derive/src/lib.rs @@ -20,6 +20,12 @@ use syn::{ mod arbitrary; mod compact; +#[derive(Clone)] +pub(crate) struct ZstdConfig { + compressor: syn::Path, + decompressor: syn::Path, +} + /// Derives the `Compact` trait for custom structs, optimizing serialization with a possible /// bitflag struct. /// @@ -51,15 +57,49 @@ mod compact; /// efficient decoding. #[proc_macro_derive(Compact, attributes(maybe_zero, reth_codecs))] pub fn derive(input: TokenStream) -> TokenStream { - let is_zstd = false; - compact::derive(input, is_zstd) + compact::derive(input, None) } /// Adds `zstd` compression to derived [`Compact`]. -#[proc_macro_derive(CompactZstd, attributes(maybe_zero, reth_codecs))] +#[proc_macro_derive(CompactZstd, attributes(maybe_zero, reth_codecs, reth_zstd))] pub fn derive_zstd(input: TokenStream) -> TokenStream { - let is_zstd = true; - compact::derive(input, is_zstd) + let derive_input = { + let input = input.clone(); + parse_macro_input!(input as DeriveInput) + }; + + let mut compressor = None; + let mut decompressor = None; + + for attr in derive_input.attrs { + if attr.path().is_ident("reth_zstd") { + if let Err(err) = attr.parse_nested_meta(|meta| { + if meta.path.is_ident("compressor") { + let value = meta.value()?; + let path: syn::Path = value.parse()?; + compressor = Some(path); + } else if meta.path.is_ident("decompressor") { + let value = meta.value()?; + let path: syn::Path = value.parse()?; + decompressor = Some(path); + } else { + return Err(meta.error("unsupported attribute")) + } + Ok(()) + }) { + return err.to_compile_error().into() + } + } + } + + let (Some(compressor), Some(decompressor)) = (compressor, decompressor) else { + return quote! { + compile_error!("missing compressor or decompressor attribute"); + } + .into() + }; + + compact::derive(input, Some(ZstdConfig { compressor, decompressor })) } /// Generates tests for given type. diff --git a/crates/storage/codecs/src/private.rs b/crates/storage/codecs/src/private.rs index 6f54d9c9ca82..440310ffd3df 100644 --- a/crates/storage/codecs/src/private.rs +++ b/crates/storage/codecs/src/private.rs @@ -1,3 +1,3 @@ pub use modular_bitfield; -pub use bytes::Buf; +pub use bytes::{self, Buf}; From e32b4ca9a7e750a73b0b819b977342f7421463d8 Mon Sep 17 00:00:00 2001 From: Arsenii Kulikov Date: Wed, 11 Dec 2024 01:26:15 +0400 Subject: [PATCH 2/3] wip --- crates/storage/codecs/derive/src/compact/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/storage/codecs/derive/src/compact/mod.rs b/crates/storage/codecs/derive/src/compact/mod.rs index 1525809dd1fa..816414a63cd9 100644 --- a/crates/storage/codecs/derive/src/compact/mod.rs +++ b/crates/storage/codecs/derive/src/compact/mod.rs @@ -300,10 +300,10 @@ mod tests { fuzz_test_test_struct(TestStruct::default()) } impl reth_codecs::Compact for TestStruct { - fn to_compact(&self, buf: &mut B) -> usize where B: bytes::BufMut + AsMut<[u8]> { + fn to_compact(&self, buf: &mut B) -> usize where B: reth_codecs::__private::bytes::BufMut + AsMut<[u8]> { let mut flags = TestStructFlags::default(); let mut total_length = 0; - let mut buffer = bytes::BytesMut::new(); + let mut buffer = reth_codecs::__private::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); From aed7fca037b34aeb1506ec1e4fb7609f12b381b9 Mon Sep 17 00:00:00 2001 From: Arsenii Kulikov Date: Wed, 11 Dec 2024 15:44:07 +0400 Subject: [PATCH 3/3] avoid clone --- crates/storage/codecs/derive/src/compact/mod.rs | 6 +++--- crates/storage/codecs/derive/src/lib.rs | 9 +++------ 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/crates/storage/codecs/derive/src/compact/mod.rs b/crates/storage/codecs/derive/src/compact/mod.rs index 816414a63cd9..ab9ed78e164e 100644 --- a/crates/storage/codecs/derive/src/compact/mod.rs +++ b/crates/storage/codecs/derive/src/compact/mod.rs @@ -1,7 +1,7 @@ use proc_macro::TokenStream; use proc_macro2::{Ident, TokenStream as TokenStream2}; use quote::{format_ident, quote}; -use syn::{parse_macro_input, Data, DeriveInput, Generics}; +use syn::{Data, DeriveInput, Generics}; mod generator; use generator::*; @@ -42,10 +42,10 @@ pub enum FieldTypes { } /// Derives the `Compact` trait and its from/to implementations. -pub fn derive(input: TokenStream, zstd: Option) -> TokenStream { +pub fn derive(input: DeriveInput, zstd: Option) -> TokenStream { let mut output = quote! {}; - let DeriveInput { ident, data, generics, attrs, .. } = parse_macro_input!(input); + let DeriveInput { ident, data, generics, attrs, .. } = input; let has_lifetime = has_lifetime(&generics); diff --git a/crates/storage/codecs/derive/src/lib.rs b/crates/storage/codecs/derive/src/lib.rs index a2799ecda8c9..a835e8fab3c2 100644 --- a/crates/storage/codecs/derive/src/lib.rs +++ b/crates/storage/codecs/derive/src/lib.rs @@ -57,21 +57,18 @@ pub(crate) struct ZstdConfig { /// efficient decoding. #[proc_macro_derive(Compact, attributes(maybe_zero, reth_codecs))] pub fn derive(input: TokenStream) -> TokenStream { - compact::derive(input, None) + compact::derive(parse_macro_input!(input as DeriveInput), None) } /// Adds `zstd` compression to derived [`Compact`]. #[proc_macro_derive(CompactZstd, attributes(maybe_zero, reth_codecs, reth_zstd))] pub fn derive_zstd(input: TokenStream) -> TokenStream { - let derive_input = { - let input = input.clone(); - parse_macro_input!(input as DeriveInput) - }; + let input = parse_macro_input!(input as DeriveInput); let mut compressor = None; let mut decompressor = None; - for attr in derive_input.attrs { + for attr in &input.attrs { if attr.path().is_ident("reth_zstd") { if let Err(err) = attr.parse_nested_meta(|meta| { if meta.path.is_ident("compressor") {