diff --git a/Cargo.lock b/Cargo.lock index e30f53c7dcf4..48367d0073e9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6448,6 +6448,7 @@ name = "reth-trie" version = "0.1.0-alpha.10" dependencies = [ "alloy-rlp", + "auto_impl", "criterion", "derive_more", "once_cell", diff --git a/crates/primitives/src/stage/checkpoints.rs b/crates/primitives/src/stage/checkpoints.rs index fa5725df55b1..8225eddd691a 100644 --- a/crates/primitives/src/stage/checkpoints.rs +++ b/crates/primitives/src/stage/checkpoints.rs @@ -17,6 +17,7 @@ pub struct MerkleCheckpoint { pub target_block: BlockNumber, /// The last hashed account key processed. pub last_account_key: B256, + // TODO: remove in the next breaking release. /// The last walker key processed. pub last_walker_key: Vec, /// Previously recorded walker stack. @@ -30,11 +31,16 @@ impl MerkleCheckpoint { pub fn new( target_block: BlockNumber, last_account_key: B256, - last_walker_key: Vec, walker_stack: Vec, state: HashBuilderState, ) -> Self { - Self { target_block, last_account_key, last_walker_key, walker_stack, state } + Self { + target_block, + last_account_key, + walker_stack, + state, + last_walker_key: Vec::default(), + } } } diff --git a/crates/stages/src/stages/merkle.rs b/crates/stages/src/stages/merkle.rs index 60d5f466dc5d..71dd050fa36e 100644 --- a/crates/stages/src/stages/merkle.rs +++ b/crates/stages/src/stages/merkle.rs @@ -224,7 +224,6 @@ impl Stage for MerkleStage { let checkpoint = MerkleCheckpoint::new( to_block, state.last_account_key, - state.last_walker_key.hex_data.to_vec(), state.walker_stack.into_iter().map(StoredSubNode::from).collect(), state.hash_builder.into(), ); diff --git a/crates/trie/Cargo.toml b/crates/trie/Cargo.toml index 5f654209a542..643988501bc9 100644 --- a/crates/trie/Cargo.toml +++ b/crates/trie/Cargo.toml @@ -27,6 +27,7 @@ tracing.workspace = true # misc thiserror.workspace = true derive_more = "0.99" +auto_impl = "1" # test-utils triehash = { version = "0.8", optional = true } diff --git a/crates/trie/src/lib.rs b/crates/trie/src/lib.rs index 436c9f9ca10a..31183f5420e8 100644 --- a/crates/trie/src/lib.rs +++ b/crates/trie/src/lib.rs @@ -34,6 +34,9 @@ pub mod walker; mod errors; pub use errors::*; +// The iterators for traversing existing intermediate hashes and updated trie leaves. +pub(crate) mod node_iter; + /// Merkle proof generation. pub mod proof; diff --git a/crates/trie/src/node_iter.rs b/crates/trie/src/node_iter.rs new file mode 100644 index 000000000000..1d433bb30109 --- /dev/null +++ b/crates/trie/src/node_iter.rs @@ -0,0 +1,208 @@ +use crate::{ + hashed_cursor::{HashedAccountCursor, HashedStorageCursor}, + trie_cursor::TrieCursor, + walker::TrieWalker, + StateRootError, StorageRootError, +}; +use reth_primitives::{trie::Nibbles, Account, StorageEntry, B256, U256}; + +#[derive(Debug)] +pub(crate) struct TrieBranchNode { + pub(crate) key: Nibbles, + pub(crate) value: B256, + pub(crate) children_are_in_trie: bool, +} + +impl TrieBranchNode { + pub(crate) fn new(key: Nibbles, value: B256, children_are_in_trie: bool) -> Self { + Self { key, value, children_are_in_trie } + } +} + +#[derive(Debug)] +pub(crate) enum AccountNode { + Branch(TrieBranchNode), + Leaf(B256, Account), +} + +#[derive(Debug)] +pub(crate) enum StorageNode { + Branch(TrieBranchNode), + Leaf(B256, U256), +} + +/// An iterator over existing intermediate branch nodes and updated leaf nodes. +#[derive(Debug)] +pub(crate) struct AccountNodeIter { + /// Underlying walker over intermediate nodes. + pub(crate) walker: TrieWalker, + /// The cursor for the hashed account entries. + pub(crate) hashed_account_cursor: H, + /// The previous account key. If the iteration was previously interrupted, this value can be + /// used to resume iterating from the last returned leaf node. + previous_account_key: Option, + + /// Current hashed account entry. + current_hashed_entry: Option<(B256, Account)>, + /// Flag indicating whether we should check the current walker key. + current_walker_key_checked: bool, +} + +impl AccountNodeIter { + pub(crate) fn new(walker: TrieWalker, hashed_account_cursor: H) -> Self { + Self { + walker, + hashed_account_cursor, + previous_account_key: None, + current_hashed_entry: None, + current_walker_key_checked: false, + } + } + + pub(crate) fn with_last_account_key(mut self, previous_account_key: B256) -> Self { + self.previous_account_key = Some(previous_account_key); + self + } +} + +impl AccountNodeIter +where + C: TrieCursor, + H: HashedAccountCursor, +{ + /// Return the next account trie node to be added to the hash builder. + /// + /// Returns the nodes using this algorithm: + /// 1. Return the current intermediate branch node if it hasn't been updated. + /// 2. Advance the trie walker to the next intermediate branch node and retrieve next + /// unprocessed key. + /// 3. Reposition the hashed account cursor on the next unprocessed key. + /// 4. Return every hashed account entry up to the key of the current intermediate branch node. + /// 5. Repeat. + /// + /// NOTE: The iteration will start from the key of the previous hashed entry if it was supplied. + pub(crate) fn try_next(&mut self) -> Result, StateRootError> { + loop { + if let Some(key) = self.walker.key() { + if !self.current_walker_key_checked && self.previous_account_key.is_none() { + self.current_walker_key_checked = true; + if self.walker.can_skip_current_node { + return Ok(Some(AccountNode::Branch(TrieBranchNode::new( + key, + self.walker.hash().unwrap(), + self.walker.children_are_in_trie(), + )))) + } + } + } + + if let Some((hashed_address, account)) = self.current_hashed_entry.take() { + if self.walker.key().map_or(false, |key| key < Nibbles::unpack(hashed_address)) { + self.current_walker_key_checked = false; + continue + } + + self.current_hashed_entry = self.hashed_account_cursor.next()?; + return Ok(Some(AccountNode::Leaf(hashed_address, account))) + } + + match self.previous_account_key.take() { + Some(account_key) => { + self.hashed_account_cursor.seek(account_key)?; + self.current_hashed_entry = self.hashed_account_cursor.next()?; + } + None => { + let seek_key = match self.walker.next_unprocessed_key() { + Some(key) => key, + None => break, // no more keys + }; + self.current_hashed_entry = self.hashed_account_cursor.seek(seek_key)?; + self.walker.advance()?; + } + } + } + + Ok(None) + } +} + +#[derive(Debug)] +pub(crate) struct StorageNodeIter { + /// Underlying walker over intermediate nodes. + pub(crate) walker: TrieWalker, + /// The cursor for the hashed storage entries. + pub(crate) hashed_storage_cursor: H, + /// The hashed address this storage trie belongs to. + hashed_address: B256, + + /// Current hashed storage entry. + current_hashed_entry: Option, + /// Flag indicating whether we should check the current walker key. + current_walker_key_checked: bool, +} + +impl StorageNodeIter { + pub(crate) fn new( + walker: TrieWalker, + hashed_storage_cursor: H, + hashed_address: B256, + ) -> Self { + Self { + walker, + hashed_storage_cursor, + hashed_address, + current_walker_key_checked: false, + current_hashed_entry: None, + } + } +} + +impl StorageNodeIter +where + C: TrieCursor, + H: HashedStorageCursor, +{ + /// Return the next storage trie node to be added to the hash builder. + /// + /// Returns the nodes using this algorithm: + /// 1. Return the current intermediate branch node if it hasn't been updated. + /// 2. Advance the trie walker to the next intermediate branch node and retrieve next + /// unprocessed key. + /// 3. Reposition the hashed storage cursor on the next unprocessed key. + /// 4. Return every hashed storage entry up to the key of the current intermediate branch node. + /// 5. Repeat. + pub(crate) fn try_next(&mut self) -> Result, StorageRootError> { + loop { + if let Some(key) = self.walker.key() { + if !self.current_walker_key_checked { + self.current_walker_key_checked = true; + if self.walker.can_skip_current_node { + return Ok(Some(StorageNode::Branch(TrieBranchNode::new( + key, + self.walker.hash().unwrap(), + self.walker.children_are_in_trie(), + )))) + } + } + } + + if let Some(StorageEntry { key: hashed_key, value }) = self.current_hashed_entry.take() + { + if self.walker.key().map_or(false, |key| key < Nibbles::unpack(hashed_key)) { + self.current_walker_key_checked = false; + continue + } + + self.current_hashed_entry = self.hashed_storage_cursor.next()?; + return Ok(Some(StorageNode::Leaf(hashed_key, value))) + } + + let Some(seek_key) = self.walker.next_unprocessed_key() else { break }; + self.current_hashed_entry = + self.hashed_storage_cursor.seek(self.hashed_address, seek_key)?; + self.walker.advance()?; + } + + Ok(None) + } +} diff --git a/crates/trie/src/progress.rs b/crates/trie/src/progress.rs index fdd7186a914e..94cd757bc477 100644 --- a/crates/trie/src/progress.rs +++ b/crates/trie/src/progress.rs @@ -1,9 +1,5 @@ use crate::{trie_cursor::CursorSubNode, updates::TrieUpdates}; -use reth_primitives::{ - stage::MerkleCheckpoint, - trie::{hash_builder::HashBuilder, Nibbles}, - B256, -}; +use reth_primitives::{stage::MerkleCheckpoint, trie::hash_builder::HashBuilder, B256}; /// The progress of the state root computation. #[derive(Debug)] @@ -24,8 +20,6 @@ pub struct IntermediateStateRootState { pub walker_stack: Vec, /// The last hashed account key processed. pub last_account_key: B256, - /// The last walker key processed. - pub last_walker_key: Nibbles, } impl From for IntermediateStateRootState { @@ -34,7 +28,6 @@ impl From for IntermediateStateRootState { hash_builder: HashBuilder::from(value.state), walker_stack: value.walker_stack.into_iter().map(CursorSubNode::from).collect(), last_account_key: value.last_account_key, - last_walker_key: Nibbles::from_hex(value.last_walker_key), } } } diff --git a/crates/trie/src/proof.rs b/crates/trie/src/proof.rs index d4ee8fa4c86f..bb7d41e3b58a 100644 --- a/crates/trie/src/proof.rs +++ b/crates/trie/src/proof.rs @@ -1,10 +1,11 @@ use crate::{ account::EthAccount, - hashed_cursor::{HashedAccountCursor, HashedCursorFactory, HashedStorageCursor}, + hashed_cursor::{HashedCursorFactory, HashedStorageCursor}, + node_iter::{AccountNode, AccountNodeIter, StorageNode, StorageNodeIter}, prefix_set::PrefixSetMut, trie_cursor::{AccountTrieCursor, StorageTrieCursor}, walker::TrieWalker, - StorageRootError, + StateRootError, StorageRootError, }; use alloy_rlp::{BufMut, Encodable}; use reth_db::{tables, transaction::DbTx}; @@ -12,7 +13,7 @@ use reth_primitives::{ keccak256, proofs::EMPTY_ROOT, trie::{AccountProof, HashBuilder, Nibbles, StorageProof}, - Address, StorageEntry, B256, + Address, B256, }; /// A struct for generating merkle proofs. @@ -45,65 +46,46 @@ where &self, address: Address, slots: &[B256], - ) -> Result { + ) -> Result { let target_hashed_address = keccak256(address); let target_nibbles = Nibbles::unpack(target_hashed_address); let mut account_proof = AccountProof::new(address); - let mut trie_cursor = - AccountTrieCursor::new(self.tx.cursor_read::()?); - let mut hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?; + let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?; + let trie_cursor = AccountTrieCursor::new(self.tx.cursor_read::()?); // Create the walker. let mut prefix_set = PrefixSetMut::default(); prefix_set.insert(target_nibbles.clone()); - let mut walker = TrieWalker::new(&mut trie_cursor, prefix_set.freeze()); + let walker = TrieWalker::new(trie_cursor, prefix_set.freeze()); // Create a hash builder to rebuild the root node since it is not available in the database. let mut hash_builder = HashBuilder::default().with_proof_retainer(Vec::from([target_nibbles.clone()])); let mut account_rlp = Vec::with_capacity(128); - while let Some(key) = walker.key() { - if walker.can_skip_current_node { - let value = walker.hash().unwrap(); - let is_in_db_trie = walker.children_are_in_trie(); - hash_builder.add_branch(key.clone(), value, is_in_db_trie); - } - - let seek_key = match walker.next_unprocessed_key() { - Some(key) => key, - None => break, // no more keys - }; - - let next_key = walker.advance()?; - let mut next_account_entry = hashed_account_cursor.seek(seek_key)?; - while let Some((hashed_address, account)) = next_account_entry { - let account_nibbles = Nibbles::unpack(hashed_address); - - if let Some(ref key) = next_key { - if key < &account_nibbles { - break - } + let mut account_node_iter = AccountNodeIter::new(walker, hashed_account_cursor); + while let Some(account_node) = account_node_iter.try_next()? { + match account_node { + AccountNode::Branch(node) => { + hash_builder.add_branch(node.key, node.value, node.children_are_in_trie); + } + AccountNode::Leaf(hashed_address, account) => { + let storage_root = if hashed_address == target_hashed_address { + let (storage_root, storage_proofs) = + self.storage_root_with_proofs(hashed_address, slots)?; + account_proof.set_account(account, storage_root, storage_proofs); + storage_root + } else { + self.storage_root(hashed_address)? + }; + + account_rlp.clear(); + let account = EthAccount::from(account).with_storage_root(storage_root); + account.encode(&mut account_rlp as &mut dyn BufMut); + + hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp); } - - let storage_root = if hashed_address == target_hashed_address { - let (storage_root, storage_proofs) = - self.storage_root_with_proofs(hashed_address, slots)?; - account_proof.set_account(account, storage_root, storage_proofs); - storage_root - } else { - self.storage_root(hashed_address)? - }; - - account_rlp.clear(); - let account = EthAccount::from(account).with_storage_root(storage_root); - account.encode(&mut &mut account_rlp as &mut dyn BufMut); - - hash_builder.add_leaf(account_nibbles, &account_rlp); - - // Move the next account entry - next_account_entry = hashed_account_cursor.next()?; } } @@ -129,11 +111,6 @@ where ) -> Result<(B256, Vec), StorageRootError> { let mut hashed_storage_cursor = self.hashed_cursor_factory.hashed_storage_cursor()?; - let mut trie_cursor = StorageTrieCursor::new( - self.tx.cursor_dup_read::()?, - hashed_address, - ); - let mut proofs = slots.iter().copied().map(StorageProof::new).collect::>(); // short circuit on empty storage @@ -143,52 +120,41 @@ where let target_nibbles = proofs.iter().map(|p| p.nibbles.clone()).collect::>(); let prefix_set = PrefixSetMut::from(target_nibbles.clone()).freeze(); - let mut walker = TrieWalker::new(&mut trie_cursor, prefix_set); + let trie_cursor = StorageTrieCursor::new( + self.tx.cursor_dup_read::()?, + hashed_address, + ); + let walker = TrieWalker::new(trie_cursor, prefix_set); let mut hash_builder = HashBuilder::default().with_proof_retainer(target_nibbles); - while let Some(key) = walker.key() { - if walker.can_skip_current_node { - hash_builder.add_branch(key, walker.hash().unwrap(), walker.children_are_in_trie()); - } - - let seek_key = match walker.next_unprocessed_key() { - Some(key) => key, - None => break, // no more keys - }; - - let next_key = walker.advance()?; - let mut storage = hashed_storage_cursor.seek(hashed_address, seek_key)?; - while let Some(StorageEntry { key: hashed_key, value }) = storage { - let hashed_key_nibbles = Nibbles::unpack(hashed_key); - if let Some(ref key) = next_key { - if key < &hashed_key_nibbles { - break - } + let mut storage_node_iter = + StorageNodeIter::new(walker, hashed_storage_cursor, hashed_address); + while let Some(node) = storage_node_iter.try_next()? { + match node { + StorageNode::Branch(node) => { + hash_builder.add_branch(node.key, node.value, node.children_are_in_trie); } - - if let Some(proof) = - proofs.iter_mut().find(|proof| proof.nibbles == hashed_key_nibbles) - { - proof.set_value(value); + StorageNode::Leaf(hashed_slot, value) => { + let nibbles = Nibbles::unpack(hashed_slot); + if let Some(proof) = proofs.iter_mut().find(|proof| proof.nibbles == nibbles) { + proof.set_value(value); + } + hash_builder.add_leaf(nibbles, alloy_rlp::encode_fixed_size(&value).as_ref()); } - - hash_builder - .add_leaf(hashed_key_nibbles, alloy_rlp::encode_fixed_size(&value).as_ref()); - storage = hashed_storage_cursor.next()?; } } let root = hash_builder.root(); - let proof_nodes = hash_builder.take_proofs(); + let all_proof_nodes = hash_builder.take_proofs(); for proof in proofs.iter_mut() { - proof.set_proof( - proof_nodes - .iter() - .filter(|(path, _)| proof.nibbles.starts_with(path)) - .map(|(_, node)| node.clone()) - .collect(), - ); + // Iterate over all proof nodes and find the matching ones. + // The filtered results are guaranteed to be in order. + let matching_proof_nodes = all_proof_nodes + .iter() + .filter(|(path, _)| proof.nibbles.starts_with(path)) + .map(|(_, node)| node.clone()); + proof.set_proof(matching_proof_nodes.collect()); } Ok((root, proofs)) diff --git a/crates/trie/src/trie.rs b/crates/trie/src/trie.rs index 335dea750bb8..89bf27fe9557 100644 --- a/crates/trie/src/trie.rs +++ b/crates/trie/src/trie.rs @@ -1,6 +1,7 @@ use crate::{ account::EthAccount, - hashed_cursor::{HashedAccountCursor, HashedCursorFactory, HashedStorageCursor}, + hashed_cursor::{HashedCursorFactory, HashedStorageCursor}, + node_iter::{AccountNode, AccountNodeIter, StorageNode, StorageNodeIter}, prefix_set::{PrefixSet, PrefixSetLoader, PrefixSetMut}, progress::{IntermediateStateRootState, StateRootProgress}, trie_cursor::{AccountTrieCursor, StorageTrieCursor}, @@ -8,13 +9,13 @@ use crate::{ walker::TrieWalker, StateRootError, StorageRootError, }; -use alloy_rlp::Encodable; +use alloy_rlp::{BufMut, Encodable}; use reth_db::{tables, transaction::DbTx}; use reth_primitives::{ keccak256, proofs::EMPTY_ROOT, trie::{HashBuilder, Nibbles}, - Address, BlockNumber, StorageEntry, B256, + Address, BlockNumber, B256, }; use std::{ collections::{HashMap, HashSet}, @@ -224,136 +225,104 @@ where tracing::debug!(target: "loader", "calculating state root"); let mut trie_updates = TrieUpdates::default(); - let mut hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?; - let mut trie_cursor = - AccountTrieCursor::new(self.tx.cursor_read::()?); - - let (mut walker, mut hash_builder, mut last_account_key, mut last_walker_key) = - match self.previous_state { - Some(state) => ( - TrieWalker::from_stack( - &mut trie_cursor, - state.walker_stack, - self.changed_account_prefixes, - ), + let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?; + let trie_cursor = AccountTrieCursor::new(self.tx.cursor_read::()?); + + let (mut hash_builder, mut account_node_iter) = match self.previous_state { + Some(state) => { + let walker = TrieWalker::from_stack( + trie_cursor, + state.walker_stack, + self.changed_account_prefixes, + ); + ( state.hash_builder, - Some(state.last_account_key), - Some(state.last_walker_key), - ), - None => ( - TrieWalker::new(&mut trie_cursor, self.changed_account_prefixes), - HashBuilder::default(), - None, - None, - ), - }; + AccountNodeIter::new(walker, hashed_account_cursor) + .with_last_account_key(state.last_account_key), + ) + } + None => { + let walker = TrieWalker::new(trie_cursor, self.changed_account_prefixes); + (HashBuilder::default(), AccountNodeIter::new(walker, hashed_account_cursor)) + } + }; - walker.set_updates(retain_updates); + account_node_iter.walker.set_updates(retain_updates); hash_builder.set_updates(retain_updates); let mut account_rlp = Vec::with_capacity(128); let mut hashed_entries_walked = 0; - - while let Some(key) = last_walker_key.take().or_else(|| walker.key()) { - // Take the last account key to make sure we take it into consideration only once. - let (next_key, mut next_account_entry) = match last_account_key.take() { - // Seek the last processed entry and take the next after. - Some(account_key) => { - hashed_account_cursor.seek(account_key)?; - (walker.key(), hashed_account_cursor.next()?) + while let Some(node) = account_node_iter.try_next()? { + match node { + AccountNode::Branch(node) => { + hash_builder.add_branch(node.key, node.value, node.children_are_in_trie); } - None => { - if walker.can_skip_current_node { - let value = walker.hash().unwrap(); - let is_in_db_trie = walker.children_are_in_trie(); - hash_builder.add_branch(key.clone(), value, is_in_db_trie); - } - - let seek_key = match walker.next_unprocessed_key() { - Some(key) => key, - None => break, // no more keys + AccountNode::Leaf(hashed_address, account) => { + hashed_entries_walked += 1; + + // We assume we can always calculate a storage root without + // OOMing. This opens us up to a potential DOS vector if + // a contract had too many storage entries and they were + // all buffered w/o us returning and committing our intermediate + // 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 = if retain_updates { + let (root, storage_slots_walked, updates) = + storage_root_calculator.root_with_updates()?; + hashed_entries_walked += storage_slots_walked; + trie_updates.extend(updates.into_iter()); + root + } else { + storage_root_calculator.root()? }; - (walker.advance()?, hashed_account_cursor.seek(seek_key)?) - } - }; + let account = EthAccount::from(account).with_storage_root(storage_root); - while let Some((hashed_address, account)) = next_account_entry { - hashed_entries_walked += 1; - let account_nibbles = Nibbles::unpack(hashed_address); + account_rlp.clear(); + account.encode(&mut account_rlp as &mut dyn BufMut); - if let Some(ref key) = next_key { - if key < &account_nibbles { - tracing::trace!(target: "loader", "breaking, already detected"); - break - } - } + hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp); - // We assume we can always calculate a storage root without - // OOMing. This opens us up to a potential DOS vector if - // a contract had too many storage entries and they were - // all buffered w/o us returning and committing our intermediate - // 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(), - ); + // Decide if we need to return intermediate progress. + let total_updates_len = trie_updates.len() + + account_node_iter.walker.updates_len() + + hash_builder.updates_len(); + if retain_updates && total_updates_len as u64 >= self.threshold { + let (walker_stack, walker_updates) = account_node_iter.walker.split(); + let (hash_builder, hash_builder_updates) = hash_builder.split(); - let storage_root = if retain_updates { - let (root, storage_slots_walked, updates) = - storage_root_calculator.root_with_updates()?; - hashed_entries_walked += storage_slots_walked; - trie_updates.extend(updates.into_iter()); - root - } else { - storage_root_calculator.root()? - }; - - let account = EthAccount::from(account).with_storage_root(storage_root); - - account_rlp.clear(); - account.encode(&mut &mut account_rlp); - - hash_builder.add_leaf(account_nibbles, &account_rlp); - - // Decide if we need to return intermediate progress. - let total_updates_len = - trie_updates.len() + walker.updates_len() + hash_builder.updates_len(); - if retain_updates && total_updates_len as u64 >= self.threshold { - let (walker_stack, walker_updates) = walker.split(); - let (hash_builder, hash_builder_updates) = hash_builder.split(); - - let state = IntermediateStateRootState { - hash_builder, - walker_stack, - last_walker_key: key, - last_account_key: hashed_address, - }; + let state = IntermediateStateRootState { + hash_builder, + walker_stack, + last_account_key: hashed_address, + }; - trie_updates.extend(walker_updates.into_iter()); - trie_updates.extend_with_account_updates(hash_builder_updates); + trie_updates.extend(walker_updates.into_iter()); + trie_updates.extend_with_account_updates(hash_builder_updates); - return Ok(StateRootProgress::Progress( - Box::new(state), - hashed_entries_walked, - trie_updates, - )) + return Ok(StateRootProgress::Progress( + Box::new(state), + hashed_entries_walked, + trie_updates, + )) + } } - - // Move the next account entry - next_account_entry = hashed_account_cursor.next()?; } } let root = hash_builder.root(); - let (_, walker_updates) = walker.split(); + let (_, walker_updates) = account_node_iter.walker.split(); let (_, hash_builder_updates) = hash_builder.split(); trie_updates.extend(walker_updates.into_iter()); @@ -464,14 +433,8 @@ where retain_updates: bool, ) -> Result<(B256, usize, TrieUpdates), StorageRootError> { tracing::debug!(target: "trie::storage_root", hashed_address = ?self.hashed_address, "calculating storage root"); - let mut hashed_storage_cursor = self.hashed_cursor_factory.hashed_storage_cursor()?; - let mut trie_cursor = StorageTrieCursor::new( - self.tx.cursor_dup_read::()?, - self.hashed_address, - ); - // short circuit on empty storage if hashed_storage_cursor.is_storage_empty(self.hashed_address)? { return Ok(( @@ -481,43 +444,37 @@ where )) } - let mut walker = TrieWalker::new(&mut trie_cursor, self.changed_prefixes.clone()) + let trie_cursor = StorageTrieCursor::new( + self.tx.cursor_dup_read::()?, + self.hashed_address, + ); + let walker = TrieWalker::new(trie_cursor, self.changed_prefixes.clone()) .with_updates(retain_updates); let mut hash_builder = HashBuilder::default().with_updates(retain_updates); let mut storage_slots_walked = 0; - while let Some(key) = walker.key() { - if walker.can_skip_current_node { - hash_builder.add_branch(key, walker.hash().unwrap(), walker.children_are_in_trie()); - } - - let seek_key = match walker.next_unprocessed_key() { - Some(key) => key, - None => break, // no more keys - }; - - let next_key = walker.advance()?; - let mut storage = hashed_storage_cursor.seek(self.hashed_address, seek_key)?; - while let Some(StorageEntry { key: hashed_key, value }) = storage { - storage_slots_walked += 1; - - let storage_key_nibbles = Nibbles::unpack(hashed_key); - if let Some(ref key) = next_key { - if key < &storage_key_nibbles { - break - } + let mut storage_node_iter = + StorageNodeIter::new(walker, hashed_storage_cursor, self.hashed_address); + while let Some(node) = storage_node_iter.try_next()? { + match node { + StorageNode::Branch(node) => { + hash_builder.add_branch(node.key, node.value, node.children_are_in_trie); + } + StorageNode::Leaf(hashed_slot, value) => { + storage_slots_walked += 1; + hash_builder.add_leaf( + Nibbles::unpack(hashed_slot), + alloy_rlp::encode_fixed_size(&value).as_ref(), + ); } - hash_builder - .add_leaf(storage_key_nibbles, alloy_rlp::encode_fixed_size(&value).as_ref()); - storage = hashed_storage_cursor.next()?; } } let root = hash_builder.root(); let (_, hash_builder_updates) = hash_builder.split(); - let (_, walker_updates) = walker.split(); + let (_, walker_updates) = storage_node_iter.walker.split(); let mut trie_updates = TrieUpdates::default(); trie_updates.extend(walker_updates.into_iter()); @@ -529,7 +486,6 @@ where } #[cfg(test)] -#[allow(clippy::mutable_key_type)] mod tests { use super::*; use crate::test_utils::{ @@ -548,7 +504,7 @@ mod tests { keccak256, proofs::triehash::KeccakHasher, trie::{BranchNodeCompact, TrieMask}, - Account, Address, B256, MAINNET, U256, + Account, Address, StorageEntry, B256, MAINNET, U256, }; use reth_provider::{DatabaseProviderRW, ProviderFactory}; use std::{collections::BTreeMap, ops::Mul, str::FromStr}; @@ -621,7 +577,6 @@ mod tests { } #[test] - // TODO: Try to find the edge case by creating some more very complex trie. fn branch_node_child_changes() { incremental_vs_full_root( &[ diff --git a/crates/trie/src/trie_cursor/account_cursor.rs b/crates/trie/src/trie_cursor/account_cursor.rs index 899152ad33c8..815396ab0ac1 100644 --- a/crates/trie/src/trie_cursor/account_cursor.rs +++ b/crates/trie/src/trie_cursor/account_cursor.rs @@ -14,20 +14,22 @@ impl AccountTrieCursor { } } -impl TrieCursor for AccountTrieCursor +impl TrieCursor for AccountTrieCursor where C: DbCursorRO, { + type Key = StoredNibbles; + fn seek_exact( &mut self, - key: StoredNibbles, + key: Self::Key, ) -> Result, BranchNodeCompact)>, DatabaseError> { Ok(self.0.seek_exact(key)?.map(|value| (value.0.inner.to_vec(), value.1))) } fn seek( &mut self, - key: StoredNibbles, + key: Self::Key, ) -> Result, BranchNodeCompact)>, DatabaseError> { Ok(self.0.seek(key)?.map(|value| (value.0.inner.to_vec(), value.1))) } diff --git a/crates/trie/src/trie_cursor/mod.rs b/crates/trie/src/trie_cursor/mod.rs index 85de68875dd5..4b50c59bebc2 100644 --- a/crates/trie/src/trie_cursor/mod.rs +++ b/crates/trie/src/trie_cursor/mod.rs @@ -1,5 +1,5 @@ use crate::updates::TrieKey; -use reth_db::{table::Key, DatabaseError}; +use reth_db::DatabaseError; use reth_primitives::trie::BranchNodeCompact; mod account_cursor; @@ -11,13 +11,22 @@ pub use self::{ }; /// A cursor for navigating a trie that works with both Tables and DupSort tables. -pub trait TrieCursor { +#[auto_impl::auto_impl(&mut)] +pub trait TrieCursor { + /// The key type of the cursor. + type Key: From>; + /// Move the cursor to the key and return if it is an exact match. - fn seek_exact(&mut self, key: K) - -> Result, BranchNodeCompact)>, DatabaseError>; + fn seek_exact( + &mut self, + key: Self::Key, + ) -> Result, BranchNodeCompact)>, DatabaseError>; /// Move the cursor to the key and return a value matching of greater than the key. - fn seek(&mut self, key: K) -> Result, BranchNodeCompact)>, DatabaseError>; + fn seek( + &mut self, + key: Self::Key, + ) -> Result, BranchNodeCompact)>, DatabaseError>; /// Get the current entry. fn current(&mut self) -> Result, DatabaseError>; diff --git a/crates/trie/src/trie_cursor/storage_cursor.rs b/crates/trie/src/trie_cursor/storage_cursor.rs index 5cc18c13a810..032ff97f64fe 100644 --- a/crates/trie/src/trie_cursor/storage_cursor.rs +++ b/crates/trie/src/trie_cursor/storage_cursor.rs @@ -24,13 +24,15 @@ impl StorageTrieCursor { } } -impl TrieCursor for StorageTrieCursor +impl TrieCursor for StorageTrieCursor where C: DbDupCursorRO + DbCursorRO, { + type Key = StoredNibblesSubKey; + fn seek_exact( &mut self, - key: StoredNibblesSubKey, + key: Self::Key, ) -> Result, BranchNodeCompact)>, DatabaseError> { Ok(self .cursor @@ -41,7 +43,7 @@ where fn seek( &mut self, - key: StoredNibblesSubKey, + key: Self::Key, ) -> Result, BranchNodeCompact)>, DatabaseError> { Ok(self .cursor diff --git a/crates/trie/src/updates.rs b/crates/trie/src/updates.rs index 02255f1828df..11cfb82204a3 100644 --- a/crates/trie/src/updates.rs +++ b/crates/trie/src/updates.rs @@ -77,7 +77,6 @@ impl TrieUpdates { } /// Extend the updates with account trie updates. - #[allow(clippy::mutable_key_type)] pub fn extend_with_account_updates(&mut self, updates: HashMap) { self.extend(updates.into_iter().map(|(nibbles, node)| { (TrieKey::AccountNode(nibbles.hex_data.to_vec().into()), TrieOp::Update(node)) @@ -85,7 +84,6 @@ impl TrieUpdates { } /// Extend the updates with storage trie updates. - #[allow(clippy::mutable_key_type)] pub fn extend_with_storage_updates( &mut self, hashed_address: B256, diff --git a/crates/trie/src/walker.rs b/crates/trie/src/walker.rs index e9832d61bd40..402977bfb1f0 100644 --- a/crates/trie/src/walker.rs +++ b/crates/trie/src/walker.rs @@ -3,19 +3,19 @@ use crate::{ trie_cursor::{CursorSubNode, TrieCursor}, updates::TrieUpdates, }; -use reth_db::{table::Key, DatabaseError}; +use reth_db::DatabaseError; use reth_primitives::{ trie::{BranchNodeCompact, Nibbles}, B256, }; -use std::marker::PhantomData; /// `TrieWalker` is a structure that enables traversal of a Merkle trie. -/// It allows moving through the trie in a depth-first manner, skipping certain branches if the . +/// It allows moving through the trie in a depth-first manner, skipping certain branches +/// if they have not changed. #[derive(Debug)] -pub struct TrieWalker<'a, K, C> { +pub struct TrieWalker { /// A mutable reference to a trie cursor instance used for navigating the trie. - pub cursor: &'a mut C, + pub cursor: C, /// A vector containing the trie nodes that have been visited. pub stack: Vec, /// A flag indicating whether the current node can be skipped when traversing the trie. This @@ -26,12 +26,11 @@ pub struct TrieWalker<'a, K, C> { pub changes: PrefixSet, /// The trie updates to be applied to the trie. trie_updates: Option, - __phantom: PhantomData, } -impl<'a, K: Key + From>, C: TrieCursor> TrieWalker<'a, K, C> { +impl TrieWalker { /// Constructs a new TrieWalker, setting up the initial state of the stack and cursor. - pub fn new(cursor: &'a mut C, changes: PrefixSet) -> Self { + pub fn new(cursor: C, changes: PrefixSet) -> Self { // Initialize the walker with a single empty stack element. let mut this = Self { cursor, @@ -39,7 +38,6 @@ impl<'a, K: Key + From>, C: TrieCursor> TrieWalker<'a, K, C> { stack: vec![CursorSubNode::default()], can_skip_current_node: false, trie_updates: None, - __phantom: PhantomData, }; // Set up the root node of the trie in the stack, if it exists. @@ -53,15 +51,9 @@ impl<'a, K: Key + From>, C: TrieCursor> TrieWalker<'a, K, C> { } /// Constructs a new TrieWalker from existing stack and a cursor. - pub fn from_stack(cursor: &'a mut C, stack: Vec, changes: PrefixSet) -> Self { - let mut this = Self { - cursor, - changes, - stack, - can_skip_current_node: false, - trie_updates: None, - __phantom: PhantomData, - }; + pub fn from_stack(cursor: C, stack: Vec, changes: PrefixSet) -> Self { + let mut this = + Self { cursor, changes, stack, can_skip_current_node: false, trie_updates: None }; this.update_skip_node(); this } @@ -255,7 +247,6 @@ impl<'a, K: Key + From>, C: TrieCursor> TrieWalker<'a, K, C> { #[cfg(test)] mod tests { - use super::*; use crate::{ prefix_set::PrefixSetMut, @@ -316,10 +307,9 @@ mod tests { test_cursor(storage_trie, &expected); } - fn test_cursor(mut trie: T, expected: &[Vec]) + fn test_cursor(mut trie: T, expected: &[Vec]) where - K: Key + From>, - T: TrieCursor, + T: TrieCursor, { let mut walker = TrieWalker::new(&mut trie, Default::default()); assert!(walker.key().unwrap().is_empty());