Files
reth/crates/trie/src/trie.rs
2024-03-24 22:52:33 +00:00

1382 lines
53 KiB
Rust
Raw Blame History

This file contains invisible Unicode characters
This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
use crate::{
hashed_cursor::{HashedCursorFactory, HashedStorageCursor},
node_iter::{AccountNode, AccountNodeIter, StorageNode, StorageNodeIter},
prefix_set::{PrefixSet, PrefixSetLoader, TriePrefixSets},
progress::{IntermediateStateRootState, StateRootProgress},
stats::TrieTracker,
trie_cursor::TrieCursorFactory,
updates::{TrieKey, TrieOp, TrieUpdates},
walker::TrieWalker,
};
use alloy_rlp::{BufMut, Encodable};
use reth_db::transaction::DbTx;
use reth_interfaces::trie::{StateRootError, StorageRootError};
use reth_primitives::{
constants::EMPTY_ROOT_HASH,
keccak256,
trie::{HashBuilder, Nibbles, TrieAccount},
Address, BlockNumber, B256,
};
use std::ops::RangeInclusive;
use tracing::{debug, trace};
#[cfg(feature = "metrics")]
use crate::metrics::{StateRootMetrics, TrieRootMetrics, TrieType};
/// StateRoot is used to compute the root node of a state trie.
#[derive(Debug)]
pub struct StateRoot<T, H> {
/// The factory for trie cursors.
pub trie_cursor_factory: T,
/// The factory for hashed cursors.
pub hashed_cursor_factory: H,
/// A set of prefix sets that have changes.
pub prefix_sets: TriePrefixSets,
/// Previous intermediate state.
previous_state: Option<IntermediateStateRootState>,
/// The number of updates after which the intermediate progress should be returned.
threshold: u64,
#[cfg(feature = "metrics")]
/// State root metrics.
metrics: StateRootMetrics,
}
impl<T, H> StateRoot<T, H> {
/// Set the prefix sets.
pub fn with_prefix_sets(mut self, prefix_sets: TriePrefixSets) -> Self {
self.prefix_sets = prefix_sets;
self
}
/// Set the threshold.
pub fn with_threshold(mut self, threshold: u64) -> Self {
self.threshold = threshold;
self
}
/// Set the threshold to maximum value so that intermediate progress is not returned.
pub fn with_no_threshold(mut self) -> Self {
self.threshold = u64::MAX;
self
}
/// Set the previously recorded intermediate state.
pub fn with_intermediate_state(mut self, state: Option<IntermediateStateRootState>) -> Self {
self.previous_state = state;
self
}
/// Set the hashed cursor factory.
pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> StateRoot<T, HF> {
StateRoot {
trie_cursor_factory: self.trie_cursor_factory,
hashed_cursor_factory,
prefix_sets: self.prefix_sets,
threshold: self.threshold,
previous_state: self.previous_state,
#[cfg(feature = "metrics")]
metrics: self.metrics,
}
}
/// Set the trie cursor factory.
pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> StateRoot<TF, H> {
StateRoot {
trie_cursor_factory,
hashed_cursor_factory: self.hashed_cursor_factory,
prefix_sets: self.prefix_sets,
threshold: self.threshold,
previous_state: self.previous_state,
#[cfg(feature = "metrics")]
metrics: self.metrics,
}
}
}
impl<'a, TX: DbTx> StateRoot<&'a TX, &'a TX> {
/// Create a new [StateRoot] instance.
pub fn from_tx(tx: &'a TX) -> Self {
Self {
trie_cursor_factory: tx,
hashed_cursor_factory: tx,
prefix_sets: TriePrefixSets::default(),
previous_state: None,
threshold: 100_000,
#[cfg(feature = "metrics")]
metrics: StateRootMetrics::default(),
}
}
/// Given a block number range, identifies all the accounts and storage keys that
/// have changed.
///
/// # Returns
///
/// An instance of state root calculator with account and storage prefixes loaded.
pub fn incremental_root_calculator(
tx: &'a TX,
range: RangeInclusive<BlockNumber>,
) -> Result<Self, StateRootError> {
let loaded_prefix_sets = PrefixSetLoader::new(tx).load(range)?;
Ok(Self::from_tx(tx).with_prefix_sets(loaded_prefix_sets))
}
/// Computes the state root of the trie with the changed account and storage prefixes and
/// existing trie nodes.
///
/// # Returns
///
/// The updated state root.
pub fn incremental_root(
tx: &'a TX,
range: RangeInclusive<BlockNumber>,
) -> Result<B256, StateRootError> {
debug!(target: "trie::loader", ?range, "incremental state root");
Self::incremental_root_calculator(tx, range)?.root()
}
/// Computes the state root of the trie with the changed account and storage prefixes and
/// existing trie nodes collecting updates in the process.
///
/// Ignores the threshold.
///
/// # Returns
///
/// The updated state root and the trie updates.
pub fn incremental_root_with_updates(
tx: &'a TX,
range: RangeInclusive<BlockNumber>,
) -> Result<(B256, TrieUpdates), StateRootError> {
debug!(target: "trie::loader", ?range, "incremental state root");
Self::incremental_root_calculator(tx, range)?.root_with_updates()
}
/// Computes the state root of the trie with the changed account and storage prefixes and
/// existing trie nodes collecting updates in the process.
///
/// # Returns
///
/// The intermediate progress of state root computation.
pub fn incremental_root_with_progress(
tx: &'a TX,
range: RangeInclusive<BlockNumber>,
) -> Result<StateRootProgress, StateRootError> {
debug!(target: "trie::loader", ?range, "incremental state root with progress");
Self::incremental_root_calculator(tx, range)?.root_with_progress()
}
}
impl<T, H> StateRoot<T, H>
where
T: TrieCursorFactory + Clone,
H: HashedCursorFactory + Clone,
{
/// Walks the intermediate nodes of existing state trie (if any) and hashed entries. Feeds the
/// nodes into the hash builder. Collects the updates in the process.
///
/// Ignores the threshold.
///
/// # Returns
///
/// The intermediate progress of state root computation and the trie updates.
pub fn root_with_updates(self) -> Result<(B256, TrieUpdates), StateRootError> {
match self.with_no_threshold().calculate(true)? {
StateRootProgress::Complete(root, _, updates) => Ok((root, updates)),
StateRootProgress::Progress(..) => unreachable!(), // unreachable threshold
}
}
/// Walks the intermediate nodes of existing state trie (if any) and hashed entries. Feeds the
/// nodes into the hash builder.
///
/// # Returns
///
/// The state root hash.
pub fn root(self) -> Result<B256, StateRootError> {
match self.calculate(false)? {
StateRootProgress::Complete(root, _, _) => Ok(root),
StateRootProgress::Progress(..) => unreachable!(), // update retenion is disabled
}
}
/// Walks the intermediate nodes of existing state trie (if any) and hashed entries. Feeds the
/// nodes into the hash builder. Collects the updates in the process.
///
/// # Returns
///
/// The intermediate progress of state root computation.
pub fn root_with_progress(self) -> Result<StateRootProgress, StateRootError> {
self.calculate(true)
}
fn calculate(self, retain_updates: bool) -> Result<StateRootProgress, StateRootError> {
trace!(target: "trie::state_root", "calculating state root");
let mut tracker = TrieTracker::default();
let mut trie_updates = TrieUpdates::default();
let trie_cursor = self.trie_cursor_factory.account_trie_cursor()?;
let (mut hash_builder, mut account_node_iter) = match self.previous_state {
Some(state) => {
let hash_builder = state.hash_builder.with_updates(retain_updates);
let walker = TrieWalker::from_stack(
trie_cursor,
state.walker_stack,
self.prefix_sets.account_prefix_set,
)
.with_updates(retain_updates);
let node_iter =
AccountNodeIter::from_factory(walker, self.hashed_cursor_factory.clone())?
.with_last_account_key(state.last_account_key);
(hash_builder, node_iter)
}
None => {
let hash_builder = HashBuilder::default().with_updates(retain_updates);
let walker = TrieWalker::new(trie_cursor, self.prefix_sets.account_prefix_set)
.with_updates(retain_updates);
let node_iter =
AccountNodeIter::from_factory(walker, self.hashed_cursor_factory.clone())?;
(hash_builder, node_iter)
}
};
let mut account_rlp = Vec::with_capacity(128);
let mut hashed_entries_walked = 0;
while let Some(node) = account_node_iter.try_next()? {
match node {
AccountNode::Branch(node) => {
tracker.inc_branch();
hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
}
AccountNode::Leaf(hashed_address, account) => {
tracker.inc_leaf();
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.trie_cursor_factory.clone(),
self.hashed_cursor_factory.clone(),
hashed_address,
#[cfg(feature = "metrics")]
self.metrics.storage_trie.clone(),
)
.with_prefix_set(
self.prefix_sets
.storage_prefix_sets
.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);
root
} else {
storage_root_calculator.root()?
};
account_rlp.clear();
let account = TrieAccount::from((account, storage_root));
account.encode(&mut account_rlp as &mut dyn BufMut);
hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp);
// 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 state = IntermediateStateRootState {
hash_builder,
walker_stack,
last_account_key: hashed_address,
};
trie_updates.extend(walker_updates);
trie_updates.extend_with_account_updates(hash_builder_updates);
return Ok(StateRootProgress::Progress(
Box::new(state),
hashed_entries_walked,
trie_updates,
))
}
}
}
}
let root = hash_builder.root();
trie_updates.finalize_state_updates(
account_node_iter.walker,
hash_builder,
self.prefix_sets.destroyed_accounts,
);
let stats = tracker.finish();
#[cfg(feature = "metrics")]
self.metrics.state_trie.record(stats);
trace!(
target: "trie::state_root",
%root,
duration = ?stats.duration(),
branches_added = stats.branches_added(),
leaves_added = stats.leaves_added(),
"calculated state root"
);
Ok(StateRootProgress::Complete(root, hashed_entries_walked, trie_updates))
}
}
/// StorageRoot is used to compute the root node of an account storage trie.
#[derive(Debug)]
pub struct StorageRoot<T, H> {
/// A reference to the database transaction.
pub trie_cursor_factory: T,
/// The factory for hashed cursors.
pub hashed_cursor_factory: H,
/// The hashed address of an account.
pub hashed_address: B256,
/// The set of storage slot prefixes that have changed.
pub prefix_set: PrefixSet,
/// Storage root metrics.
#[cfg(feature = "metrics")]
metrics: TrieRootMetrics,
}
impl<T, H> StorageRoot<T, H> {
/// Creates a new storage root calculator given a raw address.
pub fn new(
trie_cursor_factory: T,
hashed_cursor_factory: H,
address: Address,
#[cfg(feature = "metrics")] metrics: TrieRootMetrics,
) -> Self {
Self::new_hashed(
trie_cursor_factory,
hashed_cursor_factory,
keccak256(address),
#[cfg(feature = "metrics")]
metrics,
)
}
/// Creates a new storage root calculator given a hashed address.
pub fn new_hashed(
trie_cursor_factory: T,
hashed_cursor_factory: H,
hashed_address: B256,
#[cfg(feature = "metrics")] metrics: TrieRootMetrics,
) -> Self {
Self {
trie_cursor_factory,
hashed_cursor_factory,
hashed_address,
prefix_set: PrefixSet::default(),
#[cfg(feature = "metrics")]
metrics,
}
}
/// Set the changed prefixes.
pub fn with_prefix_set(mut self, prefix_set: PrefixSet) -> Self {
self.prefix_set = prefix_set;
self
}
/// Set the hashed cursor factory.
pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> StorageRoot<T, HF> {
StorageRoot {
trie_cursor_factory: self.trie_cursor_factory,
hashed_cursor_factory,
hashed_address: self.hashed_address,
prefix_set: self.prefix_set,
#[cfg(feature = "metrics")]
metrics: self.metrics,
}
}
/// Set the trie cursor factory.
pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> StorageRoot<TF, H> {
StorageRoot {
trie_cursor_factory,
hashed_cursor_factory: self.hashed_cursor_factory,
hashed_address: self.hashed_address,
prefix_set: self.prefix_set,
#[cfg(feature = "metrics")]
metrics: self.metrics,
}
}
}
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,
#[cfg(feature = "metrics")]
TrieRootMetrics::new(TrieType::Storage),
)
}
/// 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,
#[cfg(feature = "metrics")]
TrieRootMetrics::new(TrieType::Storage),
)
}
}
impl<T, H> StorageRoot<T, H>
where
T: TrieCursorFactory,
H: HashedCursorFactory,
{
/// Walks the hashed storage table entries for a given address and calculates the storage root.
///
/// # Returns
///
/// The storage root and storage trie updates for a given address.
pub fn root_with_updates(self) -> Result<(B256, usize, TrieUpdates), StorageRootError> {
self.calculate(true)
}
/// Walks the hashed storage table entries for a given address and calculates the storage root.
///
/// # Returns
///
/// The storage root.
pub fn root(self) -> Result<B256, StorageRootError> {
let (root, _, _) = self.calculate(false)?;
Ok(root)
}
/// Walks the hashed storage table entries for a given address and calculates the storage root.
///
/// # Returns
///
/// The storage root, number of walked entries and trie updates
/// for a given address if requested.
pub fn calculate(
self,
retain_updates: bool,
) -> Result<(B256, usize, TrieUpdates), StorageRootError> {
trace!(target: "trie::storage_root", hashed_address = ?self.hashed_address, "calculating storage root");
let mut hashed_storage_cursor = self.hashed_cursor_factory.hashed_storage_cursor()?;
// short circuit on empty storage
if hashed_storage_cursor.is_storage_empty(self.hashed_address)? {
return Ok((
EMPTY_ROOT_HASH,
0,
TrieUpdates::from([(TrieKey::StorageTrie(self.hashed_address), TrieOp::Delete)]),
))
}
let mut tracker = TrieTracker::default();
let trie_cursor = self.trie_cursor_factory.storage_tries_cursor(self.hashed_address)?;
let walker = TrieWalker::new(trie_cursor, self.prefix_set).with_updates(retain_updates);
let mut hash_builder = HashBuilder::default().with_updates(retain_updates);
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) => {
tracker.inc_branch();
hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
}
StorageNode::Leaf(hashed_slot, value) => {
tracker.inc_leaf();
hash_builder.add_leaf(
Nibbles::unpack(hashed_slot),
alloy_rlp::encode_fixed_size(&value).as_ref(),
);
}
}
}
let root = hash_builder.root();
let mut trie_updates = TrieUpdates::default();
trie_updates.finalize_storage_updates(
self.hashed_address,
storage_node_iter.walker,
hash_builder,
);
let stats = tracker.finish();
#[cfg(feature = "metrics")]
self.metrics.record(stats);
trace!(
target: "trie::storage_root",
%root,
hashed_address = %self.hashed_address,
duration = ?stats.duration(),
branches_added = stats.branches_added(),
leaves_added = stats.leaves_added(),
"calculated storage root"
);
let storage_slots_walked = stats.leaves_added() as usize;
Ok((root, storage_slots_walked, trie_updates))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
prefix_set::PrefixSetMut,
test_utils::{state_root, state_root_prehashed, storage_root, storage_root_prehashed},
};
use proptest::{prelude::ProptestConfig, proptest};
use reth_db::{
cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO},
tables,
test_utils::TempDatabase,
transaction::DbTxMut,
DatabaseEnv,
};
use reth_primitives::{
hex_literal::hex,
proofs::triehash::KeccakHasher,
trie::{BranchNodeCompact, TrieMask},
Account, StorageEntry, U256,
};
use reth_provider::{test_utils::create_test_provider_factory, DatabaseProviderRW};
use std::{
collections::{BTreeMap, HashMap},
ops::Mul,
str::FromStr,
sync::Arc,
};
fn insert_account(
tx: &impl DbTxMut,
address: Address,
account: Account,
storage: &BTreeMap<B256, U256>,
) {
let hashed_address = keccak256(address);
tx.put::<tables::HashedAccounts>(hashed_address, account).unwrap();
insert_storage(tx, hashed_address, storage);
}
fn insert_storage(tx: &impl DbTxMut, hashed_address: B256, storage: &BTreeMap<B256, U256>) {
for (k, v) in storage {
tx.put::<tables::HashedStorages>(
hashed_address,
StorageEntry { key: keccak256(k), value: *v },
)
.unwrap();
}
}
fn incremental_vs_full_root(inputs: &[&str], modified: &str) {
let factory = create_test_provider_factory();
let tx = factory.provider_rw().unwrap();
let hashed_address = B256::with_last_byte(1);
let mut hashed_storage_cursor =
tx.tx_ref().cursor_dup_write::<tables::HashedStorages>().unwrap();
let data = inputs.iter().map(|x| B256::from_str(x).unwrap());
let value = U256::from(0);
for key in data {
hashed_storage_cursor.upsert(hashed_address, StorageEntry { key, value }).unwrap();
}
// Generate the intermediate nodes on the receiving end of the channel
let (_, _, trie_updates) =
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();
let value = U256::from(1);
if hashed_storage_cursor.seek_by_key_subkey(hashed_address, modified_key).unwrap().is_some()
{
hashed_storage_cursor.delete_current().unwrap();
}
hashed_storage_cursor
.upsert(hashed_address, StorageEntry { key: modified_key, value })
.unwrap();
// 2. Calculate full merkle root
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
trie_updates.flush(tx.tx_ref()).unwrap();
// 3. Calculate the incremental root
let mut storage_changes = PrefixSetMut::default();
storage_changes.insert(Nibbles::unpack(modified_key));
let loader = StorageRoot::from_tx_hashed(tx.tx_ref(), hashed_address)
.with_prefix_set(storage_changes.freeze());
let incremental_root = loader.root().unwrap();
assert_eq!(modified_root, incremental_root);
}
#[test]
fn branch_node_child_changes() {
incremental_vs_full_root(
&[
"1000000000000000000000000000000000000000000000000000000000000000",
"1100000000000000000000000000000000000000000000000000000000000000",
"1110000000000000000000000000000000000000000000000000000000000000",
"1200000000000000000000000000000000000000000000000000000000000000",
"1220000000000000000000000000000000000000000000000000000000000000",
"1320000000000000000000000000000000000000000000000000000000000000",
],
"1200000000000000000000000000000000000000000000000000000000000000",
);
}
#[test]
fn arbitrary_storage_root() {
proptest!(ProptestConfig::with_cases(10), |(item: (Address, std::collections::BTreeMap<B256, U256>))| {
let (address, storage) = item;
let hashed_address = keccak256(address);
let factory = create_test_provider_factory();
let tx = factory.provider_rw().unwrap();
for (key, value) in &storage {
tx.tx_ref().put::<tables::HashedStorages>(
hashed_address,
StorageEntry { key: keccak256(key), value: *value },
)
.unwrap();
}
tx.commit().unwrap();
let tx = factory.provider_rw().unwrap();
let got = StorageRoot::from_tx(tx.tx_ref(), address).root().unwrap();
let expected = storage_root(storage.into_iter());
assert_eq!(expected, got);
});
}
#[test]
// This ensures we dont add empty accounts to the trie
fn test_empty_account() {
let state: State = BTreeMap::from([
(
Address::random(),
(
Account { nonce: 0, balance: U256::from(0), bytecode_hash: None },
BTreeMap::from([(B256::with_last_byte(0x4), U256::from(12))]),
),
),
(
Address::random(),
(
Account { nonce: 0, balance: U256::from(0), bytecode_hash: None },
BTreeMap::default(),
),
),
(
Address::random(),
(
Account {
nonce: 155,
balance: U256::from(414241124u32),
bytecode_hash: Some(keccak256("test")),
},
BTreeMap::from([
(B256::ZERO, U256::from(3)),
(B256::with_last_byte(2), U256::from(1)),
]),
),
),
]);
test_state_root_with_state(state);
}
#[test]
// This ensures we return an empty root when there are no storage entries
fn test_empty_storage_root() {
let factory = create_test_provider_factory();
let tx = factory.provider_rw().unwrap();
let address = Address::random();
let code = "el buen fla";
let account = Account {
nonce: 155,
balance: U256::from(414241124u32),
bytecode_hash: Some(keccak256(code)),
};
insert_account(tx.tx_ref(), address, account, &Default::default());
tx.commit().unwrap();
let tx = factory.provider_rw().unwrap();
let got = StorageRoot::from_tx(tx.tx_ref(), address).root().unwrap();
assert_eq!(got, EMPTY_ROOT_HASH);
}
#[test]
// This ensures that the walker goes over all the storage slots
fn test_storage_root() {
let factory = create_test_provider_factory();
let tx = factory.provider_rw().unwrap();
let address = Address::random();
let storage =
BTreeMap::from([(B256::ZERO, U256::from(3)), (B256::with_last_byte(2), U256::from(1))]);
let code = "el buen fla";
let account = Account {
nonce: 155,
balance: U256::from(414241124u32),
bytecode_hash: Some(keccak256(code)),
};
insert_account(tx.tx_ref(), address, account, &storage);
tx.commit().unwrap();
let tx = factory.provider_rw().unwrap();
let got = StorageRoot::from_tx(tx.tx_ref(), address).root().unwrap();
assert_eq!(storage_root(storage.into_iter()), got);
}
type State = BTreeMap<Address, (Account, BTreeMap<B256, U256>)>;
#[test]
fn arbitrary_state_root() {
proptest!(
ProptestConfig::with_cases(10), | (state: State) | {
test_state_root_with_state(state);
}
);
}
#[test]
fn arbitrary_state_root_with_progress() {
proptest!(
ProptestConfig::with_cases(10), | (state: State) | {
let hashed_entries_total = state.len() +
state.values().map(|(_, slots)| slots.len()).sum::<usize>();
let factory = create_test_provider_factory();
let tx = factory.provider_rw().unwrap();
for (address, (account, storage)) in &state {
insert_account(tx.tx_ref(), *address, *account, storage)
}
tx.commit().unwrap();
let tx = factory.provider_rw().unwrap();
let expected = state_root(state);
let threshold = 10;
let mut got = None;
let mut hashed_entries_walked = 0;
let mut intermediate_state: Option<Box<IntermediateStateRootState>> = None;
while got.is_none() {
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() {
StateRootProgress::Progress(state, walked, _) => {
intermediate_state = Some(state);
hashed_entries_walked += walked;
},
StateRootProgress::Complete(root, walked, _) => {
got = Some(root);
hashed_entries_walked += walked;
},
};
}
assert_eq!(expected, got.unwrap());
assert_eq!(hashed_entries_total, hashed_entries_walked)
}
);
}
fn test_state_root_with_state(state: State) {
let factory = create_test_provider_factory();
let tx = factory.provider_rw().unwrap();
for (address, (account, storage)) in &state {
insert_account(tx.tx_ref(), *address, *account, storage)
}
tx.commit().unwrap();
let expected = state_root(state);
let tx = factory.provider_rw().unwrap();
let got = StateRoot::from_tx(tx.tx_ref()).root().unwrap();
assert_eq!(expected, got);
}
fn encode_account(account: Account, storage_root: Option<B256>) -> Vec<u8> {
let account = TrieAccount::from((account, storage_root.unwrap_or(EMPTY_ROOT_HASH)));
let mut account_rlp = Vec::with_capacity(account.length());
account.encode(&mut account_rlp);
account_rlp
}
#[test]
fn storage_root_regression() {
let factory = create_test_provider_factory();
let tx = factory.provider_rw().unwrap();
// Some address whose hash starts with 0xB041
let address3 = Address::from_str("16b07afd1c635f77172e842a000ead9a2a222459").unwrap();
let key3 = keccak256(address3);
assert_eq!(key3[0], 0xB0);
assert_eq!(key3[1], 0x41);
let storage = BTreeMap::from(
[
("1200000000000000000000000000000000000000000000000000000000000000", 0x42),
("1400000000000000000000000000000000000000000000000000000000000000", 0x01),
("3000000000000000000000000000000000000000000000000000000000E00000", 0x127a89),
("3000000000000000000000000000000000000000000000000000000000E00001", 0x05),
]
.map(|(slot, val)| (B256::from_str(slot).unwrap(), U256::from(val))),
);
let mut hashed_storage_cursor =
tx.tx_ref().cursor_dup_write::<tables::HashedStorages>().unwrap();
for (hashed_slot, value) in storage.clone() {
hashed_storage_cursor.upsert(key3, StorageEntry { key: hashed_slot, value }).unwrap();
}
tx.commit().unwrap();
let tx = factory.provider_rw().unwrap();
let account3_storage_root = StorageRoot::from_tx(tx.tx_ref(), address3).root().unwrap();
let expected_root = storage_root_prehashed(storage);
assert_eq!(expected_root, account3_storage_root);
}
#[test]
fn account_and_storage_trie() {
let ether = U256::from(1e18);
let storage = BTreeMap::from(
[
("1200000000000000000000000000000000000000000000000000000000000000", 0x42),
("1400000000000000000000000000000000000000000000000000000000000000", 0x01),
("3000000000000000000000000000000000000000000000000000000000E00000", 0x127a89),
("3000000000000000000000000000000000000000000000000000000000E00001", 0x05),
]
.map(|(slot, val)| (B256::from_str(slot).unwrap(), U256::from(val))),
);
let factory = create_test_provider_factory();
let tx = factory.provider_rw().unwrap();
let mut hashed_account_cursor =
tx.tx_ref().cursor_write::<tables::HashedAccounts>().unwrap();
let mut hashed_storage_cursor =
tx.tx_ref().cursor_dup_write::<tables::HashedStorages>().unwrap();
let mut hash_builder = HashBuilder::default();
// Insert first account
let key1 =
B256::from_str("b000000000000000000000000000000000000000000000000000000000000000")
.unwrap();
let account1 = Account { nonce: 0, balance: U256::from(3).mul(ether), bytecode_hash: None };
hashed_account_cursor.upsert(key1, account1).unwrap();
hash_builder.add_leaf(Nibbles::unpack(key1), &encode_account(account1, None));
// Some address whose hash starts with 0xB040
let address2 = Address::from_str("7db3e81b72d2695e19764583f6d219dbee0f35ca").unwrap();
let key2 = keccak256(address2);
assert_eq!(key2[0], 0xB0);
assert_eq!(key2[1], 0x40);
let account2 = Account { nonce: 0, balance: ether, ..Default::default() };
hashed_account_cursor.upsert(key2, account2).unwrap();
hash_builder.add_leaf(Nibbles::unpack(key2), &encode_account(account2, None));
// Some address whose hash starts with 0xB041
let address3 = Address::from_str("16b07afd1c635f77172e842a000ead9a2a222459").unwrap();
let key3 = keccak256(address3);
assert_eq!(key3[0], 0xB0);
assert_eq!(key3[1], 0x41);
let code_hash =
B256::from_str("5be74cad16203c4905c068b012a2e9fb6d19d036c410f16fd177f337541440dd")
.unwrap();
let account3 =
Account { nonce: 0, balance: U256::from(2).mul(ether), bytecode_hash: Some(code_hash) };
hashed_account_cursor.upsert(key3, account3).unwrap();
for (hashed_slot, value) in storage {
if hashed_storage_cursor
.seek_by_key_subkey(key3, hashed_slot)
.unwrap()
.filter(|e| e.key == hashed_slot)
.is_some()
{
hashed_storage_cursor.delete_current().unwrap();
}
hashed_storage_cursor.upsert(key3, StorageEntry { key: hashed_slot, value }).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)),
);
let key4a =
B256::from_str("B1A0000000000000000000000000000000000000000000000000000000000000")
.unwrap();
let account4a =
Account { nonce: 0, balance: U256::from(4).mul(ether), ..Default::default() };
hashed_account_cursor.upsert(key4a, account4a).unwrap();
hash_builder.add_leaf(Nibbles::unpack(key4a), &encode_account(account4a, None));
let key5 =
B256::from_str("B310000000000000000000000000000000000000000000000000000000000000")
.unwrap();
let account5 =
Account { nonce: 0, balance: U256::from(8).mul(ether), ..Default::default() };
hashed_account_cursor.upsert(key5, account5).unwrap();
hash_builder.add_leaf(Nibbles::unpack(key5), &encode_account(account5, None));
let key6 =
B256::from_str("B340000000000000000000000000000000000000000000000000000000000000")
.unwrap();
let account6 =
Account { nonce: 0, balance: U256::from(1).mul(ether), ..Default::default() };
hashed_account_cursor.upsert(key6, account6).unwrap();
hash_builder.add_leaf(Nibbles::unpack(key6), &encode_account(account6, None));
// Populate account & storage trie DB tables
let expected_root =
B256::from_str("72861041bc90cd2f93777956f058a545412b56de79af5eb6b8075fe2eabbe015")
.unwrap();
let computed_expected_root: B256 = triehash::trie_root::<KeccakHasher, _, _, _>([
(key1, encode_account(account1, None)),
(key2, encode_account(account2, None)),
(key3, encode_account(account3, Some(account3_storage_root))),
(key4a, encode_account(account4a, None)),
(key5, encode_account(account5, None)),
(key6, encode_account(account6, None)),
]);
// Check computed trie root to ensure correctness
assert_eq!(computed_expected_root, expected_root);
// Check hash builder root
assert_eq!(hash_builder.root(), computed_expected_root);
// Check state root calculation from scratch
let (root, trie_updates) = StateRoot::from_tx(tx.tx_ref()).root_with_updates().unwrap();
assert_eq!(root, computed_expected_root);
// Check account trie
let mut account_updates = trie_updates
.iter()
.filter_map(|(k, v)| match (k, v) {
(TrieKey::AccountNode(nibbles), TrieOp::Update(node)) => Some((nibbles, node)),
_ => None,
})
.collect::<Vec<_>>();
account_updates.sort_unstable_by(|a, b| a.0.cmp(b.0));
assert_eq!(account_updates.len(), 2);
let (nibbles1a, node1a) = account_updates.first().unwrap();
assert_eq!(nibbles1a[..], [0xB]);
assert_eq!(node1a.state_mask, TrieMask::new(0b1011));
assert_eq!(node1a.tree_mask, TrieMask::new(0b0001));
assert_eq!(node1a.hash_mask, TrieMask::new(0b1001));
assert_eq!(node1a.root_hash, None);
assert_eq!(node1a.hashes.len(), 2);
let (nibbles2a, node2a) = account_updates.last().unwrap();
assert_eq!(nibbles2a[..], [0xB, 0x0]);
assert_eq!(node2a.state_mask, TrieMask::new(0b10001));
assert_eq!(node2a.tree_mask, TrieMask::new(0b00000));
assert_eq!(node2a.hash_mask, TrieMask::new(0b10000));
assert_eq!(node2a.root_hash, None);
assert_eq!(node2a.hashes.len(), 1);
// Check storage trie
let storage_updates = trie_updates
.iter()
.filter_map(|entry| match entry {
(TrieKey::StorageNode(_, nibbles), TrieOp::Update(node)) => Some((nibbles, node)),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(storage_updates.len(), 1);
let (nibbles3, node3) = storage_updates.first().unwrap();
assert!(nibbles3.is_empty());
assert_eq!(node3.state_mask, TrieMask::new(0b1010));
assert_eq!(node3.tree_mask, TrieMask::new(0b0000));
assert_eq!(node3.hash_mask, TrieMask::new(0b0010));
assert_eq!(node3.hashes.len(), 1);
assert_eq!(node3.root_hash, Some(account3_storage_root));
// Add an account
// Some address whose hash starts with 0xB1
let address4b = Address::from_str("4f61f2d5ebd991b85aa1677db97307caf5215c91").unwrap();
let key4b = keccak256(address4b);
assert_eq!(key4b.0[0], key4a.0[0]);
let account4b =
Account { nonce: 0, balance: U256::from(5).mul(ether), bytecode_hash: None };
hashed_account_cursor.upsert(key4b, account4b).unwrap();
let mut prefix_set = PrefixSetMut::default();
prefix_set.insert(Nibbles::unpack(key4b));
let expected_state_root =
B256::from_str("8e263cd4eefb0c3cbbb14e5541a66a755cad25bcfab1e10dd9d706263e811b28")
.unwrap();
let (root, trie_updates) = StateRoot::from_tx(tx.tx_ref())
.with_prefix_sets(TriePrefixSets {
account_prefix_set: prefix_set.freeze(),
..Default::default()
})
.root_with_updates()
.unwrap();
assert_eq!(root, expected_state_root);
let mut account_updates = trie_updates
.iter()
.filter_map(|entry| match entry {
(TrieKey::AccountNode(nibbles), TrieOp::Update(node)) => Some((nibbles, node)),
_ => None,
})
.collect::<Vec<_>>();
account_updates.sort_by(|a, b| a.0.cmp(b.0));
assert_eq!(account_updates.len(), 2);
let (nibbles1b, node1b) = account_updates.first().unwrap();
assert_eq!(nibbles1b[..], [0xB]);
assert_eq!(node1b.state_mask, TrieMask::new(0b1011));
assert_eq!(node1b.tree_mask, TrieMask::new(0b0001));
assert_eq!(node1b.hash_mask, TrieMask::new(0b1011));
assert_eq!(node1b.root_hash, None);
assert_eq!(node1b.hashes.len(), 3);
assert_eq!(node1a.hashes[0], node1b.hashes[0]);
assert_eq!(node1a.hashes[1], node1b.hashes[2]);
let (nibbles2b, node2b) = account_updates.last().unwrap();
assert_eq!(nibbles2b[..], [0xB, 0x0]);
assert_eq!(node2a, node2b);
tx.commit().unwrap();
{
let tx = factory.provider_rw().unwrap();
let mut hashed_account_cursor =
tx.tx_ref().cursor_write::<tables::HashedAccounts>().unwrap();
let account = hashed_account_cursor.seek_exact(key2).unwrap().unwrap();
hashed_account_cursor.delete_current().unwrap();
let mut account_prefix_set = PrefixSetMut::default();
account_prefix_set.insert(Nibbles::unpack(account.0));
let computed_expected_root: B256 = triehash::trie_root::<KeccakHasher, _, _, _>([
(key1, encode_account(account1, None)),
// DELETED: (key2, encode_account(account2, None)),
(key3, encode_account(account3, Some(account3_storage_root))),
(key4a, encode_account(account4a, None)),
(key4b, encode_account(account4b, None)),
(key5, encode_account(account5, None)),
(key6, encode_account(account6, None)),
]);
let (root, trie_updates) = StateRoot::from_tx(tx.tx_ref())
.with_prefix_sets(TriePrefixSets {
account_prefix_set: account_prefix_set.freeze(),
..Default::default()
})
.root_with_updates()
.unwrap();
assert_eq!(root, computed_expected_root);
assert_eq!(trie_updates.len(), 7);
assert_eq!(trie_updates.iter().filter(|(_, op)| op.is_update()).count(), 2);
let account_updates = trie_updates
.iter()
.filter_map(|entry| match entry {
(TrieKey::AccountNode(nibbles), TrieOp::Update(node)) => Some((nibbles, node)),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(account_updates.len(), 1);
let (nibbles1c, node1c) = account_updates.first().unwrap();
assert_eq!(nibbles1c[..], [0xB]);
assert_eq!(node1c.state_mask, TrieMask::new(0b1011));
assert_eq!(node1c.tree_mask, TrieMask::new(0b0000));
assert_eq!(node1c.hash_mask, TrieMask::new(0b1011));
assert_eq!(node1c.root_hash, None);
assert_eq!(node1c.hashes.len(), 3);
assert_ne!(node1c.hashes[0], node1b.hashes[0]);
assert_eq!(node1c.hashes[1], node1b.hashes[1]);
assert_eq!(node1c.hashes[2], node1b.hashes[2]);
}
{
let tx = factory.provider_rw().unwrap();
let mut hashed_account_cursor =
tx.tx_ref().cursor_write::<tables::HashedAccounts>().unwrap();
let account2 = hashed_account_cursor.seek_exact(key2).unwrap().unwrap();
hashed_account_cursor.delete_current().unwrap();
let account3 = hashed_account_cursor.seek_exact(key3).unwrap().unwrap();
hashed_account_cursor.delete_current().unwrap();
let mut account_prefix_set = PrefixSetMut::default();
account_prefix_set.insert(Nibbles::unpack(account2.0));
account_prefix_set.insert(Nibbles::unpack(account3.0));
let computed_expected_root: B256 = triehash::trie_root::<KeccakHasher, _, _, _>([
(key1, encode_account(account1, None)),
// DELETED: (key2, encode_account(account2, None)),
// DELETED: (key3, encode_account(account3, Some(account3_storage_root))),
(key4a, encode_account(account4a, None)),
(key4b, encode_account(account4b, None)),
(key5, encode_account(account5, None)),
(key6, encode_account(account6, None)),
]);
let (root, trie_updates) = StateRoot::from_tx(tx.tx_ref())
.with_prefix_sets(TriePrefixSets {
account_prefix_set: account_prefix_set.freeze(),
..Default::default()
})
.root_with_updates()
.unwrap();
assert_eq!(root, computed_expected_root);
assert_eq!(trie_updates.len(), 6);
assert_eq!(trie_updates.iter().filter(|(_, op)| op.is_update()).count(), 1); // no storage root update
let account_updates = trie_updates
.iter()
.filter_map(|entry| match entry {
(TrieKey::AccountNode(nibbles), TrieOp::Update(node)) => Some((nibbles, node)),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(account_updates.len(), 1);
let (nibbles1d, node1d) = account_updates.first().unwrap();
assert_eq!(nibbles1d[..], [0xB]);
assert_eq!(node1d.state_mask, TrieMask::new(0b1011));
assert_eq!(node1d.tree_mask, TrieMask::new(0b0000));
assert_eq!(node1d.hash_mask, TrieMask::new(0b1010));
assert_eq!(node1d.root_hash, None);
assert_eq!(node1d.hashes.len(), 2);
assert_eq!(node1d.hashes[0], node1b.hashes[1]);
assert_eq!(node1d.hashes[1], node1b.hashes[2]);
}
}
#[test]
fn account_trie_around_extension_node() {
let factory = create_test_provider_factory();
let tx = factory.provider_rw().unwrap();
let expected = extension_node_trie(&tx);
let (got, updates) = StateRoot::from_tx(tx.tx_ref()).root_with_updates().unwrap();
assert_eq!(expected, got);
// Check account trie
let account_updates = updates
.iter()
.filter_map(|entry| match entry {
(TrieKey::AccountNode(nibbles), TrieOp::Update(node)) => {
Some((nibbles.0.clone(), node.clone()))
}
_ => None,
})
.collect::<HashMap<_, _>>();
assert_trie_updates(&account_updates);
}
#[test]
fn account_trie_around_extension_node_with_dbtrie() {
let factory = create_test_provider_factory();
let tx = factory.provider_rw().unwrap();
let expected = extension_node_trie(&tx);
let (got, updates) = StateRoot::from_tx(tx.tx_ref()).root_with_updates().unwrap();
assert_eq!(expected, got);
updates.flush(tx.tx_ref()).unwrap();
// read the account updates from the db
let mut accounts_trie = tx.tx_ref().cursor_read::<tables::AccountsTrie>().unwrap();
let walker = accounts_trie.walk(None).unwrap();
let account_updates = walker
.into_iter()
.map(|item| {
let (key, node) = item.unwrap();
(key.0, node.0)
})
.collect();
assert_trie_updates(&account_updates);
}
proptest! {
#![proptest_config(ProptestConfig {
cases: 128, ..ProptestConfig::default()
})]
#[test]
fn fuzz_state_root_incremental(account_changes: [BTreeMap<B256, U256>; 5]) {
let factory = create_test_provider_factory();
let tx = factory.provider_rw().unwrap();
let mut hashed_account_cursor = tx.tx_ref().cursor_write::<tables::HashedAccounts>().unwrap();
let mut state = BTreeMap::default();
for accounts in account_changes {
let should_generate_changeset = !state.is_empty();
let mut changes = PrefixSetMut::default();
for (hashed_address, balance) in accounts.clone() {
hashed_account_cursor.upsert(hashed_address, Account { balance, ..Default::default() }).unwrap();
if should_generate_changeset {
changes.insert(Nibbles::unpack(hashed_address));
}
}
let (state_root, trie_updates) = StateRoot::from_tx(tx.tx_ref())
.with_prefix_sets(TriePrefixSets { account_prefix_set: changes.freeze(), ..Default::default() })
.root_with_updates()
.unwrap();
state.append(&mut accounts.clone());
let expected_root = state_root_prehashed(
state.iter().map(|(&key, &balance)| (key, (Account { balance, ..Default::default() }, std::iter::empty())))
);
assert_eq!(expected_root, state_root);
trie_updates.flush(tx.tx_ref()).unwrap();
}
}
}
#[test]
fn storage_trie_around_extension_node() {
let factory = create_test_provider_factory();
let tx = factory.provider_rw().unwrap();
let hashed_address = B256::random();
let (expected_root, expected_updates) = extension_node_storage_trie(&tx, hashed_address);
let (got, _, updates) =
StorageRoot::from_tx_hashed(tx.tx_ref(), hashed_address).root_with_updates().unwrap();
assert_eq!(expected_root, got);
// Check account trie
let storage_updates = updates
.iter()
.filter_map(|entry| match entry {
(TrieKey::StorageNode(_, nibbles), TrieOp::Update(node)) => {
Some((nibbles.0.clone(), node.clone()))
}
_ => None,
})
.collect::<HashMap<_, _>>();
assert_eq!(expected_updates, storage_updates);
assert_trie_updates(&storage_updates);
}
fn extension_node_storage_trie(
tx: &DatabaseProviderRW<Arc<TempDatabase<DatabaseEnv>>>,
hashed_address: B256,
) -> (B256, HashMap<Nibbles, BranchNodeCompact>) {
let value = U256::from(1);
let mut hashed_storage = tx.tx_ref().cursor_write::<tables::HashedStorages>().unwrap();
let mut hb = HashBuilder::default().with_updates(true);
for key in [
hex!("30af561000000000000000000000000000000000000000000000000000000000"),
hex!("30af569000000000000000000000000000000000000000000000000000000000"),
hex!("30af650000000000000000000000000000000000000000000000000000000000"),
hex!("30af6f0000000000000000000000000000000000000000000000000000000000"),
hex!("30af8f0000000000000000000000000000000000000000000000000000000000"),
hex!("3100000000000000000000000000000000000000000000000000000000000000"),
] {
hashed_storage
.upsert(hashed_address, StorageEntry { key: B256::new(key), value })
.unwrap();
hb.add_leaf(Nibbles::unpack(key), &alloy_rlp::encode_fixed_size(&value));
}
let root = hb.root();
let (_, updates) = hb.split();
(root, updates)
}
fn extension_node_trie(tx: &DatabaseProviderRW<Arc<TempDatabase<DatabaseEnv>>>) -> B256 {
let a =
Account { nonce: 0, balance: U256::from(1u64), bytecode_hash: Some(B256::random()) };
let val = encode_account(a, None);
let mut hashed_accounts = tx.tx_ref().cursor_write::<tables::HashedAccounts>().unwrap();
let mut hb = HashBuilder::default();
for key in [
hex!("30af561000000000000000000000000000000000000000000000000000000000"),
hex!("30af569000000000000000000000000000000000000000000000000000000000"),
hex!("30af650000000000000000000000000000000000000000000000000000000000"),
hex!("30af6f0000000000000000000000000000000000000000000000000000000000"),
hex!("30af8f0000000000000000000000000000000000000000000000000000000000"),
hex!("3100000000000000000000000000000000000000000000000000000000000000"),
] {
hashed_accounts.upsert(B256::new(key), a).unwrap();
hb.add_leaf(Nibbles::unpack(key), &val);
}
hb.root()
}
fn assert_trie_updates(account_updates: &HashMap<Nibbles, BranchNodeCompact>) {
assert_eq!(account_updates.len(), 2);
let node = account_updates.get(&[0x3][..]).unwrap();
let expected = BranchNodeCompact::new(0b0011, 0b0001, 0b0000, vec![], None);
assert_eq!(node, &expected);
let node = account_updates.get(&[0x3, 0x0, 0xA, 0xF][..]).unwrap();
assert_eq!(node.state_mask, TrieMask::new(0b101100000));
assert_eq!(node.tree_mask, TrieMask::new(0b000000000));
assert_eq!(node.hash_mask, TrieMask::new(0b001000000));
assert_eq!(node.root_hash, None);
assert_eq!(node.hashes.len(), 1);
}
}