diff --git a/Cargo.lock b/Cargo.lock index 9641d8e26..a3ccf46b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9454,6 +9454,7 @@ version = "1.1.2" dependencies = [ "alloy-primitives", "alloy-rlp", + "arbitrary", "assert_matches", "criterion", "itertools 0.13.0", diff --git a/crates/trie/sparse/Cargo.toml b/crates/trie/sparse/Cargo.toml index 330197596..4c64bf716 100644 --- a/crates/trie/sparse/Cargo.toml +++ b/crates/trie/sparse/Cargo.toml @@ -32,6 +32,7 @@ reth-testing-utils.workspace = true reth-trie = { workspace = true, features = ["test-utils"] } reth-trie-common = { workspace = true, features = ["test-utils", "arbitrary"] } +arbitrary.workspace = true assert_matches.workspace = true criterion.workspace = true itertools.workspace = true diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index 0b0db1401..b76da7937 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -6,17 +6,23 @@ use alloy_primitives::{ Bytes, B256, }; use alloy_rlp::Decodable; -use reth_trie::{Nibbles, TrieNode}; +use reth_trie::{ + updates::{StorageTrieUpdates, TrieUpdates}, + Nibbles, TrieNode, +}; /// Sparse state trie representing lazy-loaded Ethereum state trie. #[derive(Default, Debug)] pub struct SparseStateTrie { + retain_updates: bool, /// Sparse account trie. - pub(crate) state: SparseTrie, + state: SparseTrie, /// Sparse storage tries. - pub(crate) storages: HashMap, + storages: HashMap, /// Collection of revealed account and storage keys. - pub(crate) revealed: HashMap>, + revealed: HashMap>, + /// Collection of addresses that had their storage tries wiped. + wiped_storages: HashSet, } impl SparseStateTrie { @@ -25,6 +31,12 @@ impl SparseStateTrie { Self { state, ..Default::default() } } + /// Set the retention of branch node updates and deletions. + pub const fn with_updates(mut self, retain_updates: bool) -> Self { + self.retain_updates = retain_updates; + self + } + /// Returns `true` if account was already revealed. pub fn is_account_revealed(&self, account: &B256) -> bool { self.revealed.contains_key(account) @@ -42,7 +54,7 @@ impl SparseStateTrie { account: B256, proof: impl IntoIterator, ) -> SparseStateTrieResult<()> { - if self.revealed.contains_key(&account) { + if self.is_account_revealed(&account) { return Ok(()); } @@ -51,7 +63,7 @@ impl SparseStateTrie { let Some(root_node) = self.validate_proof(&mut proof)? else { return Ok(()) }; // Reveal root node if it wasn't already. - let trie = self.state.reveal_root(root_node)?; + let trie = self.state.reveal_root(root_node, self.retain_updates)?; // Reveal the remaining proof nodes. for (path, bytes) in proof { @@ -73,7 +85,7 @@ impl SparseStateTrie { slot: B256, proof: impl IntoIterator, ) -> SparseStateTrieResult<()> { - if self.revealed.get(&account).is_some_and(|v| v.contains(&slot)) { + if self.is_storage_slot_revealed(&account, &slot) { return Ok(()); } @@ -82,7 +94,11 @@ impl SparseStateTrie { let Some(root_node) = self.validate_proof(&mut proof)? else { return Ok(()) }; // Reveal root node if it wasn't already. - let trie = self.storages.entry(account).or_default().reveal_root(root_node)?; + let trie = self + .storages + .entry(account) + .or_default() + .reveal_root(root_node, self.retain_updates)?; // Reveal the remaining proof nodes. for (path, bytes) in proof { @@ -118,30 +134,98 @@ impl SparseStateTrie { Ok(Some(root_node)) } - /// Update the leaf node. - pub fn update_leaf(&mut self, path: Nibbles, value: Vec) -> SparseStateTrieResult<()> { + /// Update the account leaf node. + pub fn update_account_leaf( + &mut self, + path: Nibbles, + value: Vec, + ) -> SparseStateTrieResult<()> { self.state.update_leaf(path, value)?; Ok(()) } + /// Remove the account leaf node. + pub fn remove_account_leaf(&mut self, path: &Nibbles) -> SparseStateTrieResult<()> { + self.state.remove_leaf(path)?; + Ok(()) + } + /// Returns sparse trie root if the trie has been revealed. pub fn root(&mut self) -> Option { self.state.root() } + /// Calculates the hashes of the nodes below the provided level. + pub fn calculate_below_level(&mut self, level: usize) { + self.state.calculate_below_level(level); + } + + /// Update the leaf node of a storage trie at the provided address. + pub fn update_storage_leaf( + &mut self, + address: B256, + slot: Nibbles, + value: Vec, + ) -> SparseStateTrieResult<()> { + self.storages.entry(address).or_default().update_leaf(slot, value)?; + Ok(()) + } + + /// Wipe the storage trie at the provided address. + pub fn wipe_storage(&mut self, address: B256) -> SparseStateTrieResult<()> { + let Some(trie) = self.storages.get_mut(&address) else { return Ok(()) }; + self.wiped_storages.insert(address); + trie.wipe().map_err(Into::into) + } + /// Returns storage sparse trie root if the trie has been revealed. pub fn storage_root(&mut self, account: B256) -> Option { self.storages.get_mut(&account).and_then(|trie| trie.root()) } + + /// Returns [`TrieUpdates`] by taking the updates from the revealed sparse tries. + /// + /// Returns `None` if the accounts trie is not revealed. + pub fn take_trie_updates(&mut self) -> Option { + self.state.as_revealed_mut().map(|state| { + let updates = state.take_updates(); + TrieUpdates { + account_nodes: HashMap::from_iter(updates.updated_nodes), + removed_nodes: HashSet::from_iter(updates.removed_nodes), + storage_tries: self + .storages + .iter_mut() + .map(|(address, trie)| { + let trie = trie.as_revealed_mut().unwrap(); + let updates = trie.take_updates(); + let updates = StorageTrieUpdates { + is_deleted: self.wiped_storages.contains(address), + storage_nodes: HashMap::from_iter(updates.updated_nodes), + removed_nodes: HashSet::from_iter(updates.removed_nodes), + }; + (*address, updates) + }) + .filter(|(_, updates)| !updates.is_empty()) + .collect(), + } + }) + } } #[cfg(test)] mod tests { use super::*; - use alloy_primitives::Bytes; + use alloy_primitives::{b256, Bytes, U256}; use alloy_rlp::EMPTY_STRING_CODE; + use arbitrary::Arbitrary; use assert_matches::assert_matches; - use reth_trie::HashBuilder; + use itertools::Itertools; + use rand::{rngs::StdRng, Rng, SeedableRng}; + use reth_primitives_traits::Account; + use reth_trie::{ + updates::StorageTrieUpdates, BranchNodeCompact, HashBuilder, TrieAccount, TrieMask, + EMPTY_ROOT_HASH, + }; use reth_trie_common::proof::ProofRetainer; #[test] @@ -199,4 +283,159 @@ mod tests { HashMap::from_iter([(Default::default(), SparseTrie::revealed_empty())]) ); } + + #[test] + fn take_trie_updates() { + reth_tracing::init_test_tracing(); + + // let mut rng = generators::rng(); + let mut rng = StdRng::seed_from_u64(1); + + let mut bytes = [0u8; 1024]; + rng.fill(bytes.as_mut_slice()); + + let slot_1 = b256!("1000000000000000000000000000000000000000000000000000000000000000"); + let slot_path_1 = Nibbles::unpack(slot_1); + let value_1 = U256::from(rng.gen::()); + let slot_2 = b256!("1100000000000000000000000000000000000000000000000000000000000000"); + let slot_path_2 = Nibbles::unpack(slot_2); + let value_2 = U256::from(rng.gen::()); + let slot_3 = b256!("2000000000000000000000000000000000000000000000000000000000000000"); + let slot_path_3 = Nibbles::unpack(slot_3); + let value_3 = U256::from(rng.gen::()); + + let mut storage_hash_builder = + HashBuilder::default().with_proof_retainer(ProofRetainer::from_iter([ + slot_path_1.clone(), + slot_path_2.clone(), + ])); + storage_hash_builder.add_leaf(slot_path_1.clone(), &alloy_rlp::encode_fixed_size(&value_1)); + storage_hash_builder.add_leaf(slot_path_2.clone(), &alloy_rlp::encode_fixed_size(&value_2)); + + let storage_root = storage_hash_builder.root(); + let proof_nodes = storage_hash_builder.take_proof_nodes(); + let storage_proof_1 = proof_nodes + .iter() + .filter(|(path, _)| path.is_empty() || slot_path_1.common_prefix_length(path) > 0) + .map(|(path, proof)| (path.clone(), proof.clone())) + .sorted_by_key(|(path, _)| path.clone()) + .collect::>(); + let storage_proof_2 = proof_nodes + .iter() + .filter(|(path, _)| path.is_empty() || slot_path_2.common_prefix_length(path) > 0) + .map(|(path, proof)| (path.clone(), proof.clone())) + .sorted_by_key(|(path, _)| path.clone()) + .collect::>(); + + let address_1 = b256!("1000000000000000000000000000000000000000000000000000000000000000"); + let address_path_1 = Nibbles::unpack(address_1); + let account_1 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap(); + let mut trie_account_1 = TrieAccount::from((account_1, storage_root)); + let address_2 = b256!("1100000000000000000000000000000000000000000000000000000000000000"); + let address_path_2 = Nibbles::unpack(address_2); + let account_2 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap(); + let mut trie_account_2 = TrieAccount::from((account_2, EMPTY_ROOT_HASH)); + + let mut hash_builder = + HashBuilder::default().with_proof_retainer(ProofRetainer::from_iter([ + address_path_1.clone(), + address_path_2.clone(), + ])); + hash_builder.add_leaf(address_path_1.clone(), &alloy_rlp::encode(trie_account_1)); + hash_builder.add_leaf(address_path_2.clone(), &alloy_rlp::encode(trie_account_2)); + + let root = hash_builder.root(); + let proof_nodes = hash_builder.take_proof_nodes(); + let proof_1 = proof_nodes + .iter() + .filter(|(path, _)| path.is_empty() || address_path_1.common_prefix_length(path) > 0) + .map(|(path, proof)| (path.clone(), proof.clone())) + .sorted_by_key(|(path, _)| path.clone()) + .collect::>(); + let proof_2 = proof_nodes + .iter() + .filter(|(path, _)| path.is_empty() || address_path_2.common_prefix_length(path) > 0) + .map(|(path, proof)| (path.clone(), proof.clone())) + .sorted_by_key(|(path, _)| path.clone()) + .collect::>(); + + let mut sparse = SparseStateTrie::default().with_updates(true); + sparse.reveal_account(address_1, proof_1).unwrap(); + sparse.reveal_account(address_2, proof_2).unwrap(); + sparse.reveal_storage_slot(address_1, slot_1, storage_proof_1.clone()).unwrap(); + sparse.reveal_storage_slot(address_1, slot_2, storage_proof_2.clone()).unwrap(); + sparse.reveal_storage_slot(address_2, slot_1, storage_proof_1).unwrap(); + sparse.reveal_storage_slot(address_2, slot_2, storage_proof_2).unwrap(); + + assert_eq!(sparse.root(), Some(root)); + + let address_3 = b256!("2000000000000000000000000000000000000000000000000000000000000000"); + let address_path_3 = Nibbles::unpack(address_3); + let account_3 = Account { nonce: account_1.nonce + 1, ..account_1 }; + let trie_account_3 = TrieAccount::from((account_3, EMPTY_ROOT_HASH)); + + sparse.update_account_leaf(address_path_3, alloy_rlp::encode(trie_account_3)).unwrap(); + + sparse.update_storage_leaf(address_1, slot_path_3, alloy_rlp::encode(value_3)).unwrap(); + trie_account_1.storage_root = sparse.storage_root(address_1).unwrap(); + sparse.update_account_leaf(address_path_1, alloy_rlp::encode(trie_account_1)).unwrap(); + + sparse.wipe_storage(address_2).unwrap(); + trie_account_2.storage_root = sparse.storage_root(address_2).unwrap(); + sparse.update_account_leaf(address_path_2, alloy_rlp::encode(trie_account_2)).unwrap(); + + sparse.root(); + + let sparse_updates = sparse.take_trie_updates().unwrap(); + // TODO(alexey): assert against real state root calculation updates + pretty_assertions::assert_eq!( + sparse_updates, + TrieUpdates { + account_nodes: HashMap::from_iter([ + ( + Nibbles::default(), + BranchNodeCompact { + state_mask: TrieMask::new(0b110), + tree_mask: TrieMask::new(0b000), + hash_mask: TrieMask::new(0b010), + hashes: vec![b256!( + "4c4ffbda3569fcf2c24ea2000b4cec86ef8b92cbf9ff415db43184c0f75a212e" + )], + root_hash: Some(b256!( + "60944bd29458529c3065d19f63c6e3d5269596fd3b04ca2e7b318912dc89ca4c" + )) + }, + ), + ]), + storage_tries: HashMap::from_iter([ + ( + b256!("1000000000000000000000000000000000000000000000000000000000000000"), + StorageTrieUpdates { + is_deleted: false, + storage_nodes: HashMap::from_iter([( + Nibbles::default(), + BranchNodeCompact { + state_mask: TrieMask::new(0b110), + tree_mask: TrieMask::new(0b000), + hash_mask: TrieMask::new(0b010), + hashes: vec![b256!("5bc8b4fdf51839c1e18b8d6a4bd3e2e52c9f641860f0e4d197b68c2679b0e436")], + root_hash: Some(b256!("c44abf1a9e1a92736ac479b20328e8d7998aa8838b6ef52620324c9ce85e3201")) + } + )]), + removed_nodes: HashSet::default() + } + ), + ( + b256!("1100000000000000000000000000000000000000000000000000000000000000"), + StorageTrieUpdates { + is_deleted: true, + storage_nodes: HashMap::default(), + removed_nodes: HashSet::default() + } + ) + ]), + removed_nodes: HashSet::default() + } + ); + } } diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index dff290271..bab166d78 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -53,9 +53,13 @@ impl SparseTrie { /// # Returns /// /// Mutable reference to [`RevealedSparseTrie`]. - pub fn reveal_root(&mut self, root: TrieNode) -> SparseTrieResult<&mut RevealedSparseTrie> { + pub fn reveal_root( + &mut self, + root: TrieNode, + retain_updates: bool, + ) -> SparseTrieResult<&mut RevealedSparseTrie> { if self.is_blind() { - *self = Self::Revealed(Box::new(RevealedSparseTrie::from_root(root)?)) + *self = Self::Revealed(Box::new(RevealedSparseTrie::from_root(root, retain_updates)?)) } Ok(self.as_revealed_mut().unwrap()) } @@ -67,10 +71,29 @@ impl SparseTrie { Ok(()) } + /// Remove the leaf node. + pub fn remove_leaf(&mut self, path: &Nibbles) -> SparseTrieResult<()> { + let revealed = self.as_revealed_mut().ok_or(SparseTrieError::Blind)?; + revealed.remove_leaf(path)?; + Ok(()) + } + + /// Wipe the trie, removing all values and nodes, and replacing the root with an empty node. + pub fn wipe(&mut self) -> SparseTrieResult<()> { + let revealed = self.as_revealed_mut().ok_or(SparseTrieError::Blind)?; + revealed.wipe(); + Ok(()) + } + /// Calculates and returns the trie root if the trie has been revealed. pub fn root(&mut self) -> Option { Some(self.as_revealed_mut()?.root()) } + + /// Calculates the hashes of the nodes below the provided level. + pub fn calculate_below_level(&mut self, level: usize) { + self.as_revealed_mut().unwrap().update_rlp_node_level(level); + } } /// The representation of revealed sparse trie. @@ -120,19 +143,20 @@ impl Default for RevealedSparseTrie { impl RevealedSparseTrie { /// Create new revealed sparse trie from the given root node. - pub fn from_root(node: TrieNode) -> SparseTrieResult { + pub fn from_root(node: TrieNode, retain_updates: bool) -> SparseTrieResult { let mut this = Self { nodes: HashMap::default(), values: HashMap::default(), prefix_set: PrefixSetMut::default(), rlp_buf: Vec::new(), updates: None, - }; + } + .with_updates(retain_updates); this.reveal_node(Nibbles::default(), node)?; Ok(this) } - /// Makes the sparse trie to store updated branch nodes. + /// Set the retention of branch node updates and deletions. pub fn with_updates(mut self, retain_updates: bool) -> Self { if retain_updates { self.updates = Some(SparseTrieUpdates::default()); @@ -580,6 +604,12 @@ impl RevealedSparseTrie { Ok(nodes) } + /// Wipe the trie, removing all values and nodes, and replacing the root with an empty node. + pub fn wipe(&mut self) { + *self = Self::default(); + self.prefix_set = PrefixSetMut::all(); + } + /// Return the root of the sparse trie. /// Updates all remaining dirty nodes before calculating the root. pub fn root(&mut self) -> B256 { @@ -773,8 +803,7 @@ impl RevealedSparseTrie { } // Set the hash mask. If a child node has a hash value AND is a - // branch node, set the hash mask - // and save the hash. + // branch node, set the hash mask and save the hash. let hash = child.as_hash().filter(|_| node_type.is_branch()); hash_mask_values.push(hash.is_some()); if let Some(hash) = hash { @@ -998,8 +1027,8 @@ impl RlpNodeBuffers { /// The aggregation of sparse trie updates. #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct SparseTrieUpdates { - updated_nodes: HashMap, - removed_nodes: HashSet, + pub(crate) updated_nodes: HashMap, + pub(crate) removed_nodes: HashSet, } #[cfg(test)] @@ -1560,7 +1589,7 @@ mod tests { TrieMask::new(0b11), )); - let mut sparse = RevealedSparseTrie::from_root(branch.clone()).unwrap(); + let mut sparse = RevealedSparseTrie::from_root(branch.clone(), false).unwrap(); // Reveal a branch node and one of its children // @@ -1722,6 +1751,7 @@ mod tests { .take_proof_nodes(); let mut sparse = RevealedSparseTrie::from_root( TrieNode::decode(&mut &proof_nodes.nodes_sorted()[0].1[..]).unwrap(), + false, ) .unwrap(); @@ -1796,6 +1826,7 @@ mod tests { .take_proof_nodes(); let mut sparse = RevealedSparseTrie::from_root( TrieNode::decode(&mut &proof_nodes.nodes_sorted()[0].1[..]).unwrap(), + false, ) .unwrap(); @@ -1866,6 +1897,7 @@ mod tests { .take_proof_nodes(); let mut sparse = RevealedSparseTrie::from_root( TrieNode::decode(&mut &proof_nodes.nodes_sorted()[0].1[..]).unwrap(), + false, ) .unwrap(); @@ -1993,4 +2025,44 @@ mod tests { assert_eq!(sparse_root, hash_builder.root()); assert_eq!(sparse_updates.updated_nodes, hash_builder.updated_branch_nodes.take().unwrap()); } + + #[test] + fn sparse_trie_wipe() { + let mut sparse = RevealedSparseTrie::default().with_updates(true); + + let value = alloy_rlp::encode_fixed_size(&U256::ZERO).to_vec(); + + // Extension (Key = 5) – Level 0 + // └── Branch (Mask = 1011) – Level 1 + // ├── 0 -> Extension (Key = 23) – Level 2 + // │ └── Branch (Mask = 0101) – Level 3 + // │ ├── 1 -> Leaf (Key = 1, Path = 50231) – Level 4 + // │ └── 3 -> Leaf (Key = 3, Path = 50233) – Level 4 + // ├── 2 -> Leaf (Key = 013, Path = 52013) – Level 2 + // └── 3 -> Branch (Mask = 0101) – Level 2 + // ├── 1 -> Leaf (Key = 3102, Path = 53102) – Level 3 + // └── 3 -> Branch (Mask = 1010) – Level 3 + // ├── 0 -> Leaf (Key = 3302, Path = 53302) – Level 4 + // └── 2 -> Leaf (Key = 3320, Path = 53320) – Level 4 + sparse + .update_leaf(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]), value.clone()) + .unwrap(); + sparse + .update_leaf(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]), value.clone()) + .unwrap(); + sparse + .update_leaf(Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3]), value.clone()) + .unwrap(); + sparse + .update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2]), value.clone()) + .unwrap(); + sparse + .update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2]), value.clone()) + .unwrap(); + sparse.update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0]), value).unwrap(); + + sparse.wipe(); + + assert_eq!(sparse.root(), EMPTY_ROOT_HASH); + } } diff --git a/crates/trie/trie/src/updates.rs b/crates/trie/trie/src/updates.rs index 6d1bcab63..e7bc49064 100644 --- a/crates/trie/trie/src/updates.rs +++ b/crates/trie/trie/src/updates.rs @@ -6,11 +6,14 @@ use std::collections::{HashMap, HashSet}; #[derive(PartialEq, Eq, Clone, Default, Debug)] #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub struct TrieUpdates { + /// Collection of updated intermediate account nodes indexed by full path. #[cfg_attr(feature = "serde", serde(with = "serde_nibbles_map"))] - pub(crate) account_nodes: HashMap, + pub account_nodes: HashMap, + /// Collection of removed intermediate account nodes indexed by full path. #[cfg_attr(feature = "serde", serde(with = "serde_nibbles_set"))] - pub(crate) removed_nodes: HashSet, - pub(crate) storage_tries: HashMap, + pub removed_nodes: HashSet, + /// Collection of updated storage tries indexed by the hashed address. + pub storage_tries: HashMap, } impl TrieUpdates { @@ -113,13 +116,13 @@ impl TrieUpdates { #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub struct StorageTrieUpdates { /// Flag indicating whether the trie was deleted. - pub(crate) is_deleted: bool, + pub is_deleted: bool, /// Collection of updated storage trie nodes. #[cfg_attr(feature = "serde", serde(with = "serde_nibbles_map"))] - pub(crate) storage_nodes: HashMap, + pub storage_nodes: HashMap, /// Collection of removed storage trie nodes. #[cfg_attr(feature = "serde", serde(with = "serde_nibbles_set"))] - pub(crate) removed_nodes: HashSet, + pub removed_nodes: HashSet, } #[cfg(feature = "test-utils")]