feat(db): AccountsTrieChangeSets and StoragesTrieChangeSets tables

crates/storage/db-api/src/models/accounts.rs

@@ -7,7 +7,7 @@ use crate::{
     table::{Decode, Encode},
     DatabaseError,
 };
-use alloy_primitives::{Address, BlockNumber, StorageKey};
+use alloy_primitives::{Address, BlockNumber, StorageKey, B256};
 use serde::{Deserialize, Serialize};
 
 /// [`BlockNumber`] concatenated with [`Address`].
@@ -71,6 +71,43 @@ impl Decode for BlockNumberAddress {
     }
 }
 
+/// [`BlockNumber`] concatenated with [`B256`] (hashed address).
+///
+/// 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 BlockNumberHashedAddress(pub (BlockNumber, B256));
+
+impl From<(BlockNumber, B256)> for BlockNumberHashedAddress {
+    fn from(tpl: (BlockNumber, B256)) -> Self {
+        Self(tpl)
+    }
+}
+
+impl Encode for BlockNumberHashedAddress {
+    type Encoded = [u8; 40];
+
+    fn encode(self) -> Self::Encoded {
+        let block_number = self.0 .0;
+        let hashed_address = self.0 .1;
+
+        let mut buf = [0u8; 40];
+
+        buf[..8].copy_from_slice(&block_number.to_be_bytes());
+        buf[8..].copy_from_slice(hashed_address.as_slice());
+        buf
+    }
+}
+
+impl Decode for BlockNumberHashedAddress {
+    fn decode(value: &[u8]) -> Result<Self, DatabaseError> {
+        let num = u64::from_be_bytes(value[..8].try_into().map_err(|_| DatabaseError::Decode)?);
+        let hash = B256::from_slice(&value[8..]);
+        Ok(Self((num, hash)))
+    }
+}
+
 /// [`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.
@@ -102,7 +139,11 @@ impl Decode for AddressStorageKey {
     }
 }
 
