diff --git a/crates/trie/trie/src/constants.rs b/crates/trie/common/src/constants.rs similarity index 94% rename from crates/trie/trie/src/constants.rs rename to crates/trie/common/src/constants.rs index 7354290d9596..471b8bd9dcc6 100644 --- a/crates/trie/trie/src/constants.rs +++ b/crates/trie/common/src/constants.rs @@ -5,9 +5,9 @@ pub const TRIE_ACCOUNT_RLP_MAX_SIZE: usize = 110; #[cfg(test)] mod tests { use super::*; + use crate::TrieAccount; use alloy_primitives::{B256, U256}; use alloy_rlp::Encodable; - use reth_trie_common::TrieAccount; #[test] fn account_rlp_max_size() { diff --git a/crates/trie/common/src/lib.rs b/crates/trie/common/src/lib.rs index 04b817aab8f2..6647de67811c 100644 --- a/crates/trie/common/src/lib.rs +++ b/crates/trie/common/src/lib.rs @@ -11,6 +11,10 @@ /// The implementation of hash builder. pub mod hash_builder; +/// Constants related to the trie computation. +mod constants; +pub use constants::*; + mod account; pub use account::TrieAccount; diff --git a/crates/trie/sparse/Cargo.toml b/crates/trie/sparse/Cargo.toml index dce232fcd578..efd68020ccd7 100644 --- a/crates/trie/sparse/Cargo.toml +++ b/crates/trie/sparse/Cargo.toml @@ -14,8 +14,9 @@ workspace = true [dependencies] # reth -reth-tracing.workspace = true +reth-primitives-traits.workspace = true reth-trie-common.workspace = true +reth-tracing.workspace = true # alloy alloy-primitives.workspace = true diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index 551a47ce2bbf..aad74ac0550b 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -1,19 +1,21 @@ -use crate::{RevealedSparseTrie, SparseStateTrieError, SparseStateTrieResult, SparseTrie}; +use crate::{ + RevealedSparseTrie, SparseStateTrieError, SparseStateTrieResult, SparseTrie, SparseTrieError, +}; use alloy_primitives::{ map::{HashMap, HashSet}, Bytes, B256, }; -use alloy_rlp::Decodable; +use alloy_rlp::{Decodable, Encodable}; +use reth_primitives_traits::Account; use reth_trie_common::{ updates::{StorageTrieUpdates, TrieUpdates}, - MultiProof, Nibbles, TrieNode, + MultiProof, Nibbles, TrieAccount, TrieNode, EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE, }; use std::iter::Peekable; /// Sparse state trie representing lazy-loaded Ethereum state trie. -#[derive(Default, Debug)] +#[derive(Debug)] pub struct SparseStateTrie { - retain_updates: bool, /// Sparse account trie. state: SparseTrie, /// Sparse storage tries. @@ -22,6 +24,23 @@ pub struct SparseStateTrie { 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. + account_rlp_buf: Vec, +} + +impl Default for SparseStateTrie { + fn default() -> Self { + Self { + 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), + } + } } impl SparseStateTrie { @@ -186,6 +205,37 @@ impl SparseStateTrie { Ok(Some(root_node)) } + /// Update or remove trie account based on new account info. This method will either recompute + /// the storage root based on update storage trie or look it up from existing leaf value. + /// + /// If the new account info and storage trie are empty, the account leaf will be removed. + pub fn update_account(&mut self, address: B256, account: Account) -> SparseStateTrieResult<()> { + let nibbles = Nibbles::unpack(address); + let storage_root = if let Some(storage_trie) = self.storages.get_mut(&address) { + storage_trie.root().ok_or(SparseTrieError::Blind)? + } else if self.revealed.contains_key(&address) { + let state = self.state.as_revealed_mut().ok_or(SparseTrieError::Blind)?; + // The account was revealed, either... + if let Some(value) = state.get_leaf_value(&nibbles) { + // ..it exists and we should take it's current storage root or... + TrieAccount::decode(&mut &value[..])?.storage_root + } else { + // ...the account is newly created and the storage trie is empty. + EMPTY_ROOT_HASH + } + } else { + return Err(SparseTrieError::Blind.into()) + }; + + if account.is_empty() && storage_root == EMPTY_ROOT_HASH { + self.remove_account_leaf(&nibbles) + } else { + self.account_rlp_buf.clear(); + TrieAccount::from((account, storage_root)).encode(&mut self.account_rlp_buf); + self.update_account_leaf(nibbles, self.account_rlp_buf.clone()) + } + } + /// Update the account leaf node. pub fn update_account_leaf( &mut self, diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 21f1cf410aa6..2ecc3984445d 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -166,6 +166,11 @@ impl RevealedSparseTrie { self.updates.as_ref().map_or(Cow::Owned(SparseTrieUpdates::default()), Cow::Borrowed) } + /// Returns a reference to the leaf value if present. + pub fn get_leaf_value(&self, path: &Nibbles) -> Option<&Vec> { + self.values.get(path) + } + /// Takes and returns the retained sparse node updates pub fn take_updates(&mut self) -> SparseTrieUpdates { self.updates.take().unwrap_or_default() diff --git a/crates/trie/trie/src/lib.rs b/crates/trie/trie/src/lib.rs index 335711b8d88e..1e7eeb9b52b8 100644 --- a/crates/trie/trie/src/lib.rs +++ b/crates/trie/trie/src/lib.rs @@ -13,10 +13,6 @@ )] #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))] -/// Constants related to the trie computation. -mod constants; -pub use constants::*; - /// The implementation of forward-only in-memory cursor. pub mod forward_cursor;