diff --git a/bin/reth/src/debug_cmd/in_memory_merkle.rs b/bin/reth/src/debug_cmd/in_memory_merkle.rs index 9e6c9bd9d9..72d3ef97b6 100644 --- a/bin/reth/src/debug_cmd/in_memory_merkle.rs +++ b/bin/reth/src/debug_cmd/in_memory_merkle.rs @@ -181,7 +181,7 @@ impl Command { let (account_prefix_set, storage_prefix_set) = hashed_post_state.construct_prefix_sets(); let tx = provider.tx_ref(); let hashed_cursor_factory = HashedPostStateCursorFactory::new(tx, &hashed_post_state); - let (in_memory_state_root, in_memory_updates) = StateRoot::new(tx) + let (in_memory_state_root, in_memory_updates) = StateRoot::from_tx(tx) .with_hashed_cursor_factory(hashed_cursor_factory) .with_changed_account_prefixes(account_prefix_set) .with_changed_storage_prefixes(storage_prefix_set) diff --git a/bin/reth/src/recover/storage_tries.rs b/bin/reth/src/recover/storage_tries.rs index f1ee84474d..c9069203f5 100644 --- a/bin/reth/src/recover/storage_tries.rs +++ b/bin/reth/src/recover/storage_tries.rs @@ -76,7 +76,7 @@ impl Command { entry = storage_trie_cursor.next()?; } - let state_root = StateRoot::new(tx_mut).root()?; + let state_root = StateRoot::from_tx(tx_mut).root()?; if state_root != best_header.state_root { eyre::bail!( "Recovery failed. Incorrect state root. Expected: {:?}. Received: {:?}", diff --git a/crates/stages/benches/setup/mod.rs b/crates/stages/benches/setup/mod.rs index 17460579c5..a4422e23f4 100644 --- a/crates/stages/benches/setup/mod.rs +++ b/crates/stages/benches/setup/mod.rs @@ -126,8 +126,9 @@ pub(crate) fn txs_testdata(num_blocks: u64) -> PathBuf { db.insert_accounts_and_storages(start_state.clone()).unwrap(); // make first block after genesis have valid state root - let (root, updates) = - StateRoot::new(db.factory.provider_rw().unwrap().tx_ref()).root_with_updates().unwrap(); + let (root, updates) = StateRoot::from_tx(db.factory.provider_rw().unwrap().tx_ref()) + .root_with_updates() + .unwrap(); let second_block = blocks.get_mut(1).unwrap(); let cloned_second = second_block.clone(); let mut updated_header = cloned_second.header.unseal(); @@ -154,7 +155,7 @@ pub(crate) fn txs_testdata(num_blocks: u64) -> PathBuf { // make last block have valid state root let root = { let tx_mut = db.factory.provider_rw().unwrap(); - let root = StateRoot::new(tx_mut.tx_ref()).root().unwrap(); + let root = StateRoot::from_tx(tx_mut.tx_ref()).root().unwrap(); tx_mut.commit().unwrap(); root }; diff --git a/crates/stages/src/stages/merkle.rs b/crates/stages/src/stages/merkle.rs index 518ceac92e..fadee1107d 100644 --- a/crates/stages/src/stages/merkle.rs +++ b/crates/stages/src/stages/merkle.rs @@ -191,7 +191,7 @@ impl Stage for MerkleStage { }); let tx = provider.tx_ref(); - let progress = StateRoot::new(tx) + let progress = StateRoot::from_tx(tx) .with_intermediate_state(checkpoint.map(IntermediateStateRootState::from)) .root_with_progress() .map_err(|e| StageError::Fatal(Box::new(e)))?; diff --git a/crates/storage/provider/src/bundle_state/bundle_state_with_receipts.rs b/crates/storage/provider/src/bundle_state/bundle_state_with_receipts.rs index d196642894..0408afd090 100644 --- a/crates/storage/provider/src/bundle_state/bundle_state_with_receipts.rs +++ b/crates/storage/provider/src/bundle_state/bundle_state_with_receipts.rs @@ -159,10 +159,10 @@ impl BundleStateWithReceipts { &self, tx: &'a TX, hashed_post_state: &'b HashedPostState, - ) -> StateRoot<'a, TX, HashedPostStateCursorFactory<'a, 'b, TX>> { + ) -> StateRoot<&'a TX, HashedPostStateCursorFactory<'a, 'b, TX>> { let (account_prefix_set, storage_prefix_set) = hashed_post_state.construct_prefix_sets(); let hashed_cursor_factory = HashedPostStateCursorFactory::new(tx, hashed_post_state); - StateRoot::new(tx) + StateRoot::from_tx(tx) .with_hashed_cursor_factory(hashed_cursor_factory) .with_changed_account_prefixes(account_prefix_set) .with_changed_storage_prefixes(storage_prefix_set) @@ -1236,7 +1236,7 @@ mod tests { } } - let (_, updates) = StateRoot::new(tx).root_with_updates().unwrap(); + let (_, updates) = StateRoot::from_tx(tx).root_with_updates().unwrap(); updates.flush(tx).unwrap(); }) .unwrap(); diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs index 933ddb0123..4a2f4c03c7 100644 --- a/crates/storage/provider/src/providers/database/provider.rs +++ b/crates/storage/provider/src/providers/database/provider.rs @@ -2077,7 +2077,7 @@ impl HashingWriter for DatabaseProvider { { // This is the same as `StateRoot::incremental_root_with_updates`, only the prefix sets // are pre-loaded. - let (state_root, trie_updates) = StateRoot::new(&self.tx) + let (state_root, trie_updates) = StateRoot::from_tx(&self.tx) .with_changed_account_prefixes(account_prefix_set.freeze()) .with_changed_storage_prefixes( storage_prefix_set.into_iter().map(|(k, v)| (k, v.freeze())).collect(), @@ -2261,7 +2261,7 @@ impl BlockExecutionWriter for DatabaseProvider { // Calculate the reverted merkle root. // This is the same as `StateRoot::incremental_root_with_updates`, only the prefix sets // are pre-loaded. - let (new_state_root, trie_updates) = StateRoot::new(&self.tx) + let (new_state_root, trie_updates) = StateRoot::from_tx(&self.tx) .with_changed_account_prefixes(account_prefix_set.freeze()) .with_changed_storage_prefixes( storage_prefix_set.into_iter().map(|(k, v)| (k, v.freeze())).collect(), diff --git a/crates/trie/src/proof.rs b/crates/trie/src/proof.rs index 874c241e29..68f8767598 100644 --- a/crates/trie/src/proof.rs +++ b/crates/trie/src/proof.rs @@ -2,7 +2,7 @@ use crate::{ hashed_cursor::{HashedCursorFactory, HashedStorageCursor}, node_iter::{AccountNode, AccountNodeIter, StorageNode, StorageNodeIter}, prefix_set::PrefixSetMut, - trie_cursor::{AccountTrieCursor, StorageTrieCursor}, + trie_cursor::{DatabaseAccountTrieCursor, DatabaseStorageTrieCursor}, walker::TrieWalker, StateRootError, StorageRootError, }; @@ -51,7 +51,8 @@ where let mut account_proof = AccountProof::new(address); let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?; - let trie_cursor = AccountTrieCursor::new(self.tx.cursor_read::()?); + let trie_cursor = + DatabaseAccountTrieCursor::new(self.tx.cursor_read::()?); // Create the walker. let mut prefix_set = PrefixSetMut::default(); @@ -119,7 +120,7 @@ where let target_nibbles = proofs.iter().map(|p| p.nibbles.clone()).collect::>(); let prefix_set = PrefixSetMut::from(target_nibbles.clone()).freeze(); - let trie_cursor = StorageTrieCursor::new( + let trie_cursor = DatabaseStorageTrieCursor::new( self.tx.cursor_dup_read::()?, hashed_address, ); @@ -221,7 +222,7 @@ mod tests { }); provider.insert_storage_for_hashing(alloc_storage)?; - let (_, updates) = StateRoot::new(provider.tx_ref()) + let (_, updates) = StateRoot::from_tx(provider.tx_ref()) .root_with_updates() .map_err(Into::::into)?; updates.flush(provider.tx_mut())?; diff --git a/crates/trie/src/trie.rs b/crates/trie/src/trie.rs index b5053cbba6..7ef189cf67 100644 --- a/crates/trie/src/trie.rs +++ b/crates/trie/src/trie.rs @@ -3,13 +3,13 @@ use crate::{ node_iter::{AccountNode, AccountNodeIter, StorageNode, StorageNodeIter}, prefix_set::{PrefixSet, PrefixSetLoader, PrefixSetMut}, progress::{IntermediateStateRootState, StateRootProgress}, - trie_cursor::{AccountTrieCursor, StorageTrieCursor}, + trie_cursor::TrieCursorFactory, updates::{TrieKey, TrieOp, TrieUpdates}, walker::TrieWalker, StateRootError, StorageRootError, }; use alloy_rlp::{BufMut, Encodable}; -use reth_db::{tables, transaction::DbTx}; +use reth_db::transaction::DbTx; use reth_primitives::{ constants::EMPTY_ROOT_HASH, keccak256, @@ -24,9 +24,9 @@ use tracing::{debug, trace}; /// StateRoot is used to compute the root node of a state trie. #[derive(Debug)] -pub struct StateRoot<'a, TX, H> { - /// A reference to the database transaction. - pub tx: &'a TX, +pub struct StateRoot { + /// The factory for trie cursors. + pub trie_cursor_factory: T, /// The factory for hashed cursors. pub hashed_cursor_factory: H, /// A set of account prefixes that have changed. @@ -42,7 +42,7 @@ pub struct StateRoot<'a, TX, H> { threshold: u64, } -impl<'a, TX, H> StateRoot<'a, TX, H> { +impl StateRoot { /// Set the changed account prefixes. pub fn with_changed_account_prefixes(mut self, prefixes: PrefixSet) -> Self { self.changed_account_prefixes = prefixes; @@ -80,33 +80,43 @@ impl<'a, TX, H> StateRoot<'a, TX, H> { } /// Set the hashed cursor factory. - pub fn with_hashed_cursor_factory( - self, - hashed_cursor_factory: HF, - ) -> StateRoot<'a, TX, HF> { + pub fn with_hashed_cursor_factory(self, hashed_cursor_factory: HF) -> StateRoot { StateRoot { - tx: self.tx, + trie_cursor_factory: self.trie_cursor_factory, + hashed_cursor_factory, + changed_account_prefixes: self.changed_account_prefixes, + changed_storage_prefixes: self.changed_storage_prefixes, + destroyed_accounts: self.destroyed_accounts, + threshold: self.threshold, + previous_state: self.previous_state, + } + } + + /// Set the trie cursor factory. + pub fn with_trie_cursor_factory(self, trie_cursor_factory: TF) -> StateRoot { + StateRoot { + trie_cursor_factory, + hashed_cursor_factory: self.hashed_cursor_factory, changed_account_prefixes: self.changed_account_prefixes, changed_storage_prefixes: self.changed_storage_prefixes, destroyed_accounts: self.destroyed_accounts, threshold: self.threshold, previous_state: self.previous_state, - hashed_cursor_factory, } } } -impl<'a, TX: DbTx> StateRoot<'a, TX, &'a TX> { +impl<'a, TX: DbTx> StateRoot<&'a TX, &'a TX> { /// Create a new [StateRoot] instance. - pub fn new(tx: &'a TX) -> Self { + pub fn from_tx(tx: &'a TX) -> Self { Self { - tx, + trie_cursor_factory: tx, + hashed_cursor_factory: tx, changed_account_prefixes: PrefixSetMut::default().freeze(), changed_storage_prefixes: HashMap::default(), destroyed_accounts: HashSet::default(), previous_state: None, threshold: 100_000, - hashed_cursor_factory: tx, } } @@ -121,7 +131,7 @@ impl<'a, TX: DbTx> StateRoot<'a, TX, &'a TX> { range: RangeInclusive, ) -> Result { let loaded_prefix_sets = PrefixSetLoader::new(tx).load(range)?; - Ok(Self::new(tx) + Ok(Self::from_tx(tx) .with_changed_account_prefixes(loaded_prefix_sets.account_prefix_set.freeze()) .with_changed_storage_prefixes( loaded_prefix_sets @@ -178,9 +188,9 @@ impl<'a, TX: DbTx> StateRoot<'a, TX, &'a TX> { } } -impl<'a, TX, H> StateRoot<'a, TX, H> +impl StateRoot where - TX: DbTx, + T: TrieCursorFactory + Clone, H: HashedCursorFactory + Clone, { /// Walks the intermediate nodes of existing state trie (if any) and hashed entries. Feeds the @@ -226,7 +236,7 @@ where let mut trie_updates = TrieUpdates::default(); let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?; - let trie_cursor = AccountTrieCursor::new(self.tx.cursor_read::()?); + let trie_cursor = self.trie_cursor_factory.account_trie_cursor()?; let (mut hash_builder, mut account_node_iter) = match self.previous_state { Some(state) => { @@ -267,14 +277,17 @@ where // progress. // TODO: We can consider introducing the TrieProgress::Progress/Complete // abstraction inside StorageRoot, but let's give it a try as-is for now. - let storage_root_calculator = StorageRoot::new_hashed(self.tx, hashed_address) - .with_hashed_cursor_factory(self.hashed_cursor_factory.clone()) - .with_changed_prefixes( - self.changed_storage_prefixes - .get(&hashed_address) - .cloned() - .unwrap_or_default(), - ); + let storage_root_calculator = StorageRoot::new_hashed( + self.trie_cursor_factory.clone(), + self.hashed_cursor_factory.clone(), + hashed_address, + ) + .with_changed_prefixes( + self.changed_storage_prefixes + .get(&hashed_address) + .cloned() + .unwrap_or_default(), + ); let storage_root = if retain_updates { let (root, storage_slots_walked, updates) = @@ -336,9 +349,9 @@ where /// StorageRoot is used to compute the root node of an account storage trie. #[derive(Debug)] -pub struct StorageRoot<'a, TX, H> { +pub struct StorageRoot { /// A reference to the database transaction. - pub tx: &'a TX, + pub trie_cursor_factory: T, /// The factory for hashed cursors. pub hashed_cursor_factory: H, /// The hashed address of an account. @@ -347,40 +360,23 @@ pub struct StorageRoot<'a, TX, H> { pub changed_prefixes: PrefixSet, } -impl<'a, TX: DbTx> StorageRoot<'a, TX, &'a TX> { - /// Creates a new storage root calculator given an raw address. - pub fn new(tx: &'a TX, address: Address) -> Self { - Self::new_hashed(tx, keccak256(address)) +impl StorageRoot { + /// Creates a new storage root calculator given a raw address. + pub fn new(trie_cursor_factory: T, hashed_cursor_factory: H, address: Address) -> Self { + Self::new_hashed(trie_cursor_factory, hashed_cursor_factory, keccak256(address)) } /// Creates a new storage root calculator given a hashed address. - pub fn new_hashed(tx: &'a TX, hashed_address: B256) -> Self { - Self { - tx, - hashed_address, - changed_prefixes: PrefixSetMut::default().freeze(), - hashed_cursor_factory: tx, - } - } -} - -impl<'a, TX, H> StorageRoot<'a, TX, H> { - /// Creates a new storage root calculator given an raw address. - pub fn new_with_factory(tx: &'a TX, hashed_cursor_factory: H, address: Address) -> Self { - Self::new_hashed_with_factory(tx, hashed_cursor_factory, keccak256(address)) - } - - /// Creates a new storage root calculator given a hashed address. - pub fn new_hashed_with_factory( - tx: &'a TX, + pub fn new_hashed( + trie_cursor_factory: T, hashed_cursor_factory: H, hashed_address: B256, ) -> Self { Self { - tx, + trie_cursor_factory, + hashed_cursor_factory, hashed_address, changed_prefixes: PrefixSetMut::default().freeze(), - hashed_cursor_factory, } } @@ -391,22 +387,41 @@ impl<'a, TX, H> StorageRoot<'a, TX, H> { } /// Set the hashed cursor factory. - pub fn with_hashed_cursor_factory( - self, - hashed_cursor_factory: HF, - ) -> StorageRoot<'a, TX, HF> { + pub fn with_hashed_cursor_factory(self, hashed_cursor_factory: HF) -> StorageRoot { StorageRoot { - tx: self.tx, + trie_cursor_factory: self.trie_cursor_factory, + hashed_cursor_factory, + hashed_address: self.hashed_address, + changed_prefixes: self.changed_prefixes, + } + } + + /// Set the trie cursor factory. + pub fn with_trie_cursor_factory(self, trie_cursor_factory: TF) -> StorageRoot { + StorageRoot { + trie_cursor_factory, + hashed_cursor_factory: self.hashed_cursor_factory, hashed_address: self.hashed_address, changed_prefixes: self.changed_prefixes, - hashed_cursor_factory, } } } -impl<'a, TX, H> StorageRoot<'a, TX, H> +impl<'a, TX: DbTx> StorageRoot<&'a TX, &'a TX> { + /// Create a new storage root calculator from database transaction and raw address. + pub fn from_tx(tx: &'a TX, address: Address) -> Self { + Self::new(tx, tx, address) + } + + /// Create a new storage root calculator from database transaction and hashed address. + pub fn from_tx_hashed(tx: &'a TX, hashed_address: B256) -> Self { + Self::new_hashed(tx, tx, hashed_address) + } +} + +impl StorageRoot where - TX: DbTx, + T: TrieCursorFactory, H: HashedCursorFactory, { /// Walks the hashed storage table entries for a given address and calculates the storage root. @@ -444,10 +459,7 @@ where )) } - let trie_cursor = StorageTrieCursor::new( - self.tx.cursor_dup_read::()?, - self.hashed_address, - ); + let trie_cursor = self.trie_cursor_factory.storage_tries_cursor(self.hashed_address)?; let walker = TrieWalker::new(trie_cursor, self.changed_prefixes.clone()) .with_updates(retain_updates); @@ -545,7 +557,7 @@ mod tests { // Generate the intermediate nodes on the receiving end of the channel let (_, _, trie_updates) = - StorageRoot::new_hashed(tx.tx_ref(), hashed_address).root_with_updates().unwrap(); + StorageRoot::from_tx_hashed(tx.tx_ref(), hashed_address).root_with_updates().unwrap(); // 1. Some state transition happens, update the hashed storage to the new value let modified_key = B256::from_str(modified).unwrap(); @@ -559,7 +571,7 @@ mod tests { .unwrap(); // 2. Calculate full merkle root - let loader = StorageRoot::new_hashed(tx.tx_ref(), hashed_address); + let loader = StorageRoot::from_tx_hashed(tx.tx_ref(), hashed_address); let modified_root = loader.root().unwrap(); // Update the intermediate roots table so that we can run the incremental verification @@ -568,7 +580,7 @@ mod tests { // 3. Calculate the incremental root let mut storage_changes = PrefixSetMut::default(); storage_changes.insert(Nibbles::unpack(modified_key)); - let loader = StorageRoot::new_hashed(tx.tx_ref(), hashed_address) + let loader = StorageRoot::from_tx_hashed(tx.tx_ref(), hashed_address) .with_changed_prefixes(storage_changes.freeze()); let incremental_root = loader.root().unwrap(); @@ -608,7 +620,7 @@ mod tests { tx.commit().unwrap(); let tx = factory.provider_rw().unwrap(); - let got = StorageRoot::new(tx.tx_ref(), address).root().unwrap(); + let got = StorageRoot::from_tx(tx.tx_ref(), address).root().unwrap(); let expected = storage_root(storage.into_iter()); assert_eq!(expected, got); }); @@ -667,7 +679,7 @@ mod tests { tx.commit().unwrap(); let tx = factory.provider_rw().unwrap(); - let got = StorageRoot::new(tx.tx_ref(), address).root().unwrap(); + let got = StorageRoot::from_tx(tx.tx_ref(), address).root().unwrap(); assert_eq!(got, EMPTY_ROOT_HASH); } @@ -692,7 +704,7 @@ mod tests { tx.commit().unwrap(); let tx = factory.provider_rw().unwrap(); - let got = StorageRoot::new(tx.tx_ref(), address).root().unwrap(); + let got = StorageRoot::from_tx(tx.tx_ref(), address).root().unwrap(); assert_eq!(storage_root(storage.into_iter()), got); } @@ -732,7 +744,7 @@ mod tests { let mut intermediate_state: Option> = None; while got.is_none() { - let calculator = StateRoot::new(tx.tx_ref()) + let calculator = StateRoot::from_tx(tx.tx_ref()) .with_threshold(threshold) .with_intermediate_state(intermediate_state.take().map(|state| *state)); match calculator.root_with_progress().unwrap() { @@ -763,7 +775,7 @@ mod tests { let expected = state_root(state.into_iter()); let tx = factory.provider_rw().unwrap(); - let got = StateRoot::new(tx.tx_ref()).root().unwrap(); + let got = StateRoot::from_tx(tx.tx_ref()).root().unwrap(); assert_eq!(expected, got); } @@ -802,7 +814,7 @@ mod tests { tx.commit().unwrap(); let tx = factory.provider_rw().unwrap(); - let account3_storage_root = StorageRoot::new(tx.tx_ref(), address3).root().unwrap(); + let account3_storage_root = StorageRoot::from_tx(tx.tx_ref(), address3).root().unwrap(); let expected_root = storage_root_prehashed(storage.into_iter()); assert_eq!(expected_root, account3_storage_root); } @@ -869,7 +881,7 @@ mod tests { } hashed_storage_cursor.upsert(key3, StorageEntry { key: hashed_slot, value }).unwrap(); } - let account3_storage_root = StorageRoot::new(tx.tx_ref(), address3).root().unwrap(); + let account3_storage_root = StorageRoot::from_tx(tx.tx_ref(), address3).root().unwrap(); hash_builder.add_leaf( Nibbles::unpack(key3), &encode_account(account3, Some(account3_storage_root)), @@ -918,7 +930,7 @@ mod tests { assert_eq!(hash_builder.root(), computed_expected_root); // Check state root calculation from scratch - let (root, trie_updates) = StateRoot::new(tx.tx_ref()).root_with_updates().unwrap(); + let (root, trie_updates) = StateRoot::from_tx(tx.tx_ref()).root_with_updates().unwrap(); assert_eq!(root, computed_expected_root); // Check account trie @@ -983,7 +995,7 @@ mod tests { B256::from_str("8e263cd4eefb0c3cbbb14e5541a66a755cad25bcfab1e10dd9d706263e811b28") .unwrap(); - let (root, trie_updates) = StateRoot::new(tx.tx_ref()) + let (root, trie_updates) = StateRoot::from_tx(tx.tx_ref()) .with_changed_account_prefixes(prefix_set.freeze()) .root_with_updates() .unwrap(); @@ -1035,7 +1047,7 @@ mod tests { (key6, encode_account(account6, None)), ]); - let (root, trie_updates) = StateRoot::new(tx.tx_ref()) + let (root, trie_updates) = StateRoot::from_tx(tx.tx_ref()) .with_changed_account_prefixes(account_prefix_set.freeze()) .root_with_updates() .unwrap(); @@ -1092,7 +1104,7 @@ mod tests { (key6, encode_account(account6, None)), ]); - let (root, trie_updates) = StateRoot::new(tx.tx_ref()) + let (root, trie_updates) = StateRoot::from_tx(tx.tx_ref()) .with_changed_account_prefixes(account_prefix_set.freeze()) .root_with_updates() .unwrap(); @@ -1131,7 +1143,7 @@ mod tests { let expected = extension_node_trie(&tx); - let (got, updates) = StateRoot::new(tx.tx_ref()).root_with_updates().unwrap(); + let (got, updates) = StateRoot::from_tx(tx.tx_ref()).root_with_updates().unwrap(); assert_eq!(expected, got); // Check account trie @@ -1156,7 +1168,7 @@ mod tests { let expected = extension_node_trie(&tx); - let (got, updates) = StateRoot::new(tx.tx_ref()).root_with_updates().unwrap(); + let (got, updates) = StateRoot::from_tx(tx.tx_ref()).root_with_updates().unwrap(); assert_eq!(expected, got); updates.flush(tx.tx_ref()).unwrap(); @@ -1195,7 +1207,7 @@ mod tests { } } - let (state_root, trie_updates) = StateRoot::new(tx.tx_ref()) + let (state_root, trie_updates) = StateRoot::from_tx(tx.tx_ref()) .with_changed_account_prefixes(changes.freeze()) .root_with_updates() .unwrap(); @@ -1220,7 +1232,7 @@ mod tests { let (expected_root, expected_updates) = extension_node_storage_trie(&tx, hashed_address); let (got, _, updates) = - StorageRoot::new_hashed(tx.tx_ref(), hashed_address).root_with_updates().unwrap(); + StorageRoot::from_tx_hashed(tx.tx_ref(), hashed_address).root_with_updates().unwrap(); assert_eq!(expected_root, got); // Check account trie diff --git a/crates/trie/src/trie_cursor/account_cursor.rs b/crates/trie/src/trie_cursor/account_cursor.rs deleted file mode 100644 index 95033fae6d..0000000000 --- a/crates/trie/src/trie_cursor/account_cursor.rs +++ /dev/null @@ -1,92 +0,0 @@ -use super::TrieCursor; -use crate::updates::TrieKey; -use reth_db::{cursor::DbCursorRO, tables, DatabaseError}; -use reth_primitives::trie::{BranchNodeCompact, StoredNibbles}; - -/// A cursor over the account trie. -#[derive(Debug)] -pub struct AccountTrieCursor(C); - -impl AccountTrieCursor { - /// Create a new account trie cursor. - pub fn new(cursor: C) -> Self { - Self(cursor) - } -} - -impl TrieCursor for AccountTrieCursor -where - C: DbCursorRO, -{ - type Key = StoredNibbles; - - fn seek_exact( - &mut self, - key: Self::Key, - ) -> Result, BranchNodeCompact)>, DatabaseError> { - Ok(self.0.seek_exact(key)?.map(|value| (value.0 .0.to_vec(), value.1))) - } - - fn seek( - &mut self, - key: Self::Key, - ) -> Result, BranchNodeCompact)>, DatabaseError> { - Ok(self.0.seek(key)?.map(|value| (value.0 .0.to_vec(), value.1))) - } - - fn current(&mut self) -> Result, DatabaseError> { - Ok(self.0.current()?.map(|(k, _)| TrieKey::AccountNode(k))) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use reth_db::{ - cursor::{DbCursorRO, DbCursorRW}, - tables, - transaction::DbTxMut, - }; - use reth_primitives::hex_literal::hex; - use reth_provider::test_utils::create_test_provider_factory; - - #[test] - fn test_account_trie_order() { - let factory = create_test_provider_factory(); - let provider = factory.provider_rw().unwrap(); - let mut cursor = provider.tx_ref().cursor_write::().unwrap(); - - let data = vec![ - hex!("0303040e").to_vec(), - hex!("030305").to_vec(), - hex!("03030500").to_vec(), - hex!("0303050a").to_vec(), - ]; - - for key in data.clone() { - cursor - .upsert( - key.into(), - BranchNodeCompact::new( - 0b0000_0010_0000_0001, - 0b0000_0010_0000_0001, - 0, - Vec::default(), - None, - ), - ) - .unwrap(); - } - - let db_data = cursor.walk_range(..).unwrap().collect::, _>>().unwrap(); - assert_eq!(db_data[0].0 .0.to_vec(), data[0]); - assert_eq!(db_data[1].0 .0.to_vec(), data[1]); - assert_eq!(db_data[2].0 .0.to_vec(), data[2]); - assert_eq!(db_data[3].0 .0.to_vec(), data[3]); - - assert_eq!( - cursor.seek(hex!("0303040f").to_vec().into()).unwrap().map(|(k, _)| k.0.to_vec()), - Some(data[1].clone()) - ); - } -} diff --git a/crates/trie/src/trie_cursor/database_cursors.rs b/crates/trie/src/trie_cursor/database_cursors.rs new file mode 100644 index 0000000000..2b8484fe35 --- /dev/null +++ b/crates/trie/src/trie_cursor/database_cursors.rs @@ -0,0 +1,191 @@ +use super::{TrieCursor, TrieCursorFactory}; +use crate::updates::TrieKey; +use reth_db::{ + cursor::{DbCursorRO, DbDupCursorRO}, + tables, + transaction::DbTx, + DatabaseError, +}; +use reth_primitives::{ + trie::{BranchNodeCompact, StoredNibbles, StoredNibblesSubKey}, + B256, +}; + +/// Implementation of the trie cursor factory for a database transaction. +impl<'a, TX: DbTx> TrieCursorFactory for &'a TX { + fn account_trie_cursor( + &self, + ) -> Result + '_>, DatabaseError> { + Ok(Box::new(DatabaseAccountTrieCursor::new(self.cursor_read::()?))) + } + + fn storage_tries_cursor( + &self, + hashed_address: B256, + ) -> Result + '_>, DatabaseError> { + Ok(Box::new(DatabaseStorageTrieCursor::new( + self.cursor_dup_read::()?, + hashed_address, + ))) + } +} + +/// A cursor over the account trie. +#[derive(Debug)] +pub struct DatabaseAccountTrieCursor(C); + +impl DatabaseAccountTrieCursor { + /// Create a new account trie cursor. + pub fn new(cursor: C) -> Self { + Self(cursor) + } +} + +impl TrieCursor for DatabaseAccountTrieCursor +where + C: DbCursorRO, +{ + type Key = StoredNibbles; + + fn seek_exact( + &mut self, + key: Self::Key, + ) -> Result, BranchNodeCompact)>, DatabaseError> { + Ok(self.0.seek_exact(key)?.map(|value| (value.0 .0.to_vec(), value.1))) + } + + fn seek( + &mut self, + key: Self::Key, + ) -> Result, BranchNodeCompact)>, DatabaseError> { + Ok(self.0.seek(key)?.map(|value| (value.0 .0.to_vec(), value.1))) + } + + fn current(&mut self) -> Result, DatabaseError> { + Ok(self.0.current()?.map(|(k, _)| TrieKey::AccountNode(k))) + } +} + +/// A cursor over the storage tries stored in the database. +#[derive(Debug)] +pub struct DatabaseStorageTrieCursor { + /// The underlying cursor. + pub cursor: C, + hashed_address: B256, +} + +impl DatabaseStorageTrieCursor { + /// Create a new storage trie cursor. + pub fn new(cursor: C, hashed_address: B256) -> Self { + Self { cursor, hashed_address } + } +} + +impl TrieCursor for DatabaseStorageTrieCursor +where + C: DbDupCursorRO + DbCursorRO, +{ + type Key = StoredNibblesSubKey; + + fn seek_exact( + &mut self, + key: Self::Key, + ) -> Result, BranchNodeCompact)>, DatabaseError> { + Ok(self + .cursor + .seek_by_key_subkey(self.hashed_address, key.clone())? + .filter(|e| e.nibbles == key) + .map(|value| (value.nibbles.to_vec(), value.node))) + } + + fn seek( + &mut self, + key: Self::Key, + ) -> Result, BranchNodeCompact)>, DatabaseError> { + Ok(self + .cursor + .seek_by_key_subkey(self.hashed_address, key)? + .map(|value| (value.nibbles.to_vec(), value.node))) + } + + fn current(&mut self) -> Result, DatabaseError> { + Ok(self.cursor.current()?.map(|(k, v)| TrieKey::StorageNode(k, v.nibbles))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use reth_db::{ + cursor::{DbCursorRO, DbCursorRW}, + tables, + transaction::DbTxMut, + }; + use reth_primitives::{ + hex_literal::hex, + trie::{BranchNodeCompact, StorageTrieEntry}, + }; + use reth_provider::test_utils::create_test_provider_factory; + + #[test] + fn test_account_trie_order() { + let factory = create_test_provider_factory(); + let provider = factory.provider_rw().unwrap(); + let mut cursor = provider.tx_ref().cursor_write::().unwrap(); + + let data = vec![ + hex!("0303040e").to_vec(), + hex!("030305").to_vec(), + hex!("03030500").to_vec(), + hex!("0303050a").to_vec(), + ]; + + for key in data.clone() { + cursor + .upsert( + key.into(), + BranchNodeCompact::new( + 0b0000_0010_0000_0001, + 0b0000_0010_0000_0001, + 0, + Vec::default(), + None, + ), + ) + .unwrap(); + } + + let db_data = cursor.walk_range(..).unwrap().collect::, _>>().unwrap(); + assert_eq!(db_data[0].0 .0.to_vec(), data[0]); + assert_eq!(db_data[1].0 .0.to_vec(), data[1]); + assert_eq!(db_data[2].0 .0.to_vec(), data[2]); + assert_eq!(db_data[3].0 .0.to_vec(), data[3]); + + assert_eq!( + cursor.seek(hex!("0303040f").to_vec().into()).unwrap().map(|(k, _)| k.0.to_vec()), + Some(data[1].clone()) + ); + } + + // tests that upsert and seek match on the storage trie cursor + #[test] + fn test_storage_cursor_abstraction() { + let factory = create_test_provider_factory(); + let provider = factory.provider_rw().unwrap(); + let mut cursor = provider.tx_ref().cursor_dup_write::().unwrap(); + + let hashed_address = B256::random(); + let key = vec![0x2, 0x3]; + let value = BranchNodeCompact::new(1, 1, 1, vec![B256::random()], None); + + cursor + .upsert( + hashed_address, + StorageTrieEntry { nibbles: key.clone().into(), node: value.clone() }, + ) + .unwrap(); + + let mut cursor = DatabaseStorageTrieCursor::new(cursor, hashed_address); + assert_eq!(cursor.seek(key.into()).unwrap().unwrap().1, value); + } +} diff --git a/crates/trie/src/trie_cursor/mod.rs b/crates/trie/src/trie_cursor/mod.rs index 4b50c59beb..ee2ecec49a 100644 --- a/crates/trie/src/trie_cursor/mod.rs +++ b/crates/trie/src/trie_cursor/mod.rs @@ -1,17 +1,37 @@ use crate::updates::TrieKey; use reth_db::DatabaseError; -use reth_primitives::trie::BranchNodeCompact; - -mod account_cursor; -mod storage_cursor; -mod subnode; - -pub use self::{ - account_cursor::AccountTrieCursor, storage_cursor::StorageTrieCursor, subnode::CursorSubNode, +use reth_primitives::{ + trie::{BranchNodeCompact, StoredNibbles, StoredNibblesSubKey}, + B256, }; +mod database_cursors; +mod subnode; + +/// Noop trie cursor implementations. +pub mod noop; + +pub use self::{ + database_cursors::{DatabaseAccountTrieCursor, DatabaseStorageTrieCursor}, + subnode::CursorSubNode, +}; + +/// Factory for creating trie cursors. +pub trait TrieCursorFactory { + /// Create an account trie cursor. + fn account_trie_cursor( + &self, + ) -> Result + '_>, DatabaseError>; + + /// Create a storage tries cursor. + fn storage_tries_cursor( + &self, + hashed_address: B256, + ) -> Result + '_>, DatabaseError>; +} + /// A cursor for navigating a trie that works with both Tables and DupSort tables. -#[auto_impl::auto_impl(&mut)] +#[auto_impl::auto_impl(&mut, Box)] pub trait TrieCursor { /// The key type of the cursor. type Key: From>; diff --git a/crates/trie/src/trie_cursor/noop.rs b/crates/trie/src/trie_cursor/noop.rs new file mode 100644 index 0000000000..32487c0638 --- /dev/null +++ b/crates/trie/src/trie_cursor/noop.rs @@ -0,0 +1,80 @@ +use reth_db::DatabaseError; +use reth_primitives::trie::{BranchNodeCompact, StoredNibbles, StoredNibblesSubKey}; + +use crate::updates::TrieKey; + +use super::{TrieCursor, TrieCursorFactory}; + +/// Noop trie cursor factory. +#[derive(Default, Debug)] +#[non_exhaustive] +pub struct NoopTrieCursorFactory; + +impl TrieCursorFactory for NoopTrieCursorFactory { + fn account_trie_cursor( + &self, + ) -> Result + '_>, DatabaseError> { + Ok(Box::::default()) + } + + fn storage_tries_cursor( + &self, + _hashed_address: reth_primitives::B256, + ) -> Result + '_>, DatabaseError> { + Ok(Box::::default()) + } +} + +/// Noop account trie cursor. +#[derive(Default, Debug)] +#[non_exhaustive] +pub struct NoopAccountTrieCursor; + +impl TrieCursor for NoopAccountTrieCursor { + type Key = StoredNibbles; + + fn seek( + &mut self, + _key: Self::Key, + ) -> Result, BranchNodeCompact)>, DatabaseError> { + Ok(None) + } + + fn seek_exact( + &mut self, + _key: Self::Key, + ) -> Result, BranchNodeCompact)>, DatabaseError> { + Ok(None) + } + + fn current(&mut self) -> Result, DatabaseError> { + Ok(None) + } +} + +/// Noop account trie cursor. +#[derive(Default, Debug)] +#[non_exhaustive] +pub struct NoopStorageTrieCursor; + +impl TrieCursor for NoopStorageTrieCursor { + type Key = StoredNibblesSubKey; + + fn seek( + &mut self, + _key: Self::Key, + ) -> Result, BranchNodeCompact)>, DatabaseError> { + Ok(None) + } + + fn seek_exact( + &mut self, + _key: Self::Key, + ) -> Result, BranchNodeCompact)>, DatabaseError> { + Ok(None) + } + + fn current(&mut self) -> Result, DatabaseError> { + Ok(None) + } +} diff --git a/crates/trie/src/trie_cursor/storage_cursor.rs b/crates/trie/src/trie_cursor/storage_cursor.rs deleted file mode 100644 index 04b82ecbe4..0000000000 --- a/crates/trie/src/trie_cursor/storage_cursor.rs +++ /dev/null @@ -1,88 +0,0 @@ -use super::TrieCursor; -use crate::updates::TrieKey; -use reth_db::{ - cursor::{DbCursorRO, DbDupCursorRO}, - tables, DatabaseError, -}; -use reth_primitives::{ - trie::{BranchNodeCompact, StoredNibblesSubKey}, - B256, -}; - -/// A cursor over the storage trie. -#[derive(Debug)] -pub struct StorageTrieCursor { - /// The underlying cursor. - pub cursor: C, - hashed_address: B256, -} - -impl StorageTrieCursor { - /// Create a new storage trie cursor. - pub fn new(cursor: C, hashed_address: B256) -> Self { - Self { cursor, hashed_address } - } -} - -impl TrieCursor for StorageTrieCursor -where - C: DbDupCursorRO + DbCursorRO, -{ - type Key = StoredNibblesSubKey; - - fn seek_exact( - &mut self, - key: Self::Key, - ) -> Result, BranchNodeCompact)>, DatabaseError> { - Ok(self - .cursor - .seek_by_key_subkey(self.hashed_address, key.clone())? - .filter(|e| e.nibbles == key) - .map(|value| (value.nibbles.to_vec(), value.node))) - } - - fn seek( - &mut self, - key: Self::Key, - ) -> Result, BranchNodeCompact)>, DatabaseError> { - Ok(self - .cursor - .seek_by_key_subkey(self.hashed_address, key)? - .map(|value| (value.nibbles.to_vec(), value.node))) - } - - fn current(&mut self) -> Result, DatabaseError> { - Ok(self.cursor.current()?.map(|(k, v)| TrieKey::StorageNode(k, v.nibbles))) - } -} - -#[cfg(test)] -mod tests { - - use super::*; - use reth_db::{cursor::DbCursorRW, tables, transaction::DbTxMut}; - use reth_primitives::trie::{BranchNodeCompact, StorageTrieEntry}; - use reth_provider::test_utils::create_test_provider_factory; - - // tests that upsert and seek match on the storagetrie cursor - #[test] - fn test_storage_cursor_abstraction() { - let factory = create_test_provider_factory(); - let provider = factory.provider_rw().unwrap(); - let mut cursor = provider.tx_ref().cursor_dup_write::().unwrap(); - - let hashed_address = B256::random(); - let key = vec![0x2, 0x3]; - let value = BranchNodeCompact::new(1, 1, 1, vec![B256::random()], None); - - cursor - .upsert( - hashed_address, - StorageTrieEntry { nibbles: key.clone().into(), node: value.clone() }, - ) - .unwrap(); - - let mut cursor = StorageTrieCursor::new(cursor, hashed_address); - assert_eq!(cursor.seek(key.into()).unwrap().unwrap().1, value); - } -} diff --git a/crates/trie/src/walker.rs b/crates/trie/src/walker.rs index 67da36d742..9f80459d79 100644 --- a/crates/trie/src/walker.rs +++ b/crates/trie/src/walker.rs @@ -250,7 +250,7 @@ mod tests { use super::*; use crate::{ prefix_set::PrefixSetMut, - trie_cursor::{AccountTrieCursor, StorageTrieCursor}, + trie_cursor::{DatabaseAccountTrieCursor, DatabaseStorageTrieCursor}, }; use reth_db::{cursor::DbCursorRW, tables, transaction::DbTxMut}; use reth_primitives::trie::StorageTrieEntry; @@ -286,7 +286,7 @@ mod tests { for (k, v) in &inputs { account_cursor.upsert(k.clone().into(), v.clone()).unwrap(); } - let account_trie = AccountTrieCursor::new(account_cursor); + let account_trie = DatabaseAccountTrieCursor::new(account_cursor); test_cursor(account_trie, &expected); let hashed_address = B256::random(); @@ -299,7 +299,7 @@ mod tests { ) .unwrap(); } - let storage_trie = StorageTrieCursor::new(storage_cursor, hashed_address); + let storage_trie = DatabaseStorageTrieCursor::new(storage_cursor, hashed_address); test_cursor(storage_trie, &expected); } @@ -357,7 +357,7 @@ mod tests { cursor.upsert(hashed_address, StorageTrieEntry { nibbles: k.into(), node: v }).unwrap(); } - let mut trie = StorageTrieCursor::new(cursor, hashed_address); + let mut trie = DatabaseStorageTrieCursor::new(cursor, hashed_address); // No changes let mut cursor = TrieWalker::new(&mut trie, Default::default());