diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index aad74ac0550b..1cd13273e649 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -22,8 +22,6 @@ pub struct SparseStateTrie { storages: HashMap, /// Collection of revealed account and storage keys. revealed: HashMap>, - /// Collection of addresses that had their storage tries wiped. - wiped_storages: HashSet, /// Flag indicating whether trie updates should be retained. retain_updates: bool, /// Reusable buffer for RLP encoding of trie accounts. @@ -36,7 +34,6 @@ impl Default for SparseStateTrie { state: Default::default(), storages: Default::default(), revealed: Default::default(), - wiped_storages: Default::default(), retain_updates: false, account_rlp_buf: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE), } @@ -275,9 +272,10 @@ impl SparseStateTrie { /// 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) + if let Some(trie) = self.storages.get_mut(&address) { + trie.wipe()?; + } + Ok(()) } /// Calculates the hashes of the nodes below the provided level. @@ -302,8 +300,8 @@ impl SparseStateTrie { 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), + account_nodes: updates.updated_nodes, + removed_nodes: updates.removed_nodes, storage_tries: self .storages .iter_mut() @@ -311,9 +309,9 @@ impl SparseStateTrie { 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), + is_deleted: updates.wiped, + storage_nodes: updates.updated_nodes, + removed_nodes: updates.removed_nodes, }; (*address, updates) }) diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 2ecc3984445d..97446680df44 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -111,6 +111,7 @@ pub struct RevealedSparseTrie { prefix_set: PrefixSetMut, /// Reusable buffer for RLP encoding of nodes. rlp_buf: Vec, + /// Retained trie updates. updates: Option, } @@ -607,8 +608,10 @@ impl RevealedSparseTrie { /// Wipe the trie, removing all values and nodes, and replacing the root with an empty node. pub fn wipe(&mut self) { + let updates_retained = self.updates.is_some(); *self = Self::default(); self.prefix_set = PrefixSetMut::all(); + self.updates = updates_retained.then(SparseTrieUpdates::wiped); } /// Return the root of the sparse trie. @@ -1030,12 +1033,18 @@ impl RlpNodeBuffers { pub struct SparseTrieUpdates { pub(crate) updated_nodes: HashMap, pub(crate) removed_nodes: HashSet, + pub(crate) wiped: bool, +} + +impl SparseTrieUpdates { + /// Create new wiped sparse trie updates. + pub fn wiped() -> Self { + Self { wiped: true, ..Default::default() } + } } #[cfg(test)] mod tests { - use std::collections::BTreeMap; - use super::*; use alloy_primitives::{map::HashSet, U256}; use alloy_rlp::Encodable; @@ -1057,6 +1066,7 @@ mod tests { proof::{ProofNodes, ProofRetainer}, HashBuilder, }; + use std::collections::BTreeMap; /// Pad nibbles to the length of a B256 hash with zeros on the left. fn pad_nibbles_left(nibbles: Nibbles) -> Nibbles {