-impl_fixed_arbitrary!((BlockNumberAddress, 28), (AddressStorageKey, 52));
+impl_fixed_arbitrary!(
+    (BlockNumberAddress, 28),
+    (BlockNumberHashedAddress, 40),
+    (AddressStorageKey, 52)
+);
 
 #[cfg(test)]
 mod tests {
@@ -135,6 +176,31 @@ mod tests {
         assert_eq!(bytes, Encode::encode(key));
     }
 
+    #[test]
+    fn test_block_number_hashed_address() {
+        let num = 1u64;
+        let hash = B256::from_slice(&[0xba; 32]);
+        let key = BlockNumberHashedAddress((num, hash));
+
+        let mut bytes = [0u8; 40];
+        bytes[..8].copy_from_slice(&num.to_be_bytes());
+        bytes[8..].copy_from_slice(hash.as_slice());
+
+        let encoded = Encode::encode(key);
+        assert_eq!(encoded, bytes);
+
+        let decoded: BlockNumberHashedAddress = Decode::decode(&encoded).unwrap();
+        assert_eq!(decoded, key);
+    }
+
+    #[test]
+    fn test_block_number_hashed_address_rand() {
+        let mut bytes = [0u8; 40];
+        rng().fill(bytes.as_mut_slice());
+        let key = BlockNumberHashedAddress::arbitrary(&mut Unstructured::new(&bytes)).unwrap();
+        assert_eq!(bytes, Encode::encode(key));
+    }
+
     #[test]
     fn test_address_storage_key() {
         let storage_key = StorageKey::random();

crates/storage/db-api/src/models/mod.rs

@@ -12,7 +12,9 @@ use reth_ethereum_primitives::{Receipt, TransactionSigned, TxType};
 use reth_primitives_traits::{Account, Bytecode, StorageEntry};
 use reth_prune_types::{PruneCheckpoint, PruneSegment};
 use reth_stages_types::StageCheckpoint;
-use reth_trie_common::{StoredNibbles, StoredNibblesSubKey, *};
+use reth_trie_common::{
+    StorageTrieEntry, StoredNibbles, StoredNibblesSubKey, TrieChangeSetsEntry, *,
+};
 use serde::{Deserialize, Serialize};
 
 pub mod accounts;
@@ -219,6 +221,7 @@ impl_compression_for_compact!(
     TxType,
     StorageEntry,
     BranchNodeCompact,
+    TrieChangeSetsEntry,
     StoredNibbles,
     StoredNibblesSubKey,
     StorageTrieEntry,

crates/storage/db-api/src/tables/mod.rs

@@ -21,8 +21,8 @@ use crate::{
         accounts::BlockNumberAddress,
         blocks::{HeaderHash, StoredBlockOmmers},
         storage_sharded_key::StorageShardedKey,
-        AccountBeforeTx, ClientVersion, CompactU256, IntegerList, ShardedKey,
-        StoredBlockBodyIndices, StoredBlockWithdrawals,
+        AccountBeforeTx, BlockNumberHashedAddress, ClientVersion, CompactU256, IntegerList,
+        ShardedKey, StoredBlockBodyIndices, StoredBlockWithdrawals,
     },
     table::{Decode, DupSort, Encode, Table, TableInfo},
 };
@@ -32,7 +32,9 @@ use reth_ethereum_primitives::{Receipt, TransactionSigned};
 use reth_primitives_traits::{Account, Bytecode, StorageEntry};
 use reth_prune_types::{PruneCheckpoint, PruneSegment};
 use reth_stages_types::StageCheckpoint;
-use reth_trie_common::{BranchNodeCompact, StorageTrieEntry, StoredNibbles, StoredNibblesSubKey};
+use reth_trie_common::{
+    BranchNodeCompact, StorageTrieEntry, StoredNibbles, StoredNibblesSubKey, TrieChangeSetsEntry,
+};
 use serde::{Deserialize, Serialize};
 use std::fmt;
 
@@ -486,6 +488,20 @@ tables! {
         type SubKey = StoredNibblesSubKey;
     }
 
+    /// Stores the state of a node in the accounts trie prior to a particular block being executed.
+    table AccountsTrieChangeSets {
+        type Key = BlockNumber;
+        type Value = TrieChangeSetsEntry;
+        type SubKey = StoredNibblesSubKey;
+    }
+
+    /// Stores the state of a node in a storage trie prior to a particular block being executed.
+    table StoragesTrieChangeSets {
+        type Key = BlockNumberHashedAddress;
+        type Value = TrieChangeSetsEntry;
+        type SubKey = StoredNibblesSubKey;
+    }
+
     /// Stores the transaction sender for each canonical transaction.
     /// It is needed to speed up execution stage and allows fetching signer without doing
     /// transaction signed recovery

crates/trie/common/src/lib.rs

@@ -36,7 +36,7 @@ mod nibbles;
 pub use nibbles::{Nibbles, StoredNibbles, StoredNibblesSubKey};
 
 mod storage;
-pub use storage::StorageTrieEntry;
+pub use storage::{StorageTrieEntry, TrieChangeSetsEntry};
 
 mod subnode;
 pub use subnode::StoredSubNode;

crates/trie/common/src/storage.rs

@@ -1,6 +1,8 @@
-use super::{BranchNodeCompact, StoredNibblesSubKey};
+use super::{BranchNodeCompact, Nibbles, StoredNibblesSubKey};
 
 /// Account storage trie node.
+///
+/// `nibbles` is the subkey when used as a value in the `StorageTrie` table.
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 #[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
 pub struct StorageTrieEntry {
@@ -31,3 +33,173 @@ impl reth_codecs::Compact for StorageTrieEntry {
         (this, buf)
     }
 }
+
+/// Trie changeset entry representing the state of a trie node before a block.
+///
+/// `nibbles` is the subkey when used as a value in the changeset tables.
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
+pub struct TrieChangeSetsEntry {
+    /// The nibbles of the intermediate node
+    pub nibbles: StoredNibblesSubKey,
+    /// Node value prior to the block being processed, None indicating it didn't exist.
+    pub node: Option<BranchNodeCompact>,
+}
+
+#[cfg(any(test, feature = "reth-codec"))]
+impl reth_codecs::Compact for TrieChangeSetsEntry {
+    fn to_compact<B>(&self, buf: &mut B) -> usize
+    where
+        B: bytes::BufMut + AsMut<[u8]>,
+    {
+        let nibbles_len = self.nibbles.to_compact(buf);
+        let node_len = self.node.as_ref().map(|node| node.to_compact(buf)).unwrap_or(0);
+        nibbles_len + node_len
+    }
+
+    fn from_compact(buf: &[u8], len: usize) -> (Self, &[u8]) {
+        if len == 0 {
+            // Return an empty entry without trying to parse anything
+            return (
+                Self { nibbles: StoredNibblesSubKey::from(Nibbles::default()), node: None },
+                buf,
+            )
+        }
+
+        let (nibbles, buf) = StoredNibblesSubKey::from_compact(buf, 65);
+
+        if len <= 65 {
+            return (Self { nibbles, node: None }, buf)
+        }
+
+        let (node, buf) = BranchNodeCompact::from_compact(buf, len - 65);
+        (Self { nibbles, node: Some(node) }, buf)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use bytes::BytesMut;
+    use reth_codecs::Compact;
+
+    #[test]
+    fn test_trie_changesets_entry_full_empty() {
+        // Test a fully empty entry (empty nibbles, None node)
+        let entry = TrieChangeSetsEntry { nibbles: StoredNibblesSubKey::from(vec![]), node: None };
+
+        let mut buf = BytesMut::new();
+        let len = entry.to_compact(&mut buf);
+
+        // Empty nibbles takes 65 bytes (64 for padding + 1 for length)
+        // None node adds 0 bytes
+        assert_eq!(len, 65);
+        assert_eq!(buf.len(), 65);
+
+        // Deserialize and verify
+        let (decoded, remaining) = TrieChangeSetsEntry::from_compact(&buf, len);
+        assert_eq!(decoded.nibbles.0.to_vec(), Vec::<u8>::new());
+        assert_eq!(decoded.node, None);
+        assert_eq!(remaining.len(), 0);
+    }
+
+    #[test]
+    fn test_trie_changesets_entry_none_node() {
+        // Test non-empty nibbles with None node
+        let nibbles_data = vec![0x01, 0x02, 0x03, 0x04];
+        let entry = TrieChangeSetsEntry {
+            nibbles: StoredNibblesSubKey::from(nibbles_data.clone()),
+            node: None,
+        };
+
+        let mut buf = BytesMut::new();
+        let len = entry.to_compact(&mut buf);
+
+        // Nibbles takes 65 bytes regardless of content
+        assert_eq!(len, 65);
+
+        // Deserialize and verify
+        let (decoded, remaining) = TrieChangeSetsEntry::from_compact(&buf, len);
+        assert_eq!(decoded.nibbles.0.to_vec(), nibbles_data);
+        assert_eq!(decoded.node, None);
+        assert_eq!(remaining.len(), 0);
+    }
+
+    #[test]
+    fn test_trie_changesets_entry_empty_path_with_node() {
+        // Test empty path with Some node
+        // Using the same signature as in the codebase: (state_mask, hash_mask, tree_mask, hashes,
+        // value)
+        let test_node = BranchNodeCompact::new(
+            0b1111_1111_1111_1111, // state_mask: all children present
+            0b1111_1111_1111_1111, // hash_mask: all have hashes
+            0b0000_0000_0000_0000, // tree_mask: no embedded trees
+            vec![],                // hashes
+            None,                  // value
+        );
+
+        let entry = TrieChangeSetsEntry {
+            nibbles: StoredNibblesSubKey::from(vec![]),
+            node: Some(test_node.clone()),
+        };
+
+        let mut buf = BytesMut::new();
+        let len = entry.to_compact(&mut buf);
+
+        // Calculate expected length
+        let mut temp_buf = BytesMut::new();
+        let node_len = test_node.to_compact(&mut temp_buf);
+        assert_eq!(len, 65 + node_len);
+
+        // Deserialize and verify
+        let (decoded, remaining) = TrieChangeSetsEntry::from_compact(&buf, len);
+        assert_eq!(decoded.nibbles.0.to_vec(), Vec::<u8>::new());
+        assert_eq!(decoded.node, Some(test_node));
+        assert_eq!(remaining.len(), 0);
+    }
+
+    #[test]
+    fn test_trie_changesets_entry_normal() {
+        // Test normal case: non-empty path with Some node
+        let nibbles_data = vec![0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f];
+        // Using the same signature as in the codebase
+        let test_node = BranchNodeCompact::new(
+            0b0000_0000_1111_0000, // state_mask: some children present
+            0b0000_0000_0011_0000, // hash_mask: some have hashes
+            0b0000_0000_0000_0000, // tree_mask: no embedded trees
+            vec![],                // hashes (empty for this test)
+            None,                  // value
+        );
+
+        let entry = TrieChangeSetsEntry {
+            nibbles: StoredNibblesSubKey::from(nibbles_data.clone()),
+            node: Some(test_node.clone()),
+        };
+
+        let mut buf = BytesMut::new();
+        let len = entry.to_compact(&mut buf);
+
+        // Verify serialization length
+        let mut temp_buf = BytesMut::new();
+        let node_len = test_node.to_compact(&mut temp_buf);
+        assert_eq!(len, 65 + node_len);
+
+        // Deserialize and verify
+        let (decoded, remaining) = TrieChangeSetsEntry::from_compact(&buf, len);
+        assert_eq!(decoded.nibbles.0.to_vec(), nibbles_data);
+        assert_eq!(decoded.node, Some(test_node));
+        assert_eq!(remaining.len(), 0);
+    }
+
+    #[test]
+    fn test_trie_changesets_entry_from_compact_zero_len() {
+        // Test from_compact with zero length
+        let buf = vec![0x01, 0x02, 0x03];
+        let (decoded, remaining) = TrieChangeSetsEntry::from_compact(&buf, 0);
+
+        // Should return empty nibbles and None node
+        assert_eq!(decoded.nibbles.0.to_vec(), Vec::<u8>::new());
+        assert_eq!(decoded.node, None);
+        assert_eq!(remaining, &buf[..]); // Buffer should be unchanged
+    }
+}