feat: add ETL to HistoryStages (#7249)

crates/storage/db/src/abstraction/table.rs

@@ -16,7 +16,8 @@ pub trait Compress: Send + Sync + Sized + Debug {
         + Into<Vec<u8>>
         + Default
         + Send
-        + Sync;
+        + Sync
+        + Debug;
 
     /// If the type cannot be compressed, return its inner reference as `Some(self.as_ref())`
     fn uncompressable_ref(&self) -> Option<&[u8]> {
@@ -48,7 +49,7 @@ pub trait Decompress: Send + Sync + Sized + Debug {
 /// Trait that will transform the data to be saved in the DB.
 pub trait Encode: Send + Sync + Sized + Debug {
     /// Encoded type.
-    type Encoded: AsRef<[u8]> + Into<Vec<u8>> + Send + Sync;
+    type Encoded: AsRef<[u8]> + Into<Vec<u8>> + Send + Sync + Ord + Debug;
 
     /// Encodes data going into the database.
     fn encode(self) -> Self::Encoded;

crates/storage/db/src/tables/models/accounts.rs

@@ -8,7 +8,7 @@ use crate::{
     DatabaseError,
 };
 use reth_codecs::{derive_arbitrary, Compact};
-use reth_primitives::{Account, Address, BlockNumber, Buf};
+use reth_primitives::{Account, Address, BlockNumber, Buf, StorageKey};
 use serde::{Deserialize, Serialize};
 
 /// Account as it is saved inside [`AccountChangeSets`][crate::tables::AccountChangeSets].
@@ -121,7 +121,40 @@ impl Decode for BlockNumberAddress {
     }
 }
 
-impl_fixed_arbitrary!(BlockNumberAddress, 28);
+/// [`Address`] concatenated with [`StorageKey`]. Used by `reth_etl` and history stages.
+///
+/// Since it's used as a key, it isn't compressed when encoding it.
+#[derive(
+    Debug, Default, Copy, Clone, PartialEq, Eq, Serialize, Deserialize, Ord, PartialOrd, Hash,
+)]
+pub struct AddressStorageKey(pub (Address, StorageKey));
+
+impl Encode for AddressStorageKey {
+    type Encoded = [u8; 52];
+
+    fn encode(self) -> Self::Encoded {
+        let address = self.0 .0;
+        let storage_key = self.0 .1;
+
+        let mut buf = [0u8; 52];
+
+        buf[..20].copy_from_slice(address.as_slice());
+        buf[20..].copy_from_slice(storage_key.as_slice());
+        buf
+    }
+}
+
+impl Decode for AddressStorageKey {
+    fn decode<B: AsRef<[u8]>>(value: B) -> Result<Self, DatabaseError> {
+        let value = value.as_ref();
+        let address = Address::from_slice(&value[..20]);
+        let storage_key = StorageKey::from_slice(&value[20..]);
+
+        Ok(AddressStorageKey((address, storage_key)))
+    }
+}
+
+impl_fixed_arbitrary!((BlockNumberAddress, 28), (AddressStorageKey, 52));
 
 #[cfg(test)]
 mod tests {
@@ -153,4 +186,29 @@ mod tests {
         let key = BlockNumberAddress::arbitrary(&mut Unstructured::new(&bytes)).unwrap();
         assert_eq!(bytes, Encode::encode(key));
     }
+
+    #[test]
+    fn test_address_storage_key() {
+        let storage_key = StorageKey::random();
+        let address = Address::from_str("ba5e000000000000000000000000000000000000").unwrap();
+        let key = AddressStorageKey((address, storage_key));
+
+        let mut bytes = [0u8; 52];
+        bytes[..20].copy_from_slice(address.as_slice());
+        bytes[20..].copy_from_slice(storage_key.as_slice());
+
+        let encoded = Encode::encode(key);
+        assert_eq!(encoded, bytes);
+
+        let decoded: AddressStorageKey = Decode::decode(encoded).unwrap();
+        assert_eq!(decoded, key);
+    }
+
+    #[test]
+    fn test_address_storage_key_rand() {
+        let mut bytes = [0u8; 52];
+        thread_rng().fill(bytes.as_mut_slice());
+        let key = AddressStorageKey::arbitrary(&mut Unstructured::new(&bytes)).unwrap();
+        assert_eq!(bytes, Encode::encode(key));
+    }
 }

crates/storage/db/src/tables/utils.rs

@@ -9,34 +9,35 @@ use std::borrow::Cow;
 #[macro_export]
 /// Implements the `Arbitrary` trait for types with fixed array types.
 macro_rules! impl_fixed_arbitrary {
-    ($name:ident, $size:tt) => {
+    ($(($name:ident, $size:expr)),*) => {
         #[cfg(any(test, feature = "arbitrary"))]
         use arbitrary::{Arbitrary, Unstructured};
+        $(
+            #[cfg(any(test, feature = "arbitrary"))]
+            impl<'a> Arbitrary<'a> for $name {
+                fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self, arbitrary::Error> {
+                    let mut buffer = vec![0; $size];
+                    u.fill_buffer(buffer.as_mut_slice())?;
 
-        #[cfg(any(test, feature = "arbitrary"))]
-        impl<'a> Arbitrary<'a> for $name {
-            fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self, arbitrary::Error> {
-                let mut buffer = vec![0; $size];
-                u.fill_buffer(buffer.as_mut_slice())?;
-
-                Decode::decode(buffer).map_err(|_| arbitrary::Error::IncorrectFormat)
+                    Decode::decode(buffer).map_err(|_| arbitrary::Error::IncorrectFormat)
+                }
             }
-        }
 
-        #[cfg(any(test, feature = "arbitrary"))]
-        impl proptest::prelude::Arbitrary for $name {
-            type Parameters = ();
-            type Strategy = proptest::strategy::Map<
-                proptest::collection::VecStrategy<<u8 as proptest::arbitrary::Arbitrary>::Strategy>,
-                fn(Vec<u8>) -> Self,
-            >;
+            #[cfg(any(test, feature = "arbitrary"))]
+            impl proptest::prelude::Arbitrary for $name {
+                type Parameters = ();
+                type Strategy = proptest::strategy::Map<
+                    proptest::collection::VecStrategy<<u8 as proptest::arbitrary::Arbitrary>::Strategy>,
+                    fn(Vec<u8>) -> Self,
+                >;
 
-            fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
-                use proptest::strategy::Strategy;
-                proptest::collection::vec(proptest::arbitrary::any_with::<u8>(args), $size)
-                    .prop_map(move |vec| Decode::decode(vec).unwrap())
+                fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
+                    use proptest::strategy::Strategy;
+                    proptest::collection::vec(proptest::arbitrary::any_with::<u8>(args), $size)
+                        .prop_map(move |vec| Decode::decode(vec).unwrap())
+                }
             }
-        }
+        )+
     };
 }