diff --git a/Cargo.lock b/Cargo.lock index 9ce3d4ee4c..15555c2d83 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -649,17 +649,6 @@ dependencies = [ "inout", ] -[[package]] -name = "cita_trie" -version = "4.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "afe7baab47f510f52ca8dc9c0eb9082020c627c7f22285bea30edc3511f7ee29" -dependencies = [ - "hasher", - "parking_lot 0.12.1", - "rlp", -] - [[package]] name = "clang-sys" version = "1.6.0" @@ -2395,15 +2384,6 @@ dependencies = [ "serde", ] -[[package]] -name = "hasher" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbbba678b6567f27ce22870d951f4208e1dc2fef64993bd4521b1d497ef8a3aa" -dependencies = [ - "tiny-keccak", -] - [[package]] name = "hashers" version = "1.0.1" @@ -5001,24 +4981,18 @@ dependencies = [ name = "reth-provider" version = "0.1.0" dependencies = [ - "assert_matches", "auto_impl", - "cita_trie", - "hasher", "itertools", "parking_lot 0.12.1", - "proptest", - "reth-codecs", "reth-db", "reth-interfaces", "reth-primitives", "reth-revm-primitives", "reth-rlp", - "reth-tracing", + "reth-trie", "thiserror", "tokio", "tracing", - "triehash", ] [[package]] @@ -5262,6 +5236,7 @@ dependencies = [ "reth-provider", "reth-revm", "reth-rlp", + "reth-trie", "thiserror", "tokio", "tokio-stream", @@ -5322,8 +5297,12 @@ dependencies = [ "derive_more", "hex", "proptest", + "reth-db", + "reth-interfaces", "reth-primitives", + "reth-provider", "reth-rlp", + "thiserror", "tokio", "tokio-stream", "tracing", diff --git a/crates/primitives/src/lib.rs b/crates/primitives/src/lib.rs index c0b7d47120..5b4539d53f 100644 --- a/crates/primitives/src/lib.rs +++ b/crates/primitives/src/lib.rs @@ -67,7 +67,7 @@ pub use peer::{PeerId, WithPeerId}; pub use receipt::{Receipt, ReceiptWithBloom, ReceiptWithBloomRef}; pub use revm_primitives::JumpMap; pub use serde_helper::JsonU256; -pub use storage::{StorageEntry, StorageTrieEntry}; +pub use storage::StorageEntry; pub use transaction::{ util::secp256k1::sign_message, AccessList, AccessListItem, AccessListWithGasUsed, FromRecoveredTransaction, IntoRecoveredTransaction, InvalidTransactionError, Signature, diff --git a/crates/primitives/src/storage.rs b/crates/primitives/src/storage.rs index 67e027c34d..498045424d 100644 --- a/crates/primitives/src/storage.rs +++ b/crates/primitives/src/storage.rs @@ -37,34 +37,3 @@ impl Compact for StorageEntry { (Self { key, value }, out) } } - -/// Account storage trie node. -#[derive_arbitrary(compact)] -#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize)] -pub struct StorageTrieEntry { - /// Hashed storage key. - pub hash: H256, - /// Encoded node. - pub node: Vec, -} - -// NOTE: Removing main_codec and manually encode subkey -// and compress second part of the value. If we have compression -// over whole value (Even SubKey) that would mess up fetching of values with seek_by_key_subkey -impl Compact for StorageTrieEntry { - fn to_compact(self, buf: &mut impl bytes::BufMut) -> usize { - // for now put full bytes and later compress it. - buf.put_slice(&self.hash.to_fixed_bytes()[..]); - buf.put_slice(&self.node[..]); - self.node.len() + 32 - } - - fn from_compact(buf: &[u8], len: usize) -> (Self, &[u8]) - where - Self: Sized, - { - let key = H256::from_slice(&buf[..32]); - let node = Vec::from(&buf[32..len]); - (Self { hash: key, node }, &buf[len..]) - } -} diff --git a/crates/primitives/src/trie/branch_node.rs b/crates/primitives/src/trie/branch_node.rs index e8fc6d6ada..f56bf0e14d 100644 --- a/crates/primitives/src/trie/branch_node.rs +++ b/crates/primitives/src/trie/branch_node.rs @@ -40,12 +40,14 @@ pub struct BranchNodeCompact { impl BranchNodeCompact { /// Creates a new [BranchNodeCompact] from the given parameters. pub fn new( - state_mask: TrieMask, - tree_mask: TrieMask, - hash_mask: TrieMask, + state_mask: impl Into, + tree_mask: impl Into, + hash_mask: impl Into, hashes: Vec, root_hash: Option, ) -> Self { + let (state_mask, tree_mask, hash_mask) = + (state_mask.into(), tree_mask.into(), hash_mask.into()); assert!(tree_mask.is_subset_of(&state_mask)); assert!(hash_mask.is_subset_of(&state_mask)); assert_eq!(hash_mask.count_ones() as usize, hashes.len()); @@ -129,9 +131,9 @@ mod tests { #[test] fn node_encoding() { let n = BranchNodeCompact::new( - 0xf607.into(), - 0x0005.into(), - 0x4004.into(), + 0xf607, + 0x0005, + 0x4004, vec![ hex!("90d53cd810cc5d4243766cd4451e7b9d14b736a1148b26b3baac7617f617d321").into(), hex!("cc35c964dda53ba6c0b87798073a9628dbc9cd26b5cce88eb69655a9c609caf1").into(), diff --git a/crates/primitives/src/trie/mask.rs b/crates/primitives/src/trie/mask.rs index 1eb3fe67ff..1a53cd1db1 100644 --- a/crates/primitives/src/trie/mask.rs +++ b/crates/primitives/src/trie/mask.rs @@ -45,7 +45,7 @@ impl TrieMask { } /// Returns `true` if a given bit is set in a mask. - pub fn is_bit_set(&self, index: i32) -> bool { + pub fn is_bit_set(&self, index: u8) -> bool { self.0 & (1u16 << index) != 0 } diff --git a/crates/primitives/src/trie/mod.rs b/crates/primitives/src/trie/mod.rs index 03fc48b770..771bb94a81 100644 --- a/crates/primitives/src/trie/mod.rs +++ b/crates/primitives/src/trie/mod.rs @@ -1,10 +1,13 @@ //! Collection of trie related types. -mod nibbles; -pub use nibbles::{StoredNibbles, StoredNibblesSubKey}; - mod branch_node; -pub use branch_node::BranchNodeCompact; - mod mask; -pub use mask::TrieMask; +mod nibbles; +mod storage; + +pub use self::{ + branch_node::BranchNodeCompact, + mask::TrieMask, + nibbles::{StoredNibbles, StoredNibblesSubKey}, + storage::StorageTrieEntry, +}; diff --git a/crates/primitives/src/trie/storage.rs b/crates/primitives/src/trie/storage.rs new file mode 100644 index 0000000000..57d243fe42 --- /dev/null +++ b/crates/primitives/src/trie/storage.rs @@ -0,0 +1,33 @@ +use super::{BranchNodeCompact, StoredNibblesSubKey}; +use reth_codecs::Compact; +use serde::{Deserialize, Serialize}; + +/// Account storage trie node. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, PartialOrd, Ord)] +pub struct StorageTrieEntry { + /// The nibbles of the intermediate node + pub nibbles: StoredNibblesSubKey, + /// Encoded node. + pub node: BranchNodeCompact, +} + +// NOTE: Removing main_codec and manually encode subkey +// and compress second part of the value. If we have compression +// over whole value (Even SubKey) that would mess up fetching of values with seek_by_key_subkey +impl Compact for StorageTrieEntry { + fn to_compact(self, buf: &mut impl bytes::BufMut) -> usize { + let nibbles_len = self.nibbles.to_compact(buf); + let node_len = self.node.to_compact(buf); + nibbles_len + node_len + } + + fn from_compact(buf: &[u8], len: usize) -> (Self, &[u8]) + where + Self: Sized, + { + let (nibbles, buf) = StoredNibblesSubKey::from_compact(buf, 33); + let (node, buf) = BranchNodeCompact::from_compact(buf, len - 33); + let this = Self { nibbles, node }; + (this, buf) + } +} diff --git a/crates/rpc/rpc-builder/tests/it/http.rs b/crates/rpc/rpc-builder/tests/it/http.rs index 3a5259d308..76d7c2c27b 100644 --- a/crates/rpc/rpc-builder/tests/it/http.rs +++ b/crates/rpc/rpc-builder/tests/it/http.rs @@ -73,7 +73,6 @@ where EthApiClient::block_by_number(client, block_number, false).await.unwrap(); EthApiClient::block_transaction_count_by_number(client, block_number).await.unwrap(); EthApiClient::block_transaction_count_by_hash(client, hash).await.unwrap(); - EthApiClient::get_proof(client, address, vec![], None).await.unwrap(); EthApiClient::block_uncles_count_by_hash(client, hash).await.unwrap(); EthApiClient::block_uncles_count_by_number(client, block_number).await.unwrap(); EthApiClient::uncle_by_block_hash_and_index(client, hash, index).await.unwrap(); @@ -100,6 +99,9 @@ where EthApiClient::submit_hashrate(client, U256::default(), H256::default()).await.unwrap(); // Unimplemented + assert!(is_unimplemented( + EthApiClient::get_proof(client, address, vec![], None).await.err().unwrap() + )); assert!(is_unimplemented(EthApiClient::author(client).await.err().unwrap())); assert!(is_unimplemented(EthApiClient::gas_price(client).await.err().unwrap())); assert!(is_unimplemented(EthApiClient::max_priority_fee_per_gas(client).await.err().unwrap())); diff --git a/crates/rpc/rpc/src/eth/api/server.rs b/crates/rpc/rpc/src/eth/api/server.rs index 228d4296a6..c2a6349a7f 100644 --- a/crates/rpc/rpc/src/eth/api/server.rs +++ b/crates/rpc/rpc/src/eth/api/server.rs @@ -422,19 +422,21 @@ where /// Handler for: `eth_getProof` async fn get_proof( &self, - address: Address, - keys: Vec, - block_number: Option, + _address: Address, + _keys: Vec, + _block_number: Option, ) -> Result { - trace!(target: "rpc::eth", ?address, ?keys, ?block_number, "Serving eth_getProof"); - let res = EthApi::get_proof(self, address, keys, block_number); + // TODO: uncomment when implemented + // trace!(target: "rpc::eth", ?address, ?keys, ?block_number, "Serving eth_getProof"); + // let res = EthApi::get_proof(self, address, keys, block_number); - Ok(res.map_err(|e| match e { - EthApiError::InvalidBlockRange => { - internal_rpc_err("eth_getProof is unimplemented for historical blocks") - } - _ => e.into(), - })?) + // Ok(res.map_err(|e| match e { + // EthApiError::InvalidBlockRange => { + // internal_rpc_err("eth_getProof is unimplemented for historical blocks") + // } + // _ => e.into(), + // })?) + Err(internal_rpc_err("unimplemented")) } } diff --git a/crates/stages/Cargo.toml b/crates/stages/Cargo.toml index e22ee80b2d..1bbb19ec82 100644 --- a/crates/stages/Cargo.toml +++ b/crates/stages/Cargo.toml @@ -21,6 +21,7 @@ reth-db = { path = "../storage/db" } reth-codecs = { path = "../storage/codecs" } reth-provider = { path = "../storage/provider" } reth-metrics-derive = { path = "../metrics/metrics-derive" } +reth-trie = { path = "../trie" } # async tokio = { version = "1.21.2", features = ["sync"] } @@ -50,6 +51,7 @@ reth-eth-wire = { path = "../net/eth-wire" } # TODO(o reth-executor = { path = "../executor" } reth-rlp = { path = "../rlp" } reth-revm = { path = "../revm" } +reth-trie = { path = "../trie", features = ["test-utils"] } itertools = "0.10.5" tokio = { version = "*", features = ["rt", "sync", "macros"] } diff --git a/crates/stages/benches/setup/mod.rs b/crates/stages/benches/setup/mod.rs index eb587f893d..2f51462298 100644 --- a/crates/stages/benches/setup/mod.rs +++ b/crates/stages/benches/setup/mod.rs @@ -10,12 +10,12 @@ use reth_interfaces::test_utils::generators::{ random_transition_range, }; use reth_primitives::{Account, Address, SealedBlock, H256}; -use reth_provider::trie::DBTrieLoader; use reth_stages::{ stages::{AccountHashingStage, StorageHashingStage}, test_utils::TestTransaction, ExecInput, Stage, UnwindInput, }; +use reth_trie::StateRoot; use std::{ collections::BTreeMap, ops::Deref, @@ -70,9 +70,6 @@ pub(crate) fn unwind_hashes>>( StorageHashingStage::default().unwind(&mut db_tx, unwind).await.unwrap(); AccountHashingStage::default().unwind(&mut db_tx, unwind).await.unwrap(); - let target_root = db_tx.get_header(unwind.unwind_to).unwrap().state_root; - let _ = db_tx.delete::(target_root, None); - // Clear previous run stage.unwind(&mut db_tx, unwind).await.unwrap(); @@ -124,8 +121,7 @@ pub(crate) fn txs_testdata(num_blocks: u64) -> PathBuf { tx.insert_accounts_and_storages(start_state.clone()).unwrap(); // make first block after genesis have valid state root - let root = - DBTrieLoader::new(tx.inner().deref()).calculate_root().and_then(|e| e.root()).unwrap(); + let root = StateRoot::new(tx.inner().deref()).root(None).unwrap(); let second_block = blocks.get_mut(1).unwrap(); let cloned_second = second_block.clone(); let mut updated_header = cloned_second.header.unseal(); @@ -146,18 +142,11 @@ pub(crate) fn txs_testdata(num_blocks: u64) -> PathBuf { // make last block have valid state root let root = { let mut tx_mut = tx.inner(); - let root = - DBTrieLoader::new(tx_mut.deref()).calculate_root().and_then(|e| e.root()).unwrap(); + let root = StateRoot::new(tx.inner().deref()).root(None).unwrap(); tx_mut.commit().unwrap(); root }; - tx.query(|tx| { - assert!(tx.get::(root)?.is_some()); - Ok(()) - }) - .unwrap(); - let last_block = blocks.last_mut().unwrap(); let cloned_last = last_block.clone(); let mut updated_header = cloned_last.header.unseal(); diff --git a/crates/stages/src/stages/merkle.rs b/crates/stages/src/stages/merkle.rs index 85722bee84..e684c47cd9 100644 --- a/crates/stages/src/stages/merkle.rs +++ b/crates/stages/src/stages/merkle.rs @@ -1,10 +1,9 @@ use crate::{ExecInput, ExecOutput, Stage, StageError, StageId, UnwindInput, UnwindOutput}; -use reth_db::{database::Database, tables, transaction::DbTx}; +use reth_db::{database::Database, tables, transaction::DbTxMut}; use reth_interfaces::consensus; -use reth_provider::{ - trie::{DBTrieLoader, TrieProgress}, - Transaction, -}; +use reth_primitives::{BlockNumber, H256}; +use reth_provider::Transaction; +use reth_trie::StateRoot; use std::{fmt::Debug, ops::DerefMut}; use tracing::*; @@ -65,6 +64,24 @@ impl MerkleStage { pub fn default_unwind() -> Self { Self::Unwind } + + /// Check that the computed state root matches the expected. + fn validate_state_root( + &self, + got: H256, + expected: H256, + target_block: BlockNumber, + ) -> Result<(), StageError> { + if got == expected { + Ok(()) + } else { + warn!(target: "sync::stages::merkle", ?target_block, ?got, ?expected, "Block's root state failed verification"); + Err(StageError::Validation { + block: target_block, + error: consensus::ConsensusError::BodyStateRootDiff { got, expected }, + }) + } + } } #[async_trait::async_trait] @@ -108,43 +125,23 @@ impl Stage for MerkleStage { let trie_root = if from_transition == to_transition { block_root + } else if to_transition - from_transition > threshold || stage_progress == 0 { + // if there are more blocks than threshold it is faster to rebuild the trie + debug!(target: "sync::stages::merkle::exec", current = ?stage_progress, target = ?previous_stage_progress, "Rebuilding trie"); + tx.clear::()?; + tx.clear::()?; + StateRoot::new(tx.deref_mut()).root(None).map_err(|e| StageError::Fatal(Box::new(e)))? } else { - let res = if to_transition - from_transition > threshold || stage_progress == 0 { - debug!(target: "sync::stages::merkle::exec", current = ?stage_progress, target = ?previous_stage_progress, "Rebuilding trie"); - // if there are more blocks than threshold it is faster to rebuild the trie - let mut loader = DBTrieLoader::new(tx.deref_mut()); - loader.calculate_root().map_err(|e| StageError::Fatal(Box::new(e)))? - } else { - debug!(target: "sync::stages::merkle::exec", current = ?stage_progress, target = ?previous_stage_progress, "Updating trie"); - // Iterate over changeset (similar to Hashing stages) and take new values - let current_root = tx.get_header(stage_progress)?.state_root; - let mut loader = DBTrieLoader::new(tx.deref_mut()); - loader - .update_root(current_root, from_transition..to_transition) - .map_err(|e| StageError::Fatal(Box::new(e)))? - }; - - match res { - TrieProgress::Complete(root) => root, - TrieProgress::InProgress(_) => { - return Ok(ExecOutput { stage_progress, done: false }) - } - } + debug!(target: "sync::stages::merkle::exec", current = ?stage_progress, target = + ?previous_stage_progress, "Updating trie"); // Iterate over + StateRoot::incremental_root(tx.deref_mut(), from_transition..to_transition, None) + .map_err(|e| StageError::Fatal(Box::new(e)))? }; - if block_root != trie_root { - warn!(target: "sync::stages::merkle::exec", ?previous_stage_progress, got = ?trie_root, expected = ?block_root, "Block's root state failed verification"); - return Err(StageError::Validation { - block: previous_stage_progress, - error: consensus::ConsensusError::BodyStateRootDiff { - got: trie_root, - expected: block_root, - }, - }) - } + self.validate_state_root(trie_root, block_root, previous_stage_progress)?; info!(target: "sync::stages::merkle::exec", "Stage finished"); - Ok(ExecOutput { stage_progress: input.previous_stage_progress(), done: true }) + Ok(ExecOutput { stage_progress: previous_stage_progress, done: true }) } /// Unwind the stage. @@ -158,49 +155,24 @@ impl Stage for MerkleStage { return Ok(UnwindOutput { stage_progress: input.unwind_to }) } - let target_root = tx.get_header(input.unwind_to)?.state_root; - - // If the merkle stage fails to execute, the trie changes weren't committed - // and the root stayed the same - if tx.get::(target_root)?.is_some() { - info!(target: "sync::stages::merkle::unwind", "Stage skipped"); + if input.unwind_to == 0 { + tx.clear::()?; + tx.clear::()?; return Ok(UnwindOutput { stage_progress: input.unwind_to }) } - let current_root = tx.get_header(input.stage_progress)?.state_root; let from_transition = tx.get_block_transition(input.unwind_to)?; let to_transition = tx.get_block_transition(input.stage_progress)?; - let mut loader = DBTrieLoader::new(tx.deref_mut()); - let block_root = loop { - match loader - .update_root(current_root, from_transition..to_transition) - .map_err(|e| StageError::Fatal(Box::new(e)))? - { - TrieProgress::Complete(root) => break root, - TrieProgress::InProgress(_) => { - // Save the loader's progress & drop it to allow committing to the database, - // otherwise we're hitting the borrow checker - let progress = loader.current; - let _ = loader; - tx.commit()?; - // Reinstantiate the loader from where it was left off. - loader = DBTrieLoader::new(tx.deref_mut()); - loader.current = progress; - } - } - }; - - if block_root != target_root { - let unwind_to = input.unwind_to; - warn!(target: "sync::stages::merkle::unwind", ?unwind_to, got = ?block_root, expected = ?target_root, "Block's root state failed verification"); - return Err(StageError::Validation { - block: unwind_to, - error: consensus::ConsensusError::BodyStateRootDiff { - got: block_root, - expected: target_root, - }, - }) + // Unwind trie only if there are transitions + if from_transition < to_transition { + let block_root = + StateRoot::incremental_root(tx.deref_mut(), from_transition..to_transition, None) + .map_err(|e| StageError::Fatal(Box::new(e)))?; + let target_root = tx.get_header(input.unwind_to)?.state_root; + self.validate_state_root(block_root, target_root, input.unwind_to)?; + } else { + info!(target: "sync::stages::merkle::unwind", "Nothing to unwind"); } info!(target: "sync::stages::merkle::unwind", "Stage finished"); @@ -217,17 +189,16 @@ mod tests { }; use assert_matches::assert_matches; use reth_db::{ - cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO, DbDupCursorRW}, - database::DatabaseGAT, - mdbx::{Env, WriteMap}, + cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO}, tables, transaction::{DbTx, DbTxMut}, }; use reth_interfaces::test_utils::generators::{ random_block, random_block_range, random_contract_account_range, random_transition_range, }; - use reth_primitives::{keccak256, Account, Address, SealedBlock, StorageEntry, H256, U256}; - use std::{collections::BTreeMap, ops::Deref}; + use reth_primitives::{keccak256, SealedBlock, StorageEntry, H256, U256}; + use reth_trie::test_utils::{state_root, state_root_prehashed}; + use std::collections::BTreeMap; stage_test_suite_ext!(MerkleTestRunner, merkle); @@ -288,12 +259,6 @@ mod tests { assert!(runner.validate_execution(input, result.ok()).is_ok(), "execution validation"); } - fn create_trie_loader<'tx, 'db>( - tx: &'tx Transaction<'db, Env>, - ) -> DBTrieLoader<'tx, as DatabaseGAT<'db>>::TXMut> { - DBTrieLoader::new(tx.deref()) - } - struct MerkleTestRunner { tx: TestTransaction, clean_threshold: u64, @@ -325,24 +290,30 @@ mod tests { let stage_progress = input.stage_progress.unwrap_or_default(); let end = input.previous_stage_progress() + 1; - let n_accounts = 31; - let accounts = random_contract_account_range(&mut (0..n_accounts)) + let num_of_accounts = 31; + let accounts = random_contract_account_range(&mut (0..num_of_accounts)) .into_iter() .collect::>(); + self.tx.insert_accounts_and_storages( + accounts.iter().map(|(addr, acc)| (*addr, (*acc, std::iter::empty()))), + )?; + let SealedBlock { header, body, ommers, withdrawals } = random_block(stage_progress, None, Some(0), None); let mut header = header.unseal(); - header.state_root = - self.generate_initial_trie(accounts.iter().map(|(k, v)| (*k, *v)))?; + header.state_root = state_root( + accounts + .clone() + .into_iter() + .map(|(address, account)| (address, (account, std::iter::empty()))), + ); let sealed_head = SealedBlock { header: header.seal_slow(), body, ommers, withdrawals }; let head_hash = sealed_head.hash(); let mut blocks = vec![sealed_head]; - blocks.extend(random_block_range((stage_progress + 1)..end, head_hash, 0..3)); - self.tx.insert_blocks(blocks.iter(), None)?; let (transitions, final_state) = random_transition_range( @@ -351,13 +322,30 @@ mod tests { 0..3, 0..256, ); - self.tx.insert_transitions(transitions, None)?; - self.tx.insert_accounts_and_storages(final_state)?; + // Calculate state root + let root = self.tx.query(|tx| { + let mut accounts = BTreeMap::default(); + let mut accounts_cursor = tx.cursor_read::()?; + let mut storage_cursor = tx.cursor_dup_read::()?; + for entry in accounts_cursor.walk_range(..)? { + let (key, account) = entry?; + let storage_entries = + storage_cursor.walk_dup(Some(key), None)?.collect::, _>>()?; + let storage = storage_entries + .into_iter() + .filter(|(_, v)| v.value != U256::ZERO) + .map(|(_, v)| (v.key, v.value)) + .collect::>(); + accounts.insert(key, (account, storage)); + } + + Ok(state_root_prehashed(accounts.into_iter())) + })?; + let last_block_number = end - 1; - let root = self.state_root()?; self.tx.commit(|tx| { let mut last_header = tx.get::(last_block_number)?.unwrap(); last_header.state_root = root; @@ -369,17 +357,11 @@ mod tests { fn validate_execution( &self, - input: ExecInput, - output: Option, + _input: ExecInput, + _output: Option, ) -> Result<(), TestRunnerError> { - if let Some(output) = output { - let start_block = input.stage_progress.unwrap_or_default() + 1; - let end_block = output.stage_progress; - if start_block > end_block { - return Ok(()) - } - } - self.check_root(input.previous_stage_progress()) + // The execution is validated within the stage + Ok(()) } } @@ -394,14 +376,15 @@ mod tests { self.tx .commit(|tx| { - let mut changeset_cursor = + let mut storage_changesets_cursor = tx.cursor_dup_read::().unwrap(); - let mut hash_cursor = tx.cursor_dup_write::().unwrap(); - - let mut rev_changeset_walker = changeset_cursor.walk_back(None).unwrap(); + let mut storage_cursor = + tx.cursor_dup_write::().unwrap(); let mut tree: BTreeMap> = BTreeMap::new(); + let mut rev_changeset_walker = + storage_changesets_cursor.walk_back(None).unwrap(); while let Some((tid_address, entry)) = rev_changeset_walker.next().transpose().unwrap() { @@ -413,15 +396,18 @@ mod tests { .or_default() .insert(keccak256(entry.key), entry.value); } - for (key, val) in tree.into_iter() { - for (entry_key, entry_val) in val.into_iter() { - hash_cursor.seek_by_key_subkey(key, entry_key).unwrap(); - hash_cursor.delete_current().unwrap(); + for (hashed_address, storage) in tree.into_iter() { + for (hashed_slot, value) in storage.into_iter() { + let storage_entry = storage_cursor + .seek_by_key_subkey(hashed_address, hashed_slot) + .unwrap(); + if storage_entry.map(|v| v.key == hashed_slot).unwrap_or_default() { + storage_cursor.delete_current().unwrap(); + } - if entry_val != U256::ZERO { - let storage_entry = - StorageEntry { key: entry_key, value: entry_val }; - hash_cursor.append_dup(key, storage_entry).unwrap(); + if value != U256::ZERO { + let storage_entry = StorageEntry { key: hashed_slot, value }; + storage_cursor.upsert(hashed_address, storage_entry).unwrap(); } } } @@ -437,28 +423,18 @@ mod tests { break } - match account_before_tx.info { - Some(acc) => { - tx.put::(account_before_tx.address, acc) - .unwrap(); - tx.put::( - keccak256(account_before_tx.address), - acc, - ) - .unwrap(); - } - None => { - tx.delete::( - account_before_tx.address, - None, - ) - .unwrap(); - tx.delete::( - keccak256(account_before_tx.address), - None, - ) - .unwrap(); - } + if let Some(acc) = account_before_tx.info { + tx.put::( + keccak256(account_before_tx.address), + acc, + ) + .unwrap(); + } else { + tx.delete::( + keccak256(account_before_tx.address), + None, + ) + .unwrap(); } } Ok(()) @@ -467,48 +443,8 @@ mod tests { Ok(()) } - fn validate_unwind(&self, input: UnwindInput) -> Result<(), TestRunnerError> { - self.check_root(input.unwind_to) - } - } - - impl MerkleTestRunner { - fn state_root(&self) -> Result { - Ok(create_trie_loader(&self.tx.inner()) - .calculate_root() - .and_then(|e| e.root()) - .unwrap()) - } - - pub(crate) fn generate_initial_trie( - &self, - accounts: impl IntoIterator, - ) -> Result { - self.tx.insert_accounts_and_storages( - accounts.into_iter().map(|(addr, acc)| (addr, (acc, std::iter::empty()))), - )?; - - let mut tx = self.tx.inner(); - let root = create_trie_loader(&tx) - .calculate_root() - .and_then(|e| e.root()) - .expect("couldn't create initial trie"); - - tx.commit()?; - - Ok(root) - } - - fn check_root(&self, previous_stage_progress: u64) -> Result<(), TestRunnerError> { - if previous_stage_progress != 0 { - let block_root = - self.tx.inner().get_header(previous_stage_progress).unwrap().state_root; - let root = create_trie_loader(&self.tx().inner()) - .calculate_root() - .and_then(|e| e.root()) - .unwrap(); - assert_eq!(block_root, root); - } + fn validate_unwind(&self, _input: UnwindInput) -> Result<(), TestRunnerError> { + // The unwind is validated within the stage Ok(()) } } diff --git a/crates/stages/src/stages/mod.rs b/crates/stages/src/stages/mod.rs index b44e49309e..96f1ec5fec 100644 --- a/crates/stages/src/stages/mod.rs +++ b/crates/stages/src/stages/mod.rs @@ -14,7 +14,7 @@ mod headers; mod index_account_history; /// Index history of storage changes mod index_storage_history; -/// Intermediate hashes and creating merkle root +/// Stage for computing state root. mod merkle; /// The sender recovery stage. mod sender_recovery; diff --git a/crates/storage/db/src/tables/codecs/compact.rs b/crates/storage/db/src/tables/codecs/compact.rs index 196a3ec4b3..0b623981d6 100644 --- a/crates/storage/db/src/tables/codecs/compact.rs +++ b/crates/storage/db/src/tables/codecs/compact.rs @@ -4,10 +4,7 @@ use crate::{ Error, }; use reth_codecs::{main_codec, Compact}; -use reth_primitives::{ - trie::{StoredNibbles, StoredNibblesSubKey}, - *, -}; +use reth_primitives::{trie::*, *}; /// Implements compression for Compact type. macro_rules! impl_compression_for_compact { @@ -44,6 +41,7 @@ impl_compression_for_compact!( TxType, StorageEntry, StoredNibbles, + BranchNodeCompact, StoredNibblesSubKey, StorageTrieEntry, StoredBlockBodyIndices, diff --git a/crates/storage/db/src/tables/mod.rs b/crates/storage/db/src/tables/mod.rs index ead0861677..00105af6bb 100644 --- a/crates/storage/db/src/tables/mod.rs +++ b/crates/storage/db/src/tables/mod.rs @@ -33,8 +33,9 @@ use crate::{ }, }; use reth_primitives::{ + trie::{BranchNodeCompact, StorageTrieEntry, StoredNibbles, StoredNibblesSubKey}, Account, Address, BlockHash, BlockNumber, Bytecode, Header, IntegerList, Receipt, StorageEntry, - StorageTrieEntry, TransactionSigned, TransitionId, TxHash, TxNumber, H256, + TransactionSigned, TransitionId, TxHash, TxNumber, H256, }; /// Enum for the types of tables present in libmdbx. @@ -280,12 +281,12 @@ dupsort!( table!( /// Stores the current state's Merkle Patricia Tree. - ( AccountsTrie ) H256 | Vec + ( AccountsTrie ) StoredNibbles | BranchNodeCompact ); dupsort!( - /// Stores the Merkle Patricia Trees of each [`Account`]'s storage. - ( StoragesTrie ) H256 | [H256] StorageTrieEntry + /// From HashedAddress => NibblesSubKey => Intermediate value + ( StoragesTrie ) H256 | [StoredNibblesSubKey] StorageTrieEntry ); table!( diff --git a/crates/storage/provider/Cargo.toml b/crates/storage/provider/Cargo.toml index bf95b74f65..65bfc44428 100644 --- a/crates/storage/provider/Cargo.toml +++ b/crates/storage/provider/Cargo.toml @@ -13,17 +13,11 @@ reth-primitives = { path = "../../primitives" } reth-interfaces = { path = "../../interfaces" } reth-revm-primitives = { path = "../../revm/revm-primitives" } reth-db = { path = "../db" } -reth-codecs = { path = "../codecs" } -reth-tracing = { path = "../../tracing" } -reth-rlp = { path = "../../rlp" } +reth-trie = { path = "../../trie" } # async tokio = { version = "1.21", features = ["sync", "macros", "rt-multi-thread"] } -# trie -cita_trie = "4.0.0" -hasher = "0.1.4" - # tracing tracing = "0.1" @@ -31,18 +25,16 @@ tracing = "0.1" thiserror = "1.0.37" auto_impl = "1.0" itertools = "0.10" -parking_lot = "0.12" + +# test-utils +reth-rlp = { path = "../../rlp", optional = true } +parking_lot = { version = "0.12", optional = true } [dev-dependencies] reth-db = { path = "../db", features = ["test-utils"] } reth-primitives = { path = "../../primitives", features = ["arbitrary", "test-utils"] } parking_lot = "0.12" -proptest = { version = "1.0" } -assert_matches = "1.5" - -# trie -triehash = "0.8" [features] bench = [] -test-utils = [] +test-utils = ["reth-rlp", "parking_lot"] diff --git a/crates/storage/provider/src/lib.rs b/crates/storage/provider/src/lib.rs index e8fe0ba1b3..3a76d41e8e 100644 --- a/crates/storage/provider/src/lib.rs +++ b/crates/storage/provider/src/lib.rs @@ -25,9 +25,6 @@ pub use providers::{ LatestStateProviderRef, ShareableDatabase, }; -/// Helper type for loading Merkle Patricia Trees from the database -pub mod trie; - /// Execution result pub mod post_state; pub use post_state::PostState; diff --git a/crates/storage/provider/src/providers/state/latest.rs b/crates/storage/provider/src/providers/state/latest.rs index 8303a48b09..6c99f59ab7 100644 --- a/crates/storage/provider/src/providers/state/latest.rs +++ b/crates/storage/provider/src/providers/state/latest.rs @@ -1,6 +1,6 @@ use crate::{ - providers::state::macros::delegate_provider_impls, trie::DBTrieLoader, AccountProvider, - BlockHashProvider, StateProvider, + providers::state::macros::delegate_provider_impls, AccountProvider, BlockHashProvider, + StateProvider, }; use reth_db::{ cursor::{DbCursorRO, DbDupCursorRO}, @@ -10,7 +10,6 @@ use reth_db::{ use reth_interfaces::{provider::ProviderError, Result}; use reth_primitives::{ keccak256, Account, Address, BlockNumber, Bytecode, Bytes, StorageKey, StorageValue, H256, - KECCAK_EMPTY, }; use std::marker::PhantomData; @@ -76,11 +75,10 @@ impl<'a, 'b, TX: DbTx<'a>> StateProvider for LatestStateProviderRef<'a, 'b, TX> fn proof( &self, address: Address, - keys: &[H256], + _keys: &[H256], ) -> Result<(Vec, H256, Vec>)> { - let hashed_address = keccak256(address); - let loader = DBTrieLoader::new(self.db); - let root = self + let _hashed_address = keccak256(address); + let _root = self .db .cursor_read::()? .last()? @@ -88,25 +86,7 @@ impl<'a, 'b, TX: DbTx<'a>> StateProvider for LatestStateProviderRef<'a, 'b, TX> .1 .state_root; - let (account_proof, storage_root) = loader - .generate_account_proof(root, hashed_address) - .map_err(|_| ProviderError::StateTrie)?; - let account_proof = account_proof.into_iter().map(Bytes::from).collect(); - - let storage_proof = if storage_root == KECCAK_EMPTY { - // if there isn't storage, we return empty storage proofs - (0..keys.len()).map(|_| Vec::new()).collect() - } else { - let hashed_keys: Vec = keys.iter().map(keccak256).collect(); - loader - .generate_storage_proofs(storage_root, hashed_address, &hashed_keys) - .map_err(|_| ProviderError::StateTrie)? - .into_iter() - .map(|v| v.into_iter().map(Bytes::from).collect()) - .collect() - }; - - Ok((account_proof, storage_root, storage_proof)) + unimplemented!() } } diff --git a/crates/storage/provider/src/test_utils/blocks.rs b/crates/storage/provider/src/test_utils/blocks.rs index de1af0f66c..7e4ec1c476 100644 --- a/crates/storage/provider/src/test_utils/blocks.rs +++ b/crates/storage/provider/src/test_utils/blocks.rs @@ -3,8 +3,8 @@ use crate::{post_state::PostState, Transaction}; use reth_db::{database::Database, models::StoredBlockBodyIndices, tables}; use reth_primitives::{ - hex_literal::hex, proofs::EMPTY_ROOT, Account, Header, SealedBlock, SealedBlockWithSenders, - Withdrawal, H160, H256, U256, + hex_literal::hex, Account, Header, SealedBlock, SealedBlockWithSenders, Withdrawal, H160, H256, + U256, }; use reth_rlp::Decodable; use std::collections::BTreeMap; @@ -39,7 +39,7 @@ pub fn assert_genesis_block(tx: &Transaction<'_, DB>, g: SealedBlo assert_eq!(tx.table::().unwrap(), vec![]); assert_eq!(tx.table::().unwrap(), vec![]); assert_eq!(tx.table::().unwrap(), vec![]); - assert_eq!(tx.table::().unwrap(), vec![(EMPTY_ROOT, vec![0x80])]); + assert_eq!(tx.table::().unwrap(), vec![]); assert_eq!(tx.table::().unwrap(), vec![]); assert_eq!(tx.table::().unwrap(), vec![]); // SyncStage is not updated in tests diff --git a/crates/storage/provider/src/transaction.rs b/crates/storage/provider/src/transaction.rs index 16167f94bc..4c4b1b9826 100644 --- a/crates/storage/provider/src/transaction.rs +++ b/crates/storage/provider/src/transaction.rs @@ -1,7 +1,6 @@ use crate::{ insert_canonical_block, post_state::{Change, PostState, StorageChangeset}, - trie::{DBTrieLoader, TrieError}, }; use itertools::{izip, Itertools}; use reth_db::{ @@ -24,6 +23,7 @@ use reth_primitives::{ Header, SealedBlock, SealedBlockWithSenders, StorageEntry, TransactionSignedEcRecovered, TransitionId, TxNumber, H256, U256, }; +use reth_trie::{StateRoot, StateRootError}; use std::{ collections::{btree_map::Entry, BTreeMap, BTreeSet}, fmt::Debug, @@ -541,7 +541,6 @@ where self.insert_block(block)?; } self.insert_hashes( - fork_block_number, first_transition_id, first_transition_id + num_transitions, new_tip_number, @@ -598,7 +597,6 @@ where /// The resulting state root is compared with `expected_state_root`. pub fn insert_hashes( &mut self, - fork_block_number: BlockNumber, from_transition_id: TransitionId, to_transition_id: TransitionId, current_block_number: BlockNumber, @@ -623,14 +621,14 @@ where // merkle tree { - let current_root = self.get_header(fork_block_number)?.state_root; - let mut loader = DBTrieLoader::new(self.deref_mut()); - let root = loader - .update_root(current_root, from_transition_id..to_transition_id) - .and_then(|e| e.root())?; - if root != expected_state_root { + let state_root = StateRoot::incremental_root( + self.deref_mut(), + from_transition_id..to_transition_id, + None, + )?; + if state_root != expected_state_root { return Err(TransactionError::StateTrieRootMismatch { - got: root, + got: state_root, expected: expected_state_root, block_number: current_block_number, block_hash: current_block_hash, @@ -1111,15 +1109,7 @@ where self.unwind_storage_history_indices(transition_storage_range)?; // merkle tree - let new_state_root; - { - let (tip_number, _) = - self.cursor_read::()?.last()?.unwrap_or_default(); - let current_root = self.get_header(tip_number)?.state_root; - let mut loader = DBTrieLoader::new(self.deref()); - new_state_root = - loader.update_root(current_root, transition_range).and_then(|e| e.root())?; - } + let new_state_root = StateRoot::incremental_root(self.deref(), transition_range, None)?; // state root should be always correct as we are reverting state. // but for sake of double verification we will check it again. if new_state_root != parent_state_root { @@ -1537,14 +1527,14 @@ fn unwind_storage_history_shards( #[derive(Debug, thiserror::Error)] pub enum TransactionError { /// The transaction encountered a database error. - #[error("Database error: {0}")] + #[error(transparent)] Database(#[from] DbError), /// The transaction encountered a database integrity error. - #[error("A database integrity error occurred: {0}")] + #[error(transparent)] DatabaseIntegrity(#[from] ProviderError), - /// The transaction encountered merkle trie error. - #[error("Merkle trie calculation error: {0}")] - MerkleTrie(#[from] TrieError), + /// The trie error. + #[error(transparent)] + TrieError(#[from] StateRootError), /// Root mismatch #[error("Merkle trie root mismatch on block: #{block_number:?} {block_hash:?}. got: {got:?} expected:{expected:?}")] StateTrieRootMismatch { @@ -1565,8 +1555,8 @@ mod test { insert_canonical_block, test_utils::blocks::*, ShareableDatabase, Transaction, TransactionsProvider, }; - use reth_db::{mdbx::test_utils::create_test_rw_db, tables, transaction::DbTxMut}; - use reth_primitives::{proofs::EMPTY_ROOT, ChainSpecBuilder, TransitionId, MAINNET}; + use reth_db::mdbx::test_utils::create_test_rw_db; + use reth_primitives::{ChainSpecBuilder, TransitionId, MAINNET}; use std::{ops::DerefMut, sync::Arc}; #[test] @@ -1587,14 +1577,11 @@ mod test { let (block2, exec_res2) = data.blocks[1].clone(); insert_canonical_block(tx.deref_mut(), data.genesis.clone(), None, false).unwrap(); - - tx.put::(EMPTY_ROOT, vec![0x80]).unwrap(); assert_genesis_block(&tx, data.genesis); exec_res1.clone().write_to_db(tx.deref_mut(), 0).unwrap(); tx.insert_block(block1.clone()).unwrap(); tx.insert_hashes( - genesis.number, 0, exec_res1.transitions_count() as TransitionId, block1.number, @@ -1615,7 +1602,6 @@ mod test { exec_res1.clone().write_to_db(tx.deref_mut(), 0).unwrap(); tx.insert_block(block1.clone()).unwrap(); tx.insert_hashes( - genesis.number, 0, exec_res1.transitions_count() as TransitionId, block1.number, @@ -1630,7 +1616,6 @@ mod test { .unwrap(); tx.insert_block(block2.clone()).unwrap(); tx.insert_hashes( - block1.number, exec_res1.transitions_count() as TransitionId, (exec_res1.transitions_count() + exec_res2.transitions_count()) as TransitionId, 2, @@ -1698,8 +1683,6 @@ mod test { let (block2, exec_res2) = data.blocks[1].clone(); insert_canonical_block(tx.deref_mut(), data.genesis.clone(), None, false).unwrap(); - - tx.put::(EMPTY_ROOT, vec![0x80]).unwrap(); assert_genesis_block(&tx, data.genesis); tx.append_blocks_with_post_state(vec![block1.clone()], exec_res1.clone()).unwrap(); diff --git a/crates/storage/provider/src/trie/mod.rs b/crates/storage/provider/src/trie/mod.rs deleted file mode 100644 index bb5e6bc7ee..0000000000 --- a/crates/storage/provider/src/trie/mod.rs +++ /dev/null @@ -1,1269 +0,0 @@ -use cita_trie::{PatriciaTrie, Trie}; -use hasher::HasherKeccak; -use parking_lot::Mutex; -use reth_codecs::Compact; -use reth_db::{ - cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO}, - models::{AccountBeforeTx, TransitionIdAddress}, - tables, - transaction::{DbTx, DbTxMut, DbTxMutGAT}, -}; -use reth_primitives::{ - keccak256, proofs::EMPTY_ROOT, Account, Address, ProofCheckpoint, StorageEntry, - StorageTrieEntry, TransitionId, H256, KECCAK_EMPTY, U256, -}; -use reth_rlp::{ - encode_fixed_size, Decodable, DecodeError, Encodable, RlpDecodable, RlpEncodable, - EMPTY_STRING_CODE, -}; -use reth_tracing::tracing::*; -use std::{ - collections::{BTreeMap, BTreeSet}, - marker::PhantomData, - ops::Range, - sync::Arc, -}; - -/// Merkle Trie error types -#[allow(missing_docs)] -#[derive(Debug, thiserror::Error)] -pub enum TrieError { - /// Error returned by the underlying implementation. - #[error("Some error occurred: {0}")] - InternalError(#[from] cita_trie::TrieError), - /// The database doesn't contain the root of the trie. - #[error("The root node wasn't found in the DB")] - MissingAccountRoot(H256), - #[error("The storage root node wasn't found in the DB")] - MissingStorageRoot(H256), - /// Error returned by the database. - #[error("{0:?}")] - DatabaseError(#[from] reth_db::Error), - /// Error when encoding/decoding a value. - #[error("{0:?}")] - DecodeError(#[from] DecodeError), - #[error("Trie requires committing a checkpoint.")] - UnexpectedCheckpoint, -} - -type AccountsTrieCursor<'tx, TX> = - Arc>::CursorMut>>; - -/// Database wrapper implementing HashDB trait, with a read-write transaction. -pub struct HashDatabaseMut<'tx, TX: DbTxMutGAT<'tx>> { - accounts_trie_cursor: AccountsTrieCursor<'tx, TX>, -} - -impl<'tx, 'db, TX> cita_trie::DB for HashDatabaseMut<'tx, TX> -where - TX: DbTxMut<'db> + DbTx<'db> + Send + Sync, -{ - type Error = TrieError; - - fn get(&self, key: &[u8]) -> Result>, Self::Error> { - Ok(self.accounts_trie_cursor.lock().seek_exact(H256::from_slice(key))?.map(|(_, v)| v)) - } - - fn contains(&self, key: &[u8]) -> Result { - Ok(::get(self, key)?.is_some()) - } - - fn insert(&self, _key: Vec, _value: Vec) -> Result<(), Self::Error> { - unreachable!("Use batch instead."); - } - - // Insert a batch of data into the cache. - fn insert_batch(&self, keys: Vec>, values: Vec>) -> Result<(), Self::Error> { - let mut cursor = self.accounts_trie_cursor.lock(); - for (key, value) in keys.into_iter().zip(values.into_iter()) { - cursor.upsert(H256::from_slice(key.as_slice()), value)?; - } - Ok(()) - } - - fn remove_batch(&self, keys: &[Vec]) -> Result<(), Self::Error> { - let mut cursor = self.accounts_trie_cursor.lock(); - for key in keys { - if cursor.seek_exact(H256::from_slice(key.as_slice()))?.is_some() { - cursor.delete_current()?; - } - } - Ok(()) - } - - fn remove(&self, _key: &[u8]) -> Result<(), Self::Error> { - unreachable!("Use batch instead."); - } - - fn flush(&self) -> Result<(), Self::Error> { - Ok(()) - } -} - -impl<'tx, 'db, TX> HashDatabaseMut<'tx, TX> -where - TX: DbTxMut<'db> + DbTx<'db> + Send + Sync, -{ - /// Instantiates a new Database for the accounts trie, with an empty root - pub fn new(tx: &'tx TX) -> Result { - let mut accounts_trie_cursor = tx.cursor_write::()?; - - let root = EMPTY_ROOT; - if accounts_trie_cursor.seek_exact(root)?.is_none() { - accounts_trie_cursor.upsert(root, [EMPTY_STRING_CODE].to_vec())?; - } - - Ok(Self { accounts_trie_cursor: Arc::new(Mutex::new(accounts_trie_cursor)) }) - } - - /// Instantiates a new Database for the accounts trie, with an existing root - pub fn from_root(tx: &'tx TX, root: H256) -> Result { - let mut accounts_trie_cursor = tx.cursor_write::()?; - - if root == EMPTY_ROOT { - return Self::new(tx) - } - accounts_trie_cursor.seek_exact(root)?.ok_or(TrieError::MissingAccountRoot(root))?; - - Ok(Self { accounts_trie_cursor: Arc::new(Mutex::new(accounts_trie_cursor)) }) - } -} - -type StoragesTrieCursor<'tx, TX> = - Arc>::DupCursorMut>>; - -/// Database wrapper implementing HashDB trait, with a read-write transaction. -pub struct DupHashDatabaseMut<'tx, TX: DbTxMutGAT<'tx>> { - storages_trie_cursor: StoragesTrieCursor<'tx, TX>, - key: H256, -} - -impl<'tx, 'db, TX> cita_trie::DB for DupHashDatabaseMut<'tx, TX> -where - TX: DbTxMut<'db> + DbTx<'db> + Send + Sync, -{ - type Error = TrieError; - - fn get(&self, key: &[u8]) -> Result>, Self::Error> { - let subkey = H256::from_slice(key); - Ok(self - .storages_trie_cursor - .lock() - .seek_by_key_subkey(self.key, subkey)? - .filter(|entry| entry.hash == subkey) - .map(|entry| entry.node)) - } - - fn contains(&self, key: &[u8]) -> Result { - Ok(::get(self, key)?.is_some()) - } - - fn insert(&self, _key: Vec, _value: Vec) -> Result<(), Self::Error> { - unreachable!("Use batch instead."); - } - - /// Insert a batch of data into the cache. - fn insert_batch(&self, keys: Vec>, values: Vec>) -> Result<(), Self::Error> { - let mut cursor = self.storages_trie_cursor.lock(); - for (key, node) in keys.into_iter().zip(values.into_iter()) { - let hash = H256::from_slice(key.as_slice()); - if cursor.seek_by_key_subkey(self.key, hash)?.filter(|e| e.hash == hash).is_some() { - cursor.delete_current()?; - } - cursor.upsert(self.key, StorageTrieEntry { hash, node })?; - } - Ok(()) - } - - fn remove_batch(&self, keys: &[Vec]) -> Result<(), Self::Error> { - let mut cursor = self.storages_trie_cursor.lock(); - for key in keys { - let hash = H256::from_slice(key.as_slice()); - if cursor.seek_by_key_subkey(self.key, hash)?.filter(|e| e.hash == hash).is_some() { - cursor.delete_current()?; - } - } - Ok(()) - } - - fn remove(&self, _key: &[u8]) -> Result<(), Self::Error> { - unreachable!("Use batch instead."); - } - - fn flush(&self) -> Result<(), Self::Error> { - Ok(()) - } -} - -impl<'tx, 'db, TX> DupHashDatabaseMut<'tx, TX> -where - TX: DbTxMut<'db> + DbTx<'db> + Send + Sync, -{ - /// Instantiates a new Database for the storage trie, with an empty root - pub fn new( - storages_trie_cursor: StoragesTrieCursor<'tx, TX>, - key: H256, - ) -> Result { - let root = EMPTY_ROOT; - { - let mut cursor = storages_trie_cursor.lock(); - if cursor.seek_by_key_subkey(key, root)?.filter(|entry| entry.hash == root).is_none() { - cursor.upsert( - key, - StorageTrieEntry { hash: root, node: [EMPTY_STRING_CODE].to_vec() }, - )?; - } - } - Ok(Self { storages_trie_cursor, key }) - } - - /// Instantiates a new Database for the storage trie, with an existing root - pub fn from_root( - storages_trie_cursor: StoragesTrieCursor<'tx, TX>, - key: H256, - root: H256, - ) -> Result { - if root == EMPTY_ROOT { - return Self::new(storages_trie_cursor, key) - } - storages_trie_cursor - .lock() - .seek_by_key_subkey(key, root)? - .filter(|entry| entry.hash == root) - .ok_or(TrieError::MissingStorageRoot(root))?; - Ok(Self { storages_trie_cursor, key }) - } -} - -/// Database wrapper implementing HashDB trait, with a read-only transaction. -pub struct HashDatabase<'tx, 'itx, TX: DbTx<'itx>> { - tx: &'tx TX, - _p: PhantomData<&'itx ()>, // to suppress "unused" lifetime 'itx -} - -impl<'tx, 'itx, TX: DbTx<'itx>> HashDatabase<'tx, 'itx, TX> { - /// Creates a new Hash database with the given transaction - pub fn new(tx: &'tx TX) -> Self { - Self { tx, _p: Default::default() } - } -} - -impl<'tx, 'itx, TX> cita_trie::DB for HashDatabase<'tx, 'itx, TX> -where - TX: DbTx<'itx>, -{ - type Error = TrieError; - - fn get(&self, key: &[u8]) -> Result>, Self::Error> { - Ok(self.tx.get::(H256::from_slice(key))?) - } - - fn contains(&self, key: &[u8]) -> Result { - Ok(::get(self, key)?.is_some()) - } - - fn insert(&self, _key: Vec, _value: Vec) -> Result<(), Self::Error> { - // this could be avoided if cita_trie::DB was split into two traits - // with read and write operations respectively - unimplemented!("insert isn't valid for read-only transaction"); - } - - fn remove(&self, _key: &[u8]) -> Result<(), Self::Error> { - unimplemented!("remove isn't valid for read-only transaction"); - } - - fn flush(&self) -> Result<(), Self::Error> { - Ok(()) - } -} - -impl<'tx, 'itx, TX: DbTx<'itx>> HashDatabase<'tx, 'itx, TX> { - /// Instantiates a new Database for the accounts trie, with an existing root - fn from_root(tx: &'tx TX, root: H256) -> Result { - tx.get::(root)?.ok_or(TrieError::MissingAccountRoot(root))?; - Ok(Self { tx, _p: Default::default() }) - } -} - -/// Database wrapper implementing HashDB trait, with a read-only transaction. -pub struct DupHashDatabase<'tx, 'itx, TX: DbTx<'itx>> { - tx: &'tx TX, - key: H256, - _p: PhantomData<&'itx ()>, // to suppress "unused" lifetime 'itx -} - -impl<'tx, 'itx, TX: DbTx<'itx>> DupHashDatabase<'tx, 'itx, TX> { - /// Creates a new DupHash database with the given transaction and key. - pub fn new(tx: &'tx TX, key: H256) -> Self { - Self { tx, key, _p: Default::default() } - } -} - -impl<'tx, 'itx, TX> cita_trie::DB for DupHashDatabase<'tx, 'itx, TX> -where - TX: DbTx<'itx>, -{ - type Error = TrieError; - - fn get(&self, key: &[u8]) -> Result>, Self::Error> { - let mut cursor = self.tx.cursor_dup_read::()?; - Ok(cursor.seek_by_key_subkey(self.key, H256::from_slice(key))?.map(|entry| entry.node)) - } - - fn contains(&self, key: &[u8]) -> Result { - Ok(::get(self, key)?.is_some()) - } - - fn insert(&self, _key: Vec, _value: Vec) -> Result<(), Self::Error> { - // Caching and bulk inserting shouldn't be needed, as the data is ordered - unimplemented!("insert isn't valid for read-only transaction"); - } - - fn remove(&self, _key: &[u8]) -> Result<(), Self::Error> { - unimplemented!("remove isn't valid for read-only transaction"); - } - - fn flush(&self) -> Result<(), Self::Error> { - Ok(()) - } -} - -impl<'tx, 'itx, TX: DbTx<'itx>> DupHashDatabase<'tx, 'itx, TX> { - /// Instantiates a new Database for the storage trie, with an existing root - fn from_root(tx: &'tx TX, key: H256, root: H256) -> Result { - tx.cursor_dup_read::()? - .seek_by_key_subkey(key, root)? - .ok_or(TrieError::MissingAccountRoot(root))?; - Ok(Self { tx, key, _p: Default::default() }) - } -} - -/// An Ethereum account, for RLP encoding traits deriving. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Default, RlpEncodable, RlpDecodable)] -pub struct EthAccount { - /// Account nonce. - nonce: u64, - /// Account balance. - balance: U256, - /// Account's storage root. - storage_root: H256, - /// Hash of the account's bytecode. - code_hash: H256, -} - -impl From for EthAccount { - fn from(acc: Account) -> Self { - EthAccount { - nonce: acc.nonce, - balance: acc.balance, - storage_root: EMPTY_ROOT, - code_hash: acc.bytecode_hash.unwrap_or(KECCAK_EMPTY), - } - } -} - -impl EthAccount { - /// Set storage root on account. - pub fn with_storage_root(mut self, storage_root: H256) -> Self { - self.storage_root = storage_root; - self - } - - /// Get account's storage root. - pub fn storage_root(&self) -> H256 { - self.storage_root - } -} - -/// A merkle proof of existence (or nonexistence) of a leaf value. Consists -/// of a the encoded nodes in the path from the root of the tree to the leaf. -pub type MerkleProof = Vec>; - -/// Struct for calculating the root of a merkle patricia tree, -/// while populating the database with intermediate hashes. -#[derive(Debug)] -pub struct DBTrieLoader<'tx, TX> { - /// The maximum number of keys to insert before committing. Both from `AccountsTrie` and - /// `StoragesTrie`. - pub commit_threshold: u64, - /// The current number of inserted keys from both `AccountsTrie` and `StoragesTrie`. - pub current: u64, - /// The transaction to use for inserting the trie nodes. - pub tx: &'tx TX, -} - -/// Status of the trie calculation. -#[derive(Debug, PartialEq, Copy, Clone)] -pub enum TrieProgress { - /// Trie has finished with the passed root. - Complete(H256), - /// Trie has hit its commit threshold. - InProgress(ProofCheckpoint), -} - -impl TrieProgress { - /// Consumes the root from its `Complete` variant. If that's not possible, throw - /// `TrieError::UnexpectedCheckpoint`. - pub fn root(self) -> Result { - match self { - Self::Complete(root) => Ok(root), - _ => Err(TrieError::UnexpectedCheckpoint), - } - } -} - -impl<'tx, TX> DBTrieLoader<'tx, TX> { - /// Create new instance of trie loader. - pub fn new(tx: &'tx TX) -> Self { - Self { tx, commit_threshold: 500_000, current: 0 } - } -} - -// Read-write impls -impl<'tx, 'db, TX> DBTrieLoader<'tx, TX> -where - TX: DbTxMut<'db> + DbTx<'db> + Send + Sync, -{ - /// Calculates the root of the state trie, saving intermediate hashes in the database. - pub fn calculate_root(&mut self) -> Result { - let mut checkpoint = self.get_checkpoint()?; - - if checkpoint.hashed_address.is_none() { - self.tx.clear::()?; - self.tx.clear::()?; - } - let previous_root = checkpoint.account_root.unwrap_or(EMPTY_ROOT); - - let hasher = Arc::new(HasherKeccak::new()); - let mut trie = if let Some(root) = checkpoint.account_root { - PatriciaTrie::from( - Arc::new(HashDatabaseMut::from_root(self.tx, root)?), - hasher, - root.as_bytes(), - )? - } else { - PatriciaTrie::new(Arc::new(HashDatabaseMut::new(self.tx)?), hasher) - }; - - let mut accounts_cursor = self.tx.cursor_read::()?; - let storage_trie_cursor = - Arc::new(Mutex::new(self.tx.cursor_dup_write::()?)); - let mut walker = accounts_cursor.walk(checkpoint.hashed_address.take())?; - - while let Some((hashed_address, account)) = walker.next().transpose()? { - match self.calculate_storage_root( - hashed_address, - storage_trie_cursor.clone(), - checkpoint.storage_key.take(), - checkpoint.storage_root.take(), - )? { - TrieProgress::Complete(root) => { - let value = EthAccount::from(account).with_storage_root(root); - - let mut out = Vec::new(); - Encodable::encode(&value, &mut out); - trie.insert(hashed_address.as_bytes().to_vec(), out)?; - - if self.has_hit_threshold() { - return self.save_account_checkpoint( - ProofCheckpoint::default(), - self.replace_account_root(&mut trie, previous_root)?, - hashed_address, - ) - } - } - TrieProgress::InProgress(checkpoint) => { - return self.save_account_checkpoint( - checkpoint, - self.replace_account_root(&mut trie, previous_root)?, - hashed_address, - ) - } - } - } - - // Reset inner stage progress - self.save_checkpoint(ProofCheckpoint::default())?; - - Ok(TrieProgress::Complete(self.replace_account_root(&mut trie, previous_root)?)) - } - - fn calculate_storage_root( - &mut self, - address: H256, - storage_trie_cursor: StoragesTrieCursor<'tx, TX>, - next_storage: Option, - previous_root: Option, - ) -> Result { - let mut storage_cursor = self.tx.cursor_dup_read::()?; - - let hasher = Arc::new(HasherKeccak::new()); - let (mut current_entry, mut trie) = if let Some(entry) = next_storage { - ( - storage_cursor.seek_by_key_subkey(address, entry)?.filter(|e| e.key == entry), - PatriciaTrie::from( - Arc::new(DupHashDatabaseMut::from_root( - storage_trie_cursor, - address, - previous_root.expect("is some"), - )?), - hasher, - previous_root.expect("is some").as_bytes(), - )?, - ) - } else { - ( - storage_cursor.seek_by_key_subkey(address, H256::zero())?, - PatriciaTrie::new( - Arc::new(DupHashDatabaseMut::new(storage_trie_cursor, address)?), - hasher, - ), - ) - }; - - let previous_root = previous_root.unwrap_or(EMPTY_ROOT); - - while let Some(StorageEntry { key: storage_key, value }) = current_entry { - let out = encode_fixed_size(&value).to_vec(); - trie.insert(storage_key.to_vec(), out)?; - // Should be able to use walk_dup, but any call to next() causes an assert fail in - // mdbx.c - current_entry = storage_cursor.next_dup()?.map(|(_, v)| v); - let threshold = self.has_hit_threshold(); - if let Some(current_entry) = current_entry { - if threshold { - return Ok(TrieProgress::InProgress(ProofCheckpoint { - storage_root: Some(self.replace_storage_root( - trie, - address, - previous_root, - )?), - storage_key: Some(current_entry.key), - ..Default::default() - })) - } - } - } - - Ok(TrieProgress::Complete(self.replace_storage_root(trie, address, previous_root)?)) - } - - /// Calculates the root of the state trie by updating an existing trie. - pub fn update_root( - &mut self, - mut previous_root: H256, - tid_range: Range, - ) -> Result { - let mut checkpoint = self.get_checkpoint()?; - - if let Some(account_root) = checkpoint.account_root.take() { - previous_root = account_root; - } - - let next_acc = checkpoint.hashed_address.take(); - let changed_accounts = self - .gather_changes(tid_range)? - .into_iter() - .skip_while(|(addr, _)| next_acc.is_some() && next_acc.expect("is some") != *addr); - - let mut trie = PatriciaTrie::from( - Arc::new(HashDatabaseMut::from_root(self.tx, previous_root)?), - Arc::new(HasherKeccak::new()), - previous_root.as_bytes(), - )?; - - let mut accounts_cursor = self.tx.cursor_read::()?; - let storage_trie_cursor = - Arc::new(Mutex::new(self.tx.cursor_dup_write::()?)); - - for (hashed_address, changed_storages) in changed_accounts { - let res = if let Some(account) = trie.get(hashed_address.as_slice())? { - trie.remove(hashed_address.as_bytes())?; - - let storage_root = EthAccount::decode(&mut account.as_slice())?.storage_root; - self.update_storage_root( - checkpoint.storage_root.take().unwrap_or(storage_root), - hashed_address, - storage_trie_cursor.clone(), - changed_storages, - checkpoint.storage_key.take(), - )? - } else { - self.calculate_storage_root( - hashed_address, - storage_trie_cursor.clone(), - checkpoint.storage_key.take(), - checkpoint.storage_root.take(), - )? - }; - - let storage_root = match res { - TrieProgress::Complete(root) => root, - TrieProgress::InProgress(checkpoint) => { - return self.save_account_checkpoint( - checkpoint, - self.replace_account_root(&mut trie, previous_root)?, - hashed_address, - ) - } - }; - - if let Some((_, account)) = accounts_cursor.seek_exact(hashed_address)? { - let value = EthAccount::from(account).with_storage_root(storage_root); - - let mut out = Vec::new(); - Encodable::encode(&value, &mut out); - - trie.insert(hashed_address.as_bytes().to_vec(), out)?; - - if self.has_hit_threshold() { - return self.save_account_checkpoint( - ProofCheckpoint::default(), - self.replace_account_root(&mut trie, previous_root)?, - hashed_address, - ) - } - } - } - - // Reset inner stage progress - self.save_checkpoint(ProofCheckpoint::default())?; - - Ok(TrieProgress::Complete(self.replace_account_root(&mut trie, previous_root)?)) - } - - /// Update the account's storage root - fn update_storage_root( - &mut self, - previous_root: H256, - address: H256, - storage_trie_cursor: StoragesTrieCursor<'tx, TX>, - changed_storages: BTreeSet, - next_storage: Option, - ) -> Result { - let mut hashed_storage_cursor = self.tx.cursor_dup_read::()?; - let mut trie = PatriciaTrie::new( - Arc::new(DupHashDatabaseMut::from_root(storage_trie_cursor, address, previous_root)?), - Arc::new(HasherKeccak::new()), - ); - - let changed_storages = changed_storages - .into_iter() - .skip_while(|k| next_storage.is_some() && *k == next_storage.expect("is some")); - - for key in changed_storages { - if let Some(StorageEntry { value, .. }) = - hashed_storage_cursor.seek_by_key_subkey(address, key)?.filter(|e| e.key == key) - { - let out = encode_fixed_size(&value).to_vec(); - trie.insert(key.as_bytes().to_vec(), out)?; - if self.has_hit_threshold() { - return Ok(TrieProgress::InProgress(ProofCheckpoint { - storage_root: Some(self.replace_storage_root( - trie, - address, - previous_root, - )?), - storage_key: Some(key), - ..Default::default() - })) - } - } else { - trie.remove(key.as_bytes())?; - } - } - - Ok(TrieProgress::Complete(self.replace_storage_root(trie, address, previous_root)?)) - } - - fn gather_changes( - &self, - tid_range: Range, - ) -> Result>, TrieError> { - let mut account_cursor = self.tx.cursor_read::()?; - - let mut account_changes: BTreeMap> = BTreeMap::new(); - - let mut walker = account_cursor.walk_range(tid_range.clone())?; - - while let Some((_, AccountBeforeTx { address, .. })) = walker.next().transpose()? { - account_changes.insert(address, Default::default()); - } - - let mut storage_cursor = self.tx.cursor_dup_read::()?; - - let start = TransitionIdAddress((tid_range.start, Address::zero())); - let end = TransitionIdAddress((tid_range.end, Address::zero())); - let mut walker = storage_cursor.walk_range(start..end)?; - - while let Some((TransitionIdAddress((_, address)), StorageEntry { key, .. })) = - walker.next().transpose()? - { - account_changes.entry(address).or_default().insert(key); - } - - let hashed_changes = account_changes - .into_iter() - .map(|(address, storage)| { - (keccak256(address), storage.into_iter().map(keccak256).collect()) - }) - .collect(); - - Ok(hashed_changes) - } - - fn save_account_checkpoint( - &mut self, - mut checkpoint: ProofCheckpoint, - root: H256, - hashed_address: H256, - ) -> Result { - checkpoint.account_root = Some(root); - checkpoint.hashed_address = Some(hashed_address); - - debug!(target: "sync::stages::merkle::exec", account = ?hashed_address, storage = ?checkpoint.storage_key, "Saving inner trie checkpoint"); - - self.save_checkpoint(checkpoint)?; - - Ok(TrieProgress::InProgress(checkpoint)) - } - - fn has_hit_threshold(&mut self) -> bool { - self.current += 1; - self.current >= self.commit_threshold - } - - /// Saves the trie progress - pub fn save_checkpoint(&mut self, checkpoint: ProofCheckpoint) -> Result<(), TrieError> { - let mut buf = vec![]; - checkpoint.to_compact(&mut buf); - - // It allows unwind (which commits), to reuse this instance. - self.current = 0; - - Ok(self.tx.put::("TrieLoader".into(), buf)?) - } - - /// Gets the trie progress - pub fn get_checkpoint(&self) -> Result { - let buf = - self.tx.get::("TrieLoader".into())?.unwrap_or_default(); - - if buf.is_empty() { - return Ok(ProofCheckpoint::default()) - } - - let (checkpoint, _) = ProofCheckpoint::from_compact(&buf, buf.len()); - - if checkpoint.account_root.is_some() { - debug!(target: "sync::stages::merkle::exec", checkpoint = ?checkpoint, "Continuing inner trie checkpoint"); - } - - Ok(checkpoint) - } - - /// Finds the most recent account trie root and removes the previous one if applicable. - fn replace_account_root( - &self, - trie: &mut PatriciaTrie, HasherKeccak>, - previous_root: H256, - ) -> Result { - let new_root = H256::from_slice(trie.root()?.as_slice()); - - if new_root != previous_root { - let mut cursor = self.tx.cursor_write::()?; - if cursor.seek_exact(previous_root)?.is_some() { - cursor.delete_current()?; - } - } - - Ok(new_root) - } - - /// Finds the most recent storage trie root and removes the previous one if applicable. - fn replace_storage_root( - &self, - mut trie: PatriciaTrie, HasherKeccak>, - address: H256, - previous_root: H256, - ) -> Result { - let new_root = H256::from_slice(trie.root()?.as_slice()); - - if new_root != previous_root { - let mut trie_cursor = self.tx.cursor_dup_write::()?; - - if trie_cursor - .seek_by_key_subkey(address, previous_root)? - .filter(|entry| entry.hash == previous_root) - .is_some() - { - trie_cursor.delete_current()?; - } - } - - if new_root == EMPTY_ROOT { - self.tx.delete::(address, None)?; - } - - Ok(new_root) - } -} - -// Read-only impls -impl<'tx, 'db, TX> DBTrieLoader<'tx, TX> -where - TX: DbTx<'db> + Send + Sync, -{ - /// Returns a Merkle proof of the given account, plus its storage root hash. - pub fn generate_account_proof( - &self, - root: H256, - address: H256, - ) -> Result<(MerkleProof, H256), TrieError> { - let db = Arc::new(HashDatabase::from_root(self.tx, root)?); - let hasher = Arc::new(HasherKeccak::new()); - - let trie = PatriciaTrie::from(Arc::clone(&db), Arc::clone(&hasher), root.as_bytes())?; - let proof = trie.get_proof(address.as_bytes())?; - - let Some(account) = trie.get(address.as_slice())? else { return Ok((proof, KECCAK_EMPTY)) }; - - let storage_root = EthAccount::decode(&mut account.as_slice())?.storage_root; - - Ok((proof, storage_root)) - } - - /// Returns a Merkle proof of the given storage keys, starting at the given root hash. - pub fn generate_storage_proofs( - &self, - storage_root: H256, - address: H256, - keys: &[H256], - ) -> Result, TrieError> { - let db = Arc::new(DupHashDatabase::from_root(self.tx, address, storage_root)?); - let hasher = Arc::new(HasherKeccak::new()); - - let trie = - PatriciaTrie::from(Arc::clone(&db), Arc::clone(&hasher), storage_root.as_bytes())?; - - let proof = - keys.iter().map(|key| trie.get_proof(key.as_bytes())).collect::, _>>()?; - - Ok(proof) - } -} - -#[cfg(test)] -mod tests { - use crate::Transaction; - use std::ops::DerefMut; - - use super::*; - use assert_matches::assert_matches; - use proptest::{prelude::ProptestConfig, proptest}; - use reth_db::{ - database::{Database, DatabaseGAT}, - mdbx::{test_utils::create_test_rw_db, Env, WriteMap}, - tables, - transaction::DbTxMut, - }; - use reth_primitives::{ - hex_literal::hex, - keccak256, - proofs::{genesis_state_root, KeccakHasher, EMPTY_ROOT}, - Address, Bytes, ChainSpec, Genesis, MAINNET, - }; - use std::{collections::HashMap, ops::Deref, str::FromStr}; - use triehash::sec_trie_root; - - fn load_mainnet_genesis_root(tx: &mut Transaction<'_, DB>) -> Genesis { - let ChainSpec { genesis, .. } = MAINNET.clone(); - - // Insert account state - for (address, account) in &genesis.alloc { - tx.put::( - *address, - Account { - nonce: account.nonce.unwrap_or_default(), - balance: account.balance, - bytecode_hash: None, - }, - ) - .unwrap(); - tx.put::( - keccak256(address), - Account { - nonce: account.nonce.unwrap_or_default(), - balance: account.balance, - bytecode_hash: None, - }, - ) - .unwrap(); - } - tx.commit().unwrap(); - - genesis - } - - fn create_test_loader<'tx, 'db>( - tx: &'tx Transaction<'db, Env>, - ) -> DBTrieLoader<'tx, > as DatabaseGAT<'db>>::TXMut> { - DBTrieLoader::new(tx.deref()) - } - - #[test] - fn empty_trie() { - let db = create_test_rw_db(); - let tx = Transaction::new(db.as_ref()).unwrap(); - assert_matches!( - create_test_loader(&tx).calculate_root(), - Ok(got) if got.root().unwrap() == EMPTY_ROOT - ); - } - - #[test] - fn single_account_trie() { - let db = create_test_rw_db(); - let tx = Transaction::new(db.as_ref()).unwrap(); - let address = Address::from_str("9fe4abd71ad081f091bd06dd1c16f7e92927561e").unwrap(); - let account = Account { nonce: 0, balance: U256::ZERO, bytecode_hash: None }; - tx.put::(keccak256(address), account).unwrap(); - let mut encoded_account = Vec::new(); - EthAccount::from(account).encode(&mut encoded_account); - let expected = H256(sec_trie_root::([(address, encoded_account)]).0); - assert_matches!( - create_test_loader(&tx).calculate_root(), - Ok(got) if got.root().unwrap() == expected - ); - } - - #[test] - fn two_accounts_trie() { - let db = create_test_rw_db(); - let mut tx = Transaction::new(db.as_ref()).unwrap(); - let mut trie = DBTrieLoader::new(tx.deref_mut()); - - let accounts = [ - ( - Address::from(hex!("9fe4abd71ad081f091bd06dd1c16f7e92927561e")), - Account { nonce: 155, balance: U256::from(414241124), bytecode_hash: None }, - ), - ( - Address::from(hex!("f8a6edaad4a332e6e550d0915a7fd5300b0b12d1")), - Account { nonce: 3, balance: U256::from(78978), bytecode_hash: None }, - ), - ]; - for (address, account) in accounts { - trie.tx.put::(keccak256(address), account).unwrap(); - } - let encoded_accounts = accounts.iter().map(|(k, v)| { - let mut out = Vec::new(); - EthAccount::from(*v).encode(&mut out); - (k, out) - }); - let expected = H256(sec_trie_root::(encoded_accounts).0); - assert_matches!( - trie.calculate_root(), - Ok(got) if got.root().unwrap() == expected - ); - } - - #[test] - fn single_storage_trie() { - let db = create_test_rw_db(); - let mut tx = Transaction::new(db.as_ref()).unwrap(); - let mut trie = DBTrieLoader::new(tx.deref_mut()); - - let address = Address::from_str("9fe4abd71ad081f091bd06dd1c16f7e92927561e").unwrap(); - let hashed_address = keccak256(address); - - let storage = Vec::from([(H256::from_low_u64_be(2), U256::from(1))]); - for (k, v) in storage.clone() { - trie.tx - .put::( - hashed_address, - StorageEntry { key: keccak256(k), value: v }, - ) - .unwrap(); - } - let encoded_storage = storage.iter().map(|(k, v)| { - let out = encode_fixed_size(v).to_vec(); - (k, out) - }); - let expected = H256(sec_trie_root::(encoded_storage).0); - let storage_trie_cursor = - Arc::new(Mutex::new(trie.tx.cursor_dup_write::().unwrap())); - assert_matches!( - trie.calculate_storage_root(hashed_address, storage_trie_cursor, None, None), - Ok(got) if got.root().unwrap() == expected - ); - } - - #[test] - fn single_account_with_storage_trie() { - let db = create_test_rw_db(); - let tx = Transaction::new(db.as_ref()).unwrap(); - - let address = Address::from_str("9fe4abd71ad081f091bd06dd1c16f7e92927561e").unwrap(); - let hashed_address = keccak256(address); - - let storage = HashMap::from([ - (H256::zero(), U256::from(3)), - (H256::from_low_u64_be(2), U256::from(1)), - ]); - let code = "el buen fla"; - let account = Account { - nonce: 155, - balance: U256::from(414241124u32), - bytecode_hash: Some(keccak256(code)), - }; - tx.put::(hashed_address, account).unwrap(); - - for (k, v) in storage.clone() { - tx.put::( - hashed_address, - StorageEntry { key: keccak256(k), value: v }, - ) - .unwrap(); - } - let mut out = Vec::new(); - - let encoded_storage = storage.iter().map(|(k, v)| { - let out = encode_fixed_size(v).to_vec(); - (k, out) - }); - - let storage_root = H256(sec_trie_root::(encoded_storage).0); - let eth_account = EthAccount::from(account).with_storage_root(storage_root); - eth_account.encode(&mut out); - - let expected = H256(sec_trie_root::([(address, out)]).0); - assert_matches!( - create_test_loader(&tx).calculate_root(), - Ok(got) if got.root().unwrap() == expected - ); - } - - #[test] - fn verify_genesis() { - let db = create_test_rw_db(); - let mut tx = Transaction::new(db.as_ref()).unwrap(); - - let genesis = load_mainnet_genesis_root(&mut tx); - - let state_root = genesis_state_root(&genesis.alloc); - - assert_matches!( - create_test_loader(&tx).calculate_root(), - Ok(got) if got.root().unwrap() == state_root - ); - } - - #[test] - fn gather_changes() { - let db = create_test_rw_db(); - let tx = Transaction::new(db.as_ref()).unwrap(); - - let address = Address::from_str("9fe4abd71ad081f091bd06dd1c16f7e92927561e").unwrap(); - let hashed_address = keccak256(address); - - let storage = HashMap::from([ - (H256::zero(), U256::from(3)), - (H256::from_low_u64_be(2), U256::from(1)), - ]); - let code = "el buen fla"; - let account = Account { - nonce: 155, - balance: U256::from(414241124u32), - bytecode_hash: Some(keccak256(code)), - }; - tx.put::(hashed_address, account).unwrap(); - tx.put::(31, AccountBeforeTx { address, info: None }).unwrap(); - - for (k, v) in storage { - tx.put::( - hashed_address, - StorageEntry { key: keccak256(k), value: v }, - ) - .unwrap(); - tx.put::( - (32, address).into(), - StorageEntry { key: k, value: U256::ZERO }, - ) - .unwrap(); - } - - let expected = BTreeMap::from([( - hashed_address, - BTreeSet::from([keccak256(H256::zero()), keccak256(H256::from_low_u64_be(2))]), - )]); - assert_matches!( - create_test_loader(&tx).gather_changes(32..33), - Ok(got) if got == expected - ); - } - - fn test_with_accounts(accounts: BTreeMap)>) { - let db = create_test_rw_db(); - let tx = Transaction::new(db.as_ref()).unwrap(); - - let encoded_accounts = accounts - .into_iter() - .map(|(address, (account, storage))| { - let hashed_address = keccak256(address); - tx.put::(hashed_address, account).unwrap(); - // This is to mimic real data. Only contract accounts have storage. - let storage_root = if account.has_bytecode() { - let encoded_storage = storage.into_iter().map(|StorageEntry { key, value }| { - let hashed_key = keccak256(key); - let out = encode_fixed_size(&value).to_vec(); - tx.put::( - hashed_address, - StorageEntry { key: hashed_key, value }, - ) - .unwrap(); - (key, out) - }); - H256(sec_trie_root::(encoded_storage).0) - } else { - EMPTY_ROOT - }; - let mut out = Vec::new(); - EthAccount::from(account).with_storage_root(storage_root).encode(&mut out); - (address, out) - }) - .collect::)>>(); - - let expected = H256(sec_trie_root::(encoded_accounts).0); - assert_matches!( - create_test_loader(&tx).calculate_root(), - Ok(got) if got.root().unwrap() == expected - , "where expected is {expected:?}"); - } - - #[test] - fn arbitrary() { - proptest!(ProptestConfig::with_cases(10), |(accounts: BTreeMap)>)| { - test_with_accounts(accounts); - }); - } - - #[test] - fn get_proof() { - let db = create_test_rw_db(); - let mut tx = Transaction::new(db.as_ref()).unwrap(); - - load_mainnet_genesis_root(&mut tx); - - let root = { - let mut trie = create_test_loader(&tx); - trie.calculate_root().expect("should be able to load trie").root().unwrap() - }; - - tx.commit().unwrap(); - - let address = Address::from(hex!("000d836201318ec6899a67540690382780743280")); - - let trie = create_test_loader(&tx); - let (proof, storage_root) = trie - .generate_account_proof(root, keccak256(address)) - .expect("failed to generate proof"); - - // values extracted from geth via rpc: - // { - // "method": "eth_getProof", - // "params": ["0x000d836201318ec6899a67540690382780743280", [], "0x0"] - // } - let expected = [ - hex!("f90211a090dcaf88c40c7bbc95a912cbdde67c175767b31173df9ee4b0d733bfdd511c43a0babe369f6b12092f49181ae04ca173fb68d1a5456f18d20fa32cba73954052bda0473ecf8a7e36a829e75039a3b055e51b8332cbf03324ab4af2066bbd6fbf0021a0bbda34753d7aa6c38e603f360244e8f59611921d9e1f128372fec0d586d4f9e0a04e44caecff45c9891f74f6a2156735886eedf6f1a733628ebc802ec79d844648a0a5f3f2f7542148c973977c8a1e154c4300fec92f755f7846f1b734d3ab1d90e7a0e823850f50bf72baae9d1733a36a444ab65d0a6faaba404f0583ce0ca4dad92da0f7a00cbe7d4b30b11faea3ae61b7f1f2b315b61d9f6bd68bfe587ad0eeceb721a07117ef9fc932f1a88e908eaead8565c19b5645dc9e5b1b6e841c5edbdfd71681a069eb2de283f32c11f859d7bcf93da23990d3e662935ed4d6b39ce3673ec84472a0203d26456312bbc4da5cd293b75b840fc5045e493d6f904d180823ec22bfed8ea09287b5c21f2254af4e64fca76acc5cd87399c7f1ede818db4326c98ce2dc2208a06fc2d754e304c48ce6a517753c62b1a9c1d5925b89707486d7fc08919e0a94eca07b1c54f15e299bd58bdfef9741538c7828b5d7d11a489f9c20d052b3471df475a051f9dd3739a927c89e357580a4c97b40234aa01ed3d5e0390dc982a7975880a0a089d613f26159af43616fd9455bb461f4869bfede26f2130835ed067a8b967bfb80").as_slice(), - hex!("f90211a0dae48f5b47930c28bb116fbd55e52cd47242c71bf55373b55eb2805ee2e4a929a00f1f37f337ec800e2e5974e2e7355f10f1a4832b39b846d916c3597a460e0676a0da8f627bb8fbeead17b318e0a8e4f528db310f591bb6ab2deda4a9f7ca902ab5a0971c662648d58295d0d0aa4b8055588da0037619951217c22052802549d94a2fa0ccc701efe4b3413fd6a61a6c9f40e955af774649a8d9fd212d046a5a39ddbb67a0d607cdb32e2bd635ee7f2f9e07bc94ddbd09b10ec0901b66628e15667aec570ba05b89203dc940e6fa70ec19ad4e01d01849d3a5baa0a8f9c0525256ed490b159fa0b84227d48df68aecc772939a59afa9e1a4ab578f7b698bdb1289e29b6044668ea0fd1c992070b94ace57e48cbf6511a16aa770c645f9f5efba87bbe59d0a042913a0e16a7ccea6748ae90de92f8aef3b3dc248a557b9ac4e296934313f24f7fced5fa042373cf4a00630d94de90d0a23b8f38ced6b0f7cb818b8925fee8f0c2a28a25aa05f89d2161c1741ff428864f7889866484cef622de5023a46e795dfdec336319fa07597a017664526c8c795ce1da27b8b72455c49657113e0455552dbc068c5ba31a0d5be9089012fda2c585a1b961e988ea5efcd3a06988e150a8682091f694b37c5a0f7b0352e38c315b2d9a14d51baea4ddee1770974c806e209355233c3c89dce6ea049bf6e8df0acafd0eff86defeeb305568e44d52d2235cf340ae15c6034e2b24180").as_slice(), - hex!("f901f1a0cf67e0f5d5f8d70e53a6278056a14ddca46846f5ef69c7bde6810d058d4a9eda80a06732ada65afd192197fe7ce57792a7f25d26978e64e954b7b84a1f7857ac279da05439f8d011683a6fc07efb90afca198fd7270c795c835c7c85d91402cda992eaa0449b93033b6152d289045fdb0bf3f44926f831566faa0e616b7be1abaad2cb2da031be6c3752bcd7afb99b1bb102baf200f8567c394d464315323a363697646616a0a40e3ed11d906749aa501279392ffde868bd35102db41364d9c601fd651f974aa0044bfa4fe8dd1a58e6c7144da79326e94d1331c0b00373f6ae7f3662f45534b7a098005e3e48db68cb1dc9b9f034ff74d2392028ddf718b0f2084133017da2c2e7a02a62bc40414ee95b02e202a9e89babbabd24bef0abc3fc6dcd3e9144ceb0b725a0239facd895bbf092830390a8676f34b35b29792ae561f196f86614e0448a5792a0a4080f88925daff6b4ce26d188428841bd65655d8e93509f2106020e76d41eefa04918987904be42a6894256ca60203283d1b89139cf21f09f5719c44b8cdbb8f7a06201fc3ef0827e594d953b5e3165520af4fceb719e11cc95fd8d3481519bfd8ca05d0e353d596bd725b09de49c01ede0f29023f0153d7b6d401556aeb525b2959ba0cd367d0679950e9c5f2aa4298fd4b081ade2ea429d71ff390c50f8520e16e30880").as_slice(), - hex!("f87180808080808080a0dbee8b33c73b86df839f309f7ac92eee19836e08b39302ffa33921b3c6a09f66a06068b283d51aeeee682b8fb5458354315d0b91737441ede5e137c18b4775174a8080808080a0fe7779c7d58c2fda43eba0a6644043c86ebb9ceb4836f89e30831f23eb059ece8080").as_slice(), - hex!("f8719f20b71c90b0d523dd5004cf206f325748da347685071b34812e21801f5270c4b84ff84d80890ad78ebc5ac6200000a056e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421a0c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470").as_slice(), - ]; - - assert_eq!(storage_root, EMPTY_ROOT); - - assert_eq!(proof.len(), 5); - - for (node, expected) in proof.into_iter().zip(expected.into_iter()) { - assert_eq!(Bytes::from(node.as_slice()), Bytes::from(expected)); - } - } - - #[test] - fn get_storage_proofs() { - let db = create_test_rw_db(); - let mut tx = Transaction::new(db.as_ref()).unwrap(); - - let address = Address::from_str("9fe4abd71ad081f091bd06dd1c16f7e92927561e").unwrap(); - let hashed_address = keccak256(address); - - let storage = HashMap::from([ - (H256::zero(), U256::from(3)), - (H256::from_low_u64_be(2), U256::from(1)), - ]); - - let code = "el buen fla"; - let account = Account { - nonce: 155, - balance: U256::from(414241124u32), - bytecode_hash: Some(keccak256(code)), - }; - tx.put::(hashed_address, account).unwrap(); - - for (k, v) in storage { - tx.put::( - hashed_address, - StorageEntry { key: keccak256(k), value: v }, - ) - .unwrap(); - } - - let root = { - let mut trie = create_test_loader(&tx); - trie.calculate_root().expect("should be able to load trie").root().unwrap() - }; - - tx.commit().unwrap(); - - let trie = create_test_loader(&tx); - let (account_proof, storage_root) = - trie.generate_account_proof(root, hashed_address).expect("failed to generate proof"); - - // values extracted from geth via rpc: - let expected_account = hex!("f86fa1205126413e7857595763591580306b3f228f999498c4c5dfa74f633364936e7651b84bf849819b8418b0d164a029ff6f4d518044318d75b118cf439d8d3d7249c8afcba06ba9ecdf8959410571a02ce1a85814ad94a94ed2a1abaf7c57e9b64326622c1b8c21b4ba4d0e7df61392").as_slice(); - let expected_storage = [ - [ - // 0x0000000000000000000000000000000000000000000000000000000000000002 - hex!("f8518080a04355bd3061ad2d17e0782413925b4fd81a56bd162d91eedb2a00d6c87611471480a015503e91f9250654cf72906e38a7cb14c3f1cc06658379d37f0c5b5c32482880808080808080808080808080").as_slice(), - hex!("e2a0305787fa12a823e0f2b7631cc41b3ba8828b3321ca811111fa75cd3aa3bb5ace01").as_slice(), - ], - [ - // 0x0000000000000000000000000000000000000000000000000000000000000000 - hex!("f8518080a04355bd3061ad2d17e0782413925b4fd81a56bd162d91eedb2a00d6c87611471480a015503e91f9250654cf72906e38a7cb14c3f1cc06658379d37f0c5b5c32482880808080808080808080808080").as_slice(), - hex!("e2a0390decd9548b62a8d60345a988386fc84ba6bc95484008f6362f93160ef3e56303").as_slice(), - ] - ]; - - assert!(storage_root != EMPTY_ROOT); - - assert_eq!(account_proof.len(), 1); - assert_eq!(account_proof[0], expected_account); - - let storage_proofs = trie - .generate_storage_proofs( - storage_root, - hashed_address, - &[keccak256(H256::from_low_u64_be(2)), keccak256(H256::zero())], - ) - .expect("couldn't generate storage proof"); - - for (proof, expected) in storage_proofs.into_iter().zip(expected_storage) { - assert_eq!(proof.len(), expected.len()); - for (got_node, expected_node) in proof.into_iter().zip(expected) { - assert_eq!(got_node, expected_node); - } - } - } -} diff --git a/crates/trie/Cargo.toml b/crates/trie/Cargo.toml index 9167b06219..83962a1e9e 100644 --- a/crates/trie/Cargo.toml +++ b/crates/trie/Cargo.toml @@ -12,7 +12,9 @@ Merkle trie implementation [dependencies] # reth reth-primitives = { path = "../primitives" } +reth-interfaces = { path = "../interfaces" } reth-rlp = { path = "../rlp" } +reth-db = { path = "../storage/db" } # tokio tokio = { version = "1.21.2", default-features = false, features = ["sync"] } @@ -22,11 +24,17 @@ tracing = "0.1" # misc hex = "0.4" +thiserror = "1.0" derive_more = "0.99" +# test-utils +triehash = { version = "0.8", optional = true } + [dev-dependencies] # reth reth-primitives = { path = "../primitives", features = ["test-utils", "arbitrary"] } +reth-db = { path = "../storage/db", features = ["test-utils"] } +reth-provider = { path = "../storage/provider" } # trie triehash = "0.8" @@ -35,3 +43,6 @@ triehash = "0.8" proptest = "1.0" tokio = { version = "1.21.2", default-features = false, features = ["sync", "rt", "macros"] } tokio-stream = "0.1.10" + +[features] +test-utils = ["triehash"] \ No newline at end of file diff --git a/crates/trie/src/cursor/account_cursor.rs b/crates/trie/src/cursor/account_cursor.rs new file mode 100644 index 0000000000..aeb2a90458 --- /dev/null +++ b/crates/trie/src/cursor/account_cursor.rs @@ -0,0 +1,94 @@ +use super::TrieCursor; +use reth_db::{ + cursor::{DbCursorRO, DbCursorRW}, + tables, Error, +}; +use reth_primitives::trie::{BranchNodeCompact, StoredNibbles}; + +/// A cursor over the account trie. +pub struct AccountTrieCursor(C); + +impl AccountTrieCursor { + /// Create a new account trie cursor. + pub fn new(cursor: C) -> Self { + Self(cursor) + } +} + +impl<'a, C> TrieCursor for AccountTrieCursor +where + C: DbCursorRO<'a, tables::AccountsTrie> + DbCursorRW<'a, tables::AccountsTrie>, +{ + fn seek_exact( + &mut self, + key: StoredNibbles, + ) -> Result, BranchNodeCompact)>, Error> { + Ok(self.0.seek_exact(key)?.map(|value| (value.0.inner.to_vec(), value.1))) + } + + fn seek(&mut self, key: StoredNibbles) -> Result, BranchNodeCompact)>, Error> { + Ok(self.0.seek(key)?.map(|value| (value.0.inner.to_vec(), value.1))) + } + + fn upsert(&mut self, key: StoredNibbles, value: BranchNodeCompact) -> Result<(), Error> { + self.0.upsert(key, value) + } + + fn delete_current(&mut self) -> Result<(), Error> { + self.0.delete_current() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use reth_db::{ + cursor::{DbCursorRO, DbCursorRW}, + mdbx::test_utils::create_test_rw_db, + tables, + transaction::DbTxMut, + }; + use reth_primitives::hex_literal::hex; + use reth_provider::Transaction; + + #[test] + fn test_account_trie_order() { + let db = create_test_rw_db(); + let tx = Transaction::new(db.as_ref()).unwrap(); + let mut cursor = tx.cursor_write::().unwrap(); + + let data = vec![ + hex!("0303040e").to_vec(), + hex!("030305").to_vec(), + hex!("03030500").to_vec(), + hex!("0303050a").to_vec(), + ]; + + for key in data.clone() { + cursor + .upsert( + key.into(), + BranchNodeCompact::new( + 0b0000_0010_0000_0001, + 0b0000_0010_0000_0001, + 0, + Vec::default(), + None, + ), + ) + .unwrap(); + } + + let db_data = + cursor.walk_range(..).unwrap().collect::, _>>().unwrap(); + assert_eq!(db_data[0].0.inner.to_vec(), data[0]); + assert_eq!(db_data[1].0.inner.to_vec(), data[1]); + assert_eq!(db_data[2].0.inner.to_vec(), data[2]); + assert_eq!(db_data[3].0.inner.to_vec(), data[3]); + + assert_eq!( + cursor.seek(hex!("0303040f").to_vec().into()).unwrap().map(|(k, _)| k.inner.to_vec()), + Some(data[1].clone()) + ); + } +} diff --git a/crates/trie/src/cursor/mod.rs b/crates/trie/src/cursor/mod.rs new file mode 100644 index 0000000000..ac47b1aa75 --- /dev/null +++ b/crates/trie/src/cursor/mod.rs @@ -0,0 +1,9 @@ +mod account_cursor; +mod storage_cursor; +mod subnode; +mod trie_cursor; + +pub use self::{ + account_cursor::AccountTrieCursor, storage_cursor::StorageTrieCursor, subnode::CursorSubNode, + trie_cursor::TrieCursor, +}; diff --git a/crates/trie/src/cursor/storage_cursor.rs b/crates/trie/src/cursor/storage_cursor.rs new file mode 100644 index 0000000000..0222953dad --- /dev/null +++ b/crates/trie/src/cursor/storage_cursor.rs @@ -0,0 +1,92 @@ +use super::TrieCursor; +use reth_db::{ + cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO, DbDupCursorRW}, + tables, Error, +}; +use reth_primitives::{ + trie::{BranchNodeCompact, StorageTrieEntry, StoredNibblesSubKey}, + H256, +}; + +/// A cursor over the storage trie. +pub struct StorageTrieCursor { + /// The underlying cursor. + pub cursor: C, + hashed_address: H256, +} + +impl StorageTrieCursor { + /// Create a new storage trie cursor. + pub fn new(cursor: C, hashed_address: H256) -> Self { + Self { cursor, hashed_address } + } +} + +impl<'a, C> TrieCursor for StorageTrieCursor +where + C: DbDupCursorRO<'a, tables::StoragesTrie> + + DbDupCursorRW<'a, tables::StoragesTrie> + + DbCursorRO<'a, tables::StoragesTrie> + + DbCursorRW<'a, tables::StoragesTrie>, +{ + fn seek_exact( + &mut self, + key: StoredNibblesSubKey, + ) -> Result, BranchNodeCompact)>, Error> { + Ok(self + .cursor + .seek_by_key_subkey(self.hashed_address, key.clone())? + .filter(|e| e.nibbles == key) + .map(|value| (value.nibbles.inner.to_vec(), value.node))) + } + + fn seek( + &mut self, + key: StoredNibblesSubKey, + ) -> Result, BranchNodeCompact)>, Error> { + Ok(self + .cursor + .seek_by_key_subkey(self.hashed_address, key)? + .map(|value| (value.nibbles.inner.to_vec(), value.node))) + } + + fn upsert(&mut self, key: StoredNibblesSubKey, value: BranchNodeCompact) -> Result<(), Error> { + if let Some(entry) = self.cursor.seek_by_key_subkey(self.hashed_address, key.clone())? { + // "seek exact" + if entry.nibbles == key { + self.cursor.delete_current()?; + } + } + + self.cursor.upsert(self.hashed_address, StorageTrieEntry { nibbles: key, node: value })?; + Ok(()) + } + + fn delete_current(&mut self) -> Result<(), Error> { + self.cursor.delete_current() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use reth_db::{mdbx::test_utils::create_test_rw_db, tables, transaction::DbTxMut}; + use reth_primitives::trie::BranchNodeCompact; + use reth_provider::Transaction; + + // tests that upsert and seek match on the storagetrie cursor + #[test] + fn test_storage_cursor_abstraction() { + let db = create_test_rw_db(); + let tx = Transaction::new(db.as_ref()).unwrap(); + let cursor = tx.cursor_dup_write::().unwrap(); + + let mut cursor = StorageTrieCursor::new(cursor, H256::random()); + + let key = vec![0x2, 0x3]; + let value = BranchNodeCompact::new(1, 1, 1, vec![H256::random()], None); + + cursor.upsert(key.clone().into(), value.clone()).unwrap(); + assert_eq!(cursor.seek(key.clone().into()).unwrap().unwrap().1, value); + } +} diff --git a/crates/trie/src/cursor/subnode.rs b/crates/trie/src/cursor/subnode.rs new file mode 100644 index 0000000000..cc4f292ae6 --- /dev/null +++ b/crates/trie/src/cursor/subnode.rs @@ -0,0 +1,101 @@ +use crate::{nodes::CHILD_INDEX_RANGE, Nibbles}; +use reth_primitives::{trie::BranchNodeCompact, H256}; + +/// Cursor for iterating over a subtrie. +#[derive(Clone)] +pub struct CursorSubNode { + /// The key of the current node. + pub key: Nibbles, + /// The index of the next child to visit. + pub nibble: i8, + /// The node itself. + pub node: Option, +} + +impl Default for CursorSubNode { + fn default() -> Self { + Self::new(Nibbles::default(), None) + } +} + +impl std::fmt::Debug for CursorSubNode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CursorSubNode") + .field("key", &self.key) + .field("nibble", &self.nibble) + .field("state_flag", &self.state_flag()) + .field("tree_flag", &self.tree_flag()) + .field("hash_flag", &self.hash_flag()) + .field("hash", &self.hash()) + .finish() + } +} + +impl CursorSubNode { + /// Creates a new `CursorSubNode` from a key and an optional node. + pub fn new(key: Nibbles, node: Option) -> Self { + // Find the first nibble that is set in the state mask of the node. + let nibble = match &node { + Some(n) if n.root_hash.is_none() => { + CHILD_INDEX_RANGE.clone().find(|i| n.state_mask.is_bit_set(*i)).unwrap() as i8 + } + _ => -1, + }; + CursorSubNode { key, node, nibble } + } + + /// Returns the full key of the current node. + pub fn full_key(&self) -> Nibbles { + let mut out = self.key.clone(); + if self.nibble >= 0 { + out.extend([self.nibble as u8]); + } + out + } + + /// Returns `true` if the state flag is set for the current nibble. + pub fn state_flag(&self) -> bool { + if let Some(node) = &self.node { + if self.nibble >= 0 { + return node.state_mask.is_bit_set(self.nibble as u8) + } + } + true + } + + /// Returns `true` if the tree flag is set for the current nibble. + pub fn tree_flag(&self) -> bool { + if let Some(node) = &self.node { + if self.nibble >= 0 { + return node.tree_mask.is_bit_set(self.nibble as u8) + } + } + true + } + + /// Returns `true` if the current nibble has a root hash. + pub fn hash_flag(&self) -> bool { + match &self.node { + Some(node) => match self.nibble { + // This guy has it + -1 => node.root_hash.is_some(), + // Or get it from the children + _ => node.hash_mask.is_bit_set(self.nibble as u8), + }, + None => false, + } + } + + /// Returns the root hash of the current node, if it has one. + pub fn hash(&self) -> Option { + if self.hash_flag() { + let node = self.node.as_ref().unwrap(); + match self.nibble { + -1 => node.root_hash, + _ => Some(node.hash_for_nibble(self.nibble as u8)), + } + } else { + None + } + } +} diff --git a/crates/trie/src/cursor/trie_cursor.rs b/crates/trie/src/cursor/trie_cursor.rs new file mode 100644 index 0000000000..97f2aa9b51 --- /dev/null +++ b/crates/trie/src/cursor/trie_cursor.rs @@ -0,0 +1,17 @@ +use reth_db::{table::Key, Error}; +use reth_primitives::trie::BranchNodeCompact; + +/// A cursor for navigating a trie that works with both Tables and DupSort tables. +pub trait TrieCursor { + /// Move the cursor to the key and return if it is an exact match. + fn seek_exact(&mut self, key: K) -> Result, BranchNodeCompact)>, Error>; + + /// Move the cursor to the key and return a value matching of greater than the key. + fn seek(&mut self, key: K) -> Result, BranchNodeCompact)>, Error>; + + /// Upsert the key/value pair. + fn upsert(&mut self, key: K, value: BranchNodeCompact) -> Result<(), Error>; + + /// Delete the key/value pair at the current cursor position. + fn delete_current(&mut self) -> Result<(), Error>; +} diff --git a/crates/trie/src/errors.rs b/crates/trie/src/errors.rs new file mode 100644 index 0000000000..3fc9f53aed --- /dev/null +++ b/crates/trie/src/errors.rs @@ -0,0 +1,20 @@ +use thiserror::Error; + +/// State root error. +#[derive(Error, Debug)] +pub enum StateRootError { + /// Internal database error. + #[error(transparent)] + DB(#[from] reth_db::Error), + /// Storage root error. + #[error(transparent)] + StorageRootError(#[from] StorageRootError), +} + +/// Storage root error. +#[derive(Error, Debug)] +pub enum StorageRootError { + /// Internal database error. + #[error(transparent)] + DB(#[from] reth_db::Error), +} diff --git a/crates/trie/src/hash_builder/mod.rs b/crates/trie/src/hash_builder/mod.rs index e3c612f493..25b0e76f01 100644 --- a/crates/trie/src/hash_builder/mod.rs +++ b/crates/trie/src/hash_builder/mod.rs @@ -8,14 +8,14 @@ use reth_primitives::{ trie::{BranchNodeCompact, TrieMask}, H256, }; -use std::fmt::Debug; -use tokio::sync::mpsc; +use std::{fmt::Debug, sync::mpsc}; mod value; use value::HashBuilderValue; /// A type alias for a sender of branch nodes. -pub type BranchNodeSender = mpsc::UnboundedSender<(Nibbles, BranchNodeCompact)>; +/// Branch nodes are sent by the Hash Builder to be stored in the database. +pub type BranchNodeSender = mpsc::Sender<(Nibbles, BranchNodeCompact)>; /// A component used to construct the root hash of the trie. The primary purpose of a Hash Builder /// is to build the Merkle proof that is essential for verifying the integrity and authenticity of @@ -363,8 +363,6 @@ mod tests { use proptest::prelude::*; use reth_primitives::{hex_literal::hex, proofs::KeccakHasher, H256, U256}; use std::collections::{BTreeMap, HashMap}; - use tokio::sync::mpsc::unbounded_channel; - use tokio_stream::{wrappers::UnboundedReceiverStream, StreamExt}; fn trie_root(iter: I) -> H256 where @@ -429,9 +427,9 @@ mod tests { }); } - #[tokio::test] - async fn test_generates_branch_node() { - let (sender, recv) = unbounded_channel(); + #[test] + fn test_generates_branch_node() { + let (sender, recv) = mpsc::channel(); let mut hb = HashBuilder::new(Some(sender)); // We have 1 branch node update to be stored at 0x01, indicated by the first nibble. @@ -481,8 +479,7 @@ mod tests { let root = hb.root(); drop(hb); - let receiver = UnboundedReceiverStream::new(recv); - let updates = receiver.collect::>().await; + let updates = recv.iter().collect::>(); let updates = updates.iter().cloned().collect::>(); let update = updates.get(&Nibbles::from(hex!("01").as_slice())).unwrap(); diff --git a/crates/trie/src/lib.rs b/crates/trie/src/lib.rs index 853030eb09..3380444f6c 100644 --- a/crates/trie/src/lib.rs +++ b/crates/trie/src/lib.rs @@ -24,3 +24,20 @@ pub mod hash_builder; /// The implementation of a container for storing intermediate changes to a trie. /// The container indicates when the trie has been modified. pub mod prefix_set; + +/// The cursor implementations for navigating account and storage tries. +pub mod cursor; + +/// The trie walker for iterating over the trie nodes. +pub mod walker; + +mod errors; +pub use errors::{StateRootError, StorageRootError}; + +/// The implementation of the Merkle Patricia Trie. +mod trie; +pub use trie::{BranchNodeUpdate, BranchNodeUpdateSender, StateRoot, StorageRoot}; + +/// Collection of trie-related test utilities. +#[cfg(any(test, feature = "test-utils"))] +pub mod test_utils; diff --git a/crates/trie/src/nibbles.rs b/crates/trie/src/nibbles.rs index 6896b9b5fc..b3699dce68 100644 --- a/crates/trie/src/nibbles.rs +++ b/crates/trie/src/nibbles.rs @@ -217,8 +217,8 @@ impl Nibbles { } /// Extend the current nibbles with another nibbles. - pub fn extend(&mut self, b: &Nibbles) { - self.hex_data.extend_from_slice(b); + pub fn extend(&mut self, b: impl AsRef<[u8]>) { + self.hex_data.extend_from_slice(b.as_ref()); } /// Truncate the current nibbles to the given length. diff --git a/crates/trie/src/nodes/branch.rs b/crates/trie/src/nodes/branch.rs index be3f34e52c..c6b58dfeb1 100644 --- a/crates/trie/src/nodes/branch.rs +++ b/crates/trie/src/nodes/branch.rs @@ -1,4 +1,4 @@ -use super::rlp_node; +use super::{rlp_node, CHILD_INDEX_RANGE}; use reth_primitives::{bytes::BytesMut, trie::TrieMask, H256}; use reth_rlp::{BufMut, EMPTY_STRING_CODE}; @@ -25,7 +25,7 @@ impl<'a> BranchNode<'a> { hash_mask: TrieMask, ) -> impl Iterator + '_ { let mut index = self.stack.len() - state_mask.count_ones() as usize; - (0..16).filter_map(move |digit| { + CHILD_INDEX_RANGE.filter_map(move |digit| { let mut child = None; if state_mask.is_bit_set(digit) { if hash_mask.is_bit_set(digit) { @@ -44,7 +44,7 @@ impl<'a> BranchNode<'a> { // Create the RLP header from the mask elements present. let mut i = first_child_idx; - let header = (0..16).fold( + let header = CHILD_INDEX_RANGE.fold( reth_rlp::Header { list: true, payload_length: 1 }, |mut header, digit| { if state_mask.is_bit_set(digit) { @@ -60,7 +60,7 @@ impl<'a> BranchNode<'a> { // Extend the RLP buffer with the present children let mut i = first_child_idx; - (0..16).for_each(|idx| { + CHILD_INDEX_RANGE.for_each(|idx| { if state_mask.is_bit_set(idx) { buf.extend_from_slice(&self.stack[i]); i += 1; diff --git a/crates/trie/src/nodes/mod.rs b/crates/trie/src/nodes/mod.rs index 347f0f722a..5cac694d50 100644 --- a/crates/trie/src/nodes/mod.rs +++ b/crates/trie/src/nodes/mod.rs @@ -1,5 +1,6 @@ use reth_primitives::{keccak256, H256}; use reth_rlp::EMPTY_STRING_CODE; +use std::ops::Range; mod branch; mod extension; @@ -7,6 +8,9 @@ mod leaf; pub use self::{branch::BranchNode, extension::ExtensionNode, leaf::LeafNode}; +/// The range of valid child indexes. +pub const CHILD_INDEX_RANGE: Range = 0..16; + /// Given an RLP encoded node, returns either RLP(Node) or RLP(keccak(RLP(node))) fn rlp_node(rlp: &[u8]) -> Vec { if rlp.len() < H256::len_bytes() { diff --git a/crates/trie/src/prefix_set/loader.rs b/crates/trie/src/prefix_set/loader.rs new file mode 100644 index 0000000000..3b628163e3 --- /dev/null +++ b/crates/trie/src/prefix_set/loader.rs @@ -0,0 +1,62 @@ +use super::PrefixSet; +use crate::Nibbles; +use derive_more::Deref; +use reth_db::{ + cursor::DbCursorRO, + models::{AccountBeforeTx, TransitionIdAddress}, + tables, + transaction::DbTx, + Error, +}; +use reth_primitives::{keccak256, Address, StorageEntry, TransitionId, H256}; +use std::{collections::HashMap, ops::Range}; + +/// A wrapper around a database transaction that loads prefix sets within a given transition range. +#[derive(Deref)] +pub struct PrefixSetLoader<'a, TX>(&'a TX); + +impl<'a, TX> PrefixSetLoader<'a, TX> { + /// Create a new loader. + pub fn new(tx: &'a TX) -> Self { + Self(tx) + } +} + +impl<'a, 'b, TX> PrefixSetLoader<'a, TX> +where + TX: DbTx<'b>, +{ + /// Load all account and storage changes for the given transition id range. + pub fn load( + self, + tid_range: Range, + ) -> Result<(PrefixSet, HashMap), Error> { + // Initialize prefix sets. + let mut account_prefix_set = PrefixSet::default(); + let mut storage_prefix_set: HashMap = HashMap::default(); + + // Walk account changeset and insert account prefixes. + let mut account_cursor = self.cursor_read::()?; + for account_entry in account_cursor.walk_range(tid_range.clone())? { + let (_, AccountBeforeTx { address, .. }) = account_entry?; + account_prefix_set.insert(Nibbles::unpack(keccak256(address))); + } + + // Walk storage changeset and insert storage prefixes as well as account prefixes if missing + // from the account prefix set. + let mut storage_cursor = self.cursor_dup_read::()?; + let start = TransitionIdAddress((tid_range.start, Address::zero())); + let end = TransitionIdAddress((tid_range.end, Address::zero())); + for storage_entry in storage_cursor.walk_range(start..end)? { + let (TransitionIdAddress((_, address)), StorageEntry { key, .. }) = storage_entry?; + let hashed_address = keccak256(address); + account_prefix_set.insert(Nibbles::unpack(hashed_address)); + storage_prefix_set + .entry(hashed_address) + .or_default() + .insert(Nibbles::unpack(keccak256(key))); + } + + Ok((account_prefix_set, storage_prefix_set)) + } +} diff --git a/crates/trie/src/prefix_set.rs b/crates/trie/src/prefix_set/mod.rs similarity index 97% rename from crates/trie/src/prefix_set.rs rename to crates/trie/src/prefix_set/mod.rs index 26210882df..21b4ec273c 100644 --- a/crates/trie/src/prefix_set.rs +++ b/crates/trie/src/prefix_set/mod.rs @@ -1,6 +1,9 @@ use crate::Nibbles; use std::collections::BTreeSet; +mod loader; +pub use loader::PrefixSetLoader; + /// A container for efficiently storing and checking for the presence of key prefixes. /// /// This data structure stores a set of `Nibbles` and provides methods to insert diff --git a/crates/trie/src/test_utils.rs b/crates/trie/src/test_utils.rs new file mode 100644 index 0000000000..690514acdf --- /dev/null +++ b/crates/trie/src/test_utils.rs @@ -0,0 +1,51 @@ +use crate::account::EthAccount; +use reth_primitives::{proofs::KeccakHasher, Account, Address, H256, U256}; +use reth_rlp::{encode_fixed_size, Encodable}; + +/// Re-export of [triehash]. +pub use triehash; + +/// Compute the state root of a given set of accounts using [triehash::sec_trie_root]. +pub fn state_root(accounts: I) -> H256 +where + I: Iterator, + S: IntoIterator, +{ + let encoded_accounts = accounts.map(|(address, (account, storage))| { + let storage_root = storage_root(storage.into_iter()); + let mut out = Vec::new(); + EthAccount::from(account).with_storage_root(storage_root).encode(&mut out); + (address, out) + }); + + triehash::sec_trie_root::(encoded_accounts) +} + +/// Compute the storage root for a given account using [triehash::sec_trie_root]. +pub fn storage_root>(storage: I) -> H256 { + let encoded_storage = storage.map(|(k, v)| (k, encode_fixed_size(&v).to_vec())); + triehash::sec_trie_root::(encoded_storage) +} + +/// Compute the state root of a given set of accounts with prehashed keys using +/// [triehash::trie_root]. +pub fn state_root_prehashed(accounts: I) -> H256 +where + I: Iterator, + S: IntoIterator, +{ + let encoded_accounts = accounts.map(|(address, (account, storage))| { + let storage_root = storage_root_prehashed(storage.into_iter()); + let mut out = Vec::new(); + EthAccount::from(account).with_storage_root(storage_root).encode(&mut out); + (address, out) + }); + + triehash::trie_root::(encoded_accounts) +} + +/// Compute the storage root for a given account with prehashed slots using [triehash::trie_root]. +pub fn storage_root_prehashed>(storage: I) -> H256 { + let encoded_storage = storage.map(|(k, v)| (k, encode_fixed_size(&v).to_vec())); + triehash::trie_root::(encoded_storage) +} diff --git a/crates/trie/src/trie.rs b/crates/trie/src/trie.rs new file mode 100644 index 0000000000..53023d9cdd --- /dev/null +++ b/crates/trie/src/trie.rs @@ -0,0 +1,1167 @@ +use crate::{ + account::EthAccount, + cursor::{AccountTrieCursor, StorageTrieCursor}, + hash_builder::HashBuilder, + nibbles::Nibbles, + prefix_set::{PrefixSet, PrefixSetLoader}, + walker::TrieWalker, + StateRootError, StorageRootError, +}; +use reth_db::{ + cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO, DbDupCursorRW}, + tables, + transaction::{DbTx, DbTxMut}, +}; +use reth_primitives::{ + keccak256, + proofs::EMPTY_ROOT, + trie::{BranchNodeCompact, StorageTrieEntry, StoredNibblesSubKey}, + Address, StorageEntry, TransitionId, H256, +}; +use reth_rlp::Encodable; +use std::{collections::HashMap, ops::Range, sync::mpsc}; + +/// The branch node update sender +pub type BranchNodeUpdateSender = mpsc::Sender; + +/// The branch node message to update the database. +#[derive(Debug, Clone)] +pub enum BranchNodeUpdate { + /// The account trie branch node. + Account(Nibbles, BranchNodeCompact), + /// The storage trie branch node with the hashed key of the account. + Storage(H256, Nibbles, BranchNodeCompact), +} + +/// StateRoot is used to compute the root node of a state trie. +pub struct StateRoot<'a, TX> { + /// A reference to the database transaction. + pub tx: &'a TX, + /// A set of account prefixes that have changed. + pub changed_account_prefixes: PrefixSet, + /// A map containing storage changes with the hashed address as key and a set of storage key + /// prefixes as the value. + pub changed_storage_prefixes: HashMap, +} + +impl<'a, TX> StateRoot<'a, TX> { + /// Create a new [StateRoot] instance. + pub fn new(tx: &'a TX) -> Self { + Self { + tx, + changed_account_prefixes: PrefixSet::default(), + changed_storage_prefixes: HashMap::default(), + } + } + + /// Set the changed account prefixes. + pub fn with_changed_account_prefixes(mut self, prefixes: PrefixSet) -> Self { + self.changed_account_prefixes = prefixes; + self + } + + /// Set the changed storage prefixes. + pub fn with_changed_storage_prefixes(mut self, prefixes: HashMap) -> Self { + self.changed_storage_prefixes = prefixes; + self + } +} + +impl<'a, 'tx, TX: DbTx<'tx> + DbTxMut<'tx>> StateRoot<'a, TX> { + /// Given a transition id range, identifies all the accounts and storage keys that + /// have changed. Calculates the new state root using existing unchanged intermediate nodes and + /// updating the nodes that are present in the prefix set. + /// + /// # Returns + /// + /// The updated state root hash. + pub fn incremental_root( + tx: &'a TX, + tid_range: Range, + branch_node_sender: Option, + ) -> Result { + tracing::debug!(target: "loader", "incremental state root"); + let (account_prefixes, storage_prefixes) = PrefixSetLoader::new(tx).load(tid_range)?; + let this = Self::new(tx) + .with_changed_account_prefixes(account_prefixes) + .with_changed_storage_prefixes(storage_prefixes); + + let root = this.root(branch_node_sender)?; + + Ok(root) + } + + /// 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, + branch_node_sender: Option, + ) -> Result { + tracing::debug!(target: "loader", "calculating state root"); + + let (sender, maybe_receiver) = match branch_node_sender { + Some(sender) => (sender, None), + None => { + let (sender, recv) = mpsc::channel(); + (sender, Some(recv)) + } + }; + + let mut hashed_account_cursor = self.tx.cursor_read::()?; + let mut trie_cursor = + AccountTrieCursor::new(self.tx.cursor_write::()?); + let mut walker = TrieWalker::new(&mut trie_cursor, self.changed_account_prefixes.clone()); + + let (account_branch_node_tx, account_branch_node_rx) = mpsc::channel(); + let mut hash_builder = + HashBuilder::default().with_branch_node_sender(account_branch_node_tx); + + 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, 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 { + tracing::trace!(target: "loader", "breaking, already detected"); + break + } + } + + // 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 intermeditate + // 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 = StorageRoot::new_hashed(self.tx, hashed_address) + .with_changed_prefixes( + self.changed_storage_prefixes + .get(&hashed_address) + .cloned() + .unwrap_or_default(), + ) + .root(Some(sender.clone()))?; + + let account = EthAccount::from(account).with_storage_root(storage_root); + let mut account_rlp = Vec::with_capacity(account.length()); + account.encode(&mut account_rlp); + + hash_builder.add_leaf(account_nibbles, &account_rlp); + + next_account_entry = hashed_account_cursor.next()?; + } + } + + let root = hash_builder.root(); + drop(hash_builder); + + for (nibbles, branch_node) in account_branch_node_rx.iter() { + let _ = sender.send(BranchNodeUpdate::Account(nibbles, branch_node)); + } + drop(sender); + + if let Some(receiver) = maybe_receiver { + let mut account_cursor = self.tx.cursor_write::()?; + let mut storage_cursor = self.tx.cursor_dup_write::()?; + + for update in receiver.iter() { + match update { + BranchNodeUpdate::Account(nibbles, branch_node) => { + if !nibbles.is_empty() { + account_cursor.upsert(nibbles.hex_data.into(), branch_node)?; + } + } + BranchNodeUpdate::Storage(hashed_address, nibbles, node) => { + let key: StoredNibblesSubKey = nibbles.hex_data.into(); + if let Some(entry) = + storage_cursor.seek_by_key_subkey(hashed_address, key.clone())? + { + // "seek exact" + if entry.nibbles == key { + storage_cursor.delete_current()?; + } + } + + storage_cursor + .upsert(hashed_address, StorageTrieEntry { nibbles: key, node })?; + } + } + } + } + + Ok(root) + } +} + +/// StorageRoot is used to compute the root node of an account storage trie. +pub struct StorageRoot<'a, TX> { + /// A reference to the database transaction. + pub tx: &'a TX, + /// The hashed address of an account. + pub hashed_address: H256, + /// The set of storage slot prefixes that have changed. + pub changed_prefixes: PrefixSet, +} + +impl<'a, TX> StorageRoot<'a, TX> { + /// Creates a new storage root calculator given an raw address. + pub fn new(tx: &'a TX, address: Address) -> Self { + Self::new_hashed(tx, keccak256(address)) + } + + /// Creates a new storage root calculator given a hashed address. + pub fn new_hashed(tx: &'a TX, hashed_address: H256) -> Self { + Self { tx, hashed_address, changed_prefixes: PrefixSet::default() } + } + + /// Set the changed prefixes. + pub fn with_changed_prefixes(mut self, prefixes: PrefixSet) -> Self { + self.changed_prefixes = prefixes; + self + } +} + +impl<'a, 'tx, TX: DbTx<'tx> + DbTxMut<'tx>> StorageRoot<'a, TX> { + /// Walks the hashed storage table entries for a given address and calculates the storage root. + pub fn root( + &self, + branch_node_update_sender: Option, + ) -> Result { + tracing::debug!(target: "trie::storage_root", hashed_address = ?self.hashed_address, "calculating storage root"); + + let mut hashed_storage_cursor = self.tx.cursor_dup_read::()?; + + let mut trie_cursor = StorageTrieCursor::new( + self.tx.cursor_dup_write::()?, + self.hashed_address, + ); + + // do not add a branch node on empty storage + if hashed_storage_cursor.seek_exact(self.hashed_address)?.is_none() { + if trie_cursor.cursor.seek_exact(self.hashed_address)?.is_some() { + trie_cursor.cursor.delete_current_duplicates()?; + } + return Ok(EMPTY_ROOT) + } + + let mut walker = TrieWalker::new(&mut trie_cursor, self.changed_prefixes.clone()); + + let (storage_branch_node_tx, storage_branch_node_rx) = mpsc::channel(); + let mut hash_builder = + HashBuilder::default().with_branch_node_sender(storage_branch_node_tx); + + 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_by_key_subkey(self.hashed_address, seek_key)?; + while let Some(StorageEntry { key: hashed_key, value }) = storage { + let storage_key_nibbles = Nibbles::unpack(hashed_key); + if let Some(ref key) = next_key { + if key < &storage_key_nibbles { + break + } + } + hash_builder + .add_leaf(storage_key_nibbles, reth_rlp::encode_fixed_size(&value).as_ref()); + storage = hashed_storage_cursor.next_dup_val()?; + } + } + + let root = hash_builder.root(); + drop(hash_builder); + + if let Some(sender) = branch_node_update_sender { + for (nibbles, branch_node) in storage_branch_node_rx.iter() { + let _ = sender.send(BranchNodeUpdate::Storage( + self.hashed_address, + nibbles, + branch_node, + )); + } + } + + tracing::debug!(target: "trie::storage_root", ?root, hashed_address = ?self.hashed_address, "calculated storage root"); + Ok(root) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::{ + state_root, state_root_prehashed, storage_root, storage_root_prehashed, + }; + use proptest::{prelude::ProptestConfig, proptest}; + use reth_db::{ + cursor::DbCursorRW, + mdbx::{test_utils::create_test_rw_db, Env, WriteMap}, + tables, + transaction::DbTxMut, + }; + use reth_primitives::{ + hex_literal::hex, keccak256, proofs::KeccakHasher, trie::TrieMask, Account, Address, H256, + U256, + }; + use reth_provider::Transaction; + use std::{ + collections::BTreeMap, + ops::{Deref, DerefMut, Mul}, + str::FromStr, + }; + + fn insert_account<'a, TX: DbTxMut<'a>>( + tx: &mut TX, + address: Address, + account: Account, + storage: &BTreeMap, + ) { + let hashed_address = keccak256(address); + tx.put::(hashed_address, account).unwrap(); + insert_storage(tx, hashed_address, storage); + } + + fn insert_storage<'a, TX: DbTxMut<'a>>( + tx: &mut TX, + hashed_address: H256, + storage: &BTreeMap, + ) { + for (k, v) in storage { + tx.put::( + hashed_address, + StorageEntry { key: keccak256(k), value: *v }, + ) + .unwrap(); + } + } + + fn incremental_vs_full_root(inputs: &[&str], modified: &str) { + let db = create_test_rw_db(); + let mut tx = Transaction::new(db.as_ref()).unwrap(); + let hashed_address = H256::from_low_u64_be(1); + + let mut hashed_storage_cursor = tx.cursor_dup_write::().unwrap(); + let data = inputs.iter().map(|x| H256::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 (branch_node_tx, branch_node_rx) = mpsc::channel(); + let _ = + StorageRoot::new_hashed(tx.deref(), hashed_address).root(Some(branch_node_tx)).unwrap(); + + // 1. Some state transition happens, update the hashed storage to the new value + let modified_key = H256::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::new_hashed(tx.deref(), hashed_address); + let modified_root = loader.root(None).unwrap(); + + // Update the intermediate roots table so that we can run the incremental verification + let mut trie_cursor = tx.cursor_dup_write::().unwrap(); + let updates = branch_node_rx.iter().collect::>(); + for update in updates { + match update { + BranchNodeUpdate::Storage(_, nibbles, node) => { + trie_cursor + .upsert( + hashed_address, + StorageTrieEntry { nibbles: nibbles.hex_data.into(), node }, + ) + .unwrap(); + } + _ => unreachable!(), + } + } + + // 3. Calculate the incremental root + let mut storage_changes = PrefixSet::default(); + storage_changes.insert(Nibbles::unpack(modified_key)); + let loader = StorageRoot::new_hashed(tx.deref_mut(), hashed_address) + .with_changed_prefixes(storage_changes); + let incremental_root = loader.root(None).unwrap(); + + assert_eq!(modified_root, incremental_root); + } + + #[test] + // TODO: Try to find the edge case by creating some more very complex trie. + 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))| { + tokio::runtime::Runtime::new().unwrap().block_on(async { + let (address, storage) = item; + + let hashed_address = keccak256(address); + let db = create_test_rw_db(); + let mut tx = Transaction::new(db.as_ref()).unwrap(); + for (key, value) in &storage { + tx.put::( + hashed_address, + StorageEntry { key: keccak256(key), value: *value }, + ) + .unwrap(); + } + tx.commit().unwrap(); + + let got = StorageRoot::new(tx.deref_mut(), address).root(None).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([(H256::from_low_u64_be(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([ + (H256::zero(), U256::from(3)), + (H256::from_low_u64_be(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 db = create_test_rw_db(); + let mut tx = Transaction::new(db.as_ref()).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(&mut *tx, address, account, &Default::default()); + tx.commit().unwrap(); + + let got = StorageRoot::new(tx.deref_mut(), address).root(None).unwrap(); + assert_eq!(got, EMPTY_ROOT); + } + + #[test] + // This ensures that the walker goes over all the storage slots + fn test_storage_root() { + let db = create_test_rw_db(); + let mut tx = Transaction::new(db.as_ref()).unwrap(); + + let address = Address::random(); + let storage = BTreeMap::from([ + (H256::zero(), U256::from(3)), + (H256::from_low_u64_be(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(&mut *tx, address, account, &storage); + tx.commit().unwrap(); + + let got = StorageRoot::new(tx.deref_mut(), address).root(None).unwrap(); + + assert_eq!(storage_root(storage.into_iter()), got); + } + + type State = BTreeMap)>; + + #[test] + fn arbitrary_state_root() { + proptest!( + ProptestConfig::with_cases(10), | (state: State) | { + // set the bytecodehash for the accounts so that storage root is computed + // this is needed because proptest will generate accs with empty bytecodehash + // but non-empty storage, which is obviously invalid + let state = state + .into_iter() + .map(|(addr, (mut acc, storage))| { + if !storage.is_empty() { + acc.bytecode_hash = Some(H256::random()); + } + (addr, (acc, storage)) + }) + .collect::>(); + test_state_root_with_state(state); + } + ); + } + + fn test_state_root_with_state(state: State) { + let db = create_test_rw_db(); + let mut tx = Transaction::new(db.as_ref()).unwrap(); + + for (address, (account, storage)) in &state { + insert_account(&mut *tx, *address, *account, storage) + } + tx.commit().unwrap(); + let expected = state_root(state.into_iter()); + + let got = StateRoot::new(tx.deref_mut()).root(None).unwrap(); + assert_eq!(expected, got); + } + + fn encode_account(account: Account, storage_root: Option) -> Vec { + let mut account = EthAccount::from(account); + if let Some(storage_root) = storage_root { + account = account.with_storage_root(storage_root); + } + let mut account_rlp = Vec::with_capacity(account.length()); + account.encode(&mut account_rlp); + account_rlp + } + + #[test] + fn storage_root_regression() { + let db = create_test_rw_db(); + let mut tx = Transaction::new(db.as_ref()).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)| (H256::from_str(slot).unwrap(), U256::from(val))), + ); + + let mut hashed_storage_cursor = tx.cursor_dup_write::().unwrap(); + for (hashed_slot, value) in storage.clone() { + hashed_storage_cursor.upsert(key3, StorageEntry { key: hashed_slot, value }).unwrap(); + } + tx.commit().unwrap(); + + let account3_storage_root = StorageRoot::new(tx.deref_mut(), address3).root(None).unwrap(); + let expected_root = storage_root_prehashed(storage.into_iter()); + 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)| (H256::from_str(slot).unwrap(), U256::from(val))), + ); + + let db = create_test_rw_db(); + let mut tx = Transaction::new(db.as_ref()).unwrap(); + + let mut hashed_account_cursor = tx.cursor_write::().unwrap(); + let mut hashed_storage_cursor = tx.cursor_dup_write::().unwrap(); + + let mut hash_builder = HashBuilder::default(); + + // Insert first account + let key1 = + H256::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.clone(), ..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 = + H256::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::new(tx.deref_mut(), address3).root(None).unwrap(); + hash_builder.add_leaf( + Nibbles::unpack(key3), + &encode_account(account3, Some(account3_storage_root)), + ); + + let key4a = + H256::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 = + H256::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 = + H256::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 = + H256::from_str("72861041bc90cd2f93777956f058a545412b56de79af5eb6b8075fe2eabbe015") + .unwrap(); + let computed_expected_root: H256 = triehash::trie_root::([ + (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 (branch_node_tx, branch_node_rx) = mpsc::channel(); + let loader = StateRoot::new(tx.deref()); + assert_eq!(loader.root(Some(branch_node_tx)).unwrap(), computed_expected_root); + + // Check account trie + drop(loader); + let updates = branch_node_rx.iter().collect::>(); + + let account_updates = updates + .iter() + .filter_map(|u| { + if let BranchNodeUpdate::Account(nibbles, node) = u { + Some((nibbles, node)) + } else { + None + } + }) + .collect::>(); + assert_eq!(account_updates.len(), 2); + + let (nibbles1a, node1a) = account_updates.last().unwrap(); + assert_eq!(**nibbles1a, Nibbles::from(&[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.first().unwrap(); + assert_eq!(**nibbles2a, Nibbles::from(&[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 = updates + .iter() + .filter_map(|u| { + if let BranchNodeUpdate::Storage(_, nibbles, node) = u { + Some((nibbles, node)) + } else { + None + } + }) + .collect::>(); + 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 = PrefixSet::default(); + prefix_set.insert(Nibbles::unpack(key4b)); + + let expected_state_root = + H256::from_str("8e263cd4eefb0c3cbbb14e5541a66a755cad25bcfab1e10dd9d706263e811b28") + .unwrap(); + + let (branch_node_tx, branch_node_rx) = mpsc::channel(); + let loader = StateRoot::new(tx.deref()).with_changed_account_prefixes(prefix_set); + assert_eq!(loader.root(Some(branch_node_tx)).unwrap(), expected_state_root); + + drop(loader); + let updates = branch_node_rx.iter().collect::>(); + + let account_updates = updates + .iter() + .filter_map(|u| { + if let BranchNodeUpdate::Account(nibbles, node) = u { + Some((nibbles, node)) + } else { + None + } + }) + .collect::>(); + assert_eq!(account_updates.len(), 2); + + let (nibbles1b, node1b) = account_updates.last().unwrap(); + assert_eq!(**nibbles1b, Nibbles::from(&[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.first().unwrap(); + assert_eq!(**nibbles2b, Nibbles::from(&[0xB, 0x0])); + assert_eq!(node2a, node2b); + tx.commit().unwrap(); + + { + let mut hashed_account_cursor = tx.cursor_write::().unwrap(); + + let account = hashed_account_cursor.seek_exact(key2).unwrap().unwrap(); + hashed_account_cursor.delete_current().unwrap(); + + let mut account_prefix_set = PrefixSet::default(); + account_prefix_set.insert(Nibbles::unpack(account.0)); + + let computed_expected_root: H256 = triehash::trie_root::([ + (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 (branch_node_tx, branch_node_rx) = mpsc::channel(); + let loader = + StateRoot::new(tx.deref_mut()).with_changed_account_prefixes(account_prefix_set); + assert_eq!(loader.root(Some(branch_node_tx)).unwrap(), computed_expected_root); + drop(loader); + + let updates = branch_node_rx.iter().collect::>(); + assert_eq!(updates.len(), 2); + + let account_updates = updates + .iter() + .filter_map(|u| { + if let BranchNodeUpdate::Account(nibbles, node) = u { + Some((nibbles, node)) + } else { + None + } + }) + .collect::>(); + assert_eq!(account_updates.len(), 1); + + let (nibbles1c, node1c) = account_updates.first().unwrap(); + assert_eq!(**nibbles1c, Nibbles::from(&[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]); + tx.drop().unwrap(); + } + + { + let mut hashed_account_cursor = tx.cursor_write::().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 = PrefixSet::default(); + account_prefix_set.insert(Nibbles::unpack(account2.0)); + account_prefix_set.insert(Nibbles::unpack(account3.0)); + + let computed_expected_root: H256 = triehash::trie_root::([ + (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 (branch_node_tx, branch_node_rx) = mpsc::channel(); + let loader = + StateRoot::new(tx.deref_mut()).with_changed_account_prefixes(account_prefix_set); + assert_eq!(loader.root(Some(branch_node_tx)).unwrap(), computed_expected_root); + drop(loader); + + let updates = branch_node_rx.iter().collect::>(); + assert_eq!(updates.len(), 1); // no storage root update + + let account_updates = updates + .iter() + .filter_map(|u| { + if let BranchNodeUpdate::Account(nibbles, node) = u { + Some((nibbles, node)) + } else { + None + } + }) + .collect::>(); + assert_eq!(account_updates.len(), 1); + + let (nibbles1d, node1d) = account_updates.first().unwrap(); + assert_eq!(**nibbles1d, Nibbles::from(&[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 db = create_test_rw_db(); + let mut tx = Transaction::new(db.as_ref()).unwrap(); + + let expected = extension_node_trie(&mut tx); + + let (sender, recv) = mpsc::channel(); + let loader = StateRoot::new(tx.deref_mut()); + let got = loader.root(Some(sender)).unwrap(); + assert_eq!(expected, got); + + // Check account trie + drop(loader); + let updates = recv.iter().collect::>(); + + let account_updates = updates + .into_iter() + .filter_map(|u| { + if let BranchNodeUpdate::Account(nibbles, node) = u { + Some((nibbles, node)) + } else { + None + } + }) + .collect::>(); + + assert_trie_updates(&account_updates); + } + + #[test] + + fn account_trie_around_extension_node_with_dbtrie() { + let db = create_test_rw_db(); + let mut tx = Transaction::new(db.as_ref()).unwrap(); + + let expected = extension_node_trie(&mut tx); + + let loader = StateRoot::new(tx.deref_mut()); + let got = loader.root(None).unwrap(); + assert_eq!(expected, got); + + drop(loader); + + // read the account updates from the db + let mut accounts_trie = tx.cursor_read::().unwrap(); + let mut walker = accounts_trie.walk(None).unwrap(); + let mut account_updates = BTreeMap::new(); + while let Some(item) = walker.next() { + let (key, node) = item.unwrap(); + account_updates.insert(Nibbles::from(key.inner.0.as_ref()), node); + } + + assert_trie_updates(&account_updates); + } + + // TODO: limit the thumber of test cases? + proptest! { + #[test] + fn fuzz_state_root_incremental(account_changes: [BTreeMap; 5]) { + tokio::runtime::Runtime::new().unwrap().block_on(async { + + let db = create_test_rw_db(); + let mut tx = Transaction::new(db.as_ref()).unwrap(); + let mut hashed_account_cursor = tx.cursor_write::().unwrap(); + + let mut state = BTreeMap::default(); + for accounts in account_changes { + let mut account_trie = tx.cursor_write::().unwrap(); + + let should_generate_changeset = !state.is_empty(); + let mut changes = PrefixSet::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 (branch_node_rx, branch_node_tx) = mpsc::channel(); + let account_storage_root = StateRoot::new(tx.deref_mut()).with_changed_account_prefixes(changes).root(Some(branch_node_rx)).unwrap(); + + state.append(&mut accounts.clone()); + let expected_root = state_root_prehashed( + state.clone().into_iter().map(|(key, balance)| (key, (Account { balance, ..Default::default() }, std::iter::empty()))), + ); + assert_eq!(expected_root, account_storage_root); + + let updates = branch_node_tx.iter().collect::>(); + for update in updates { + match update { + BranchNodeUpdate::Account(nibbles, node) => { + if !nibbles.is_empty() { + account_trie.upsert(nibbles.hex_data.into(), node).unwrap(); + } + } + BranchNodeUpdate::Storage(..) => unreachable!(), + }; + } + } + }); + } + } + + #[test] + fn storage_trie_around_extension_node() { + let db = create_test_rw_db(); + let mut tx = Transaction::new(db.as_ref()).unwrap(); + + let hashed_address = H256::random(); + let (expected_root, expected_updates) = + extension_node_storage_trie(&mut tx, hashed_address); + + let (sender, recv) = mpsc::channel(); + let loader = StorageRoot::new_hashed(tx.deref_mut(), hashed_address); + let got = loader.root(Some(sender)).unwrap(); + assert_eq!(expected_root, got); + + // Check account trie + drop(loader); + let updates = recv.iter().collect::>(); + + let storage_updates = updates + .into_iter() + .filter_map(|u| { + if let BranchNodeUpdate::Storage(_, nibbles, node) = u { + Some((nibbles, node)) + } else { + None + } + }) + .collect::>(); + assert_eq!(expected_updates, storage_updates); + + assert_trie_updates(&storage_updates); + } + + fn extension_node_storage_trie( + tx: &mut Transaction<'_, Env>, + hashed_address: H256, + ) -> (H256, BTreeMap) { + let value = U256::from(1); + + let mut hashed_storage = tx.cursor_write::().unwrap(); + + let (sender, receiver) = mpsc::channel(); + let mut hb = HashBuilder::new(Some(sender)); + + for key in [ + hex!("30af561000000000000000000000000000000000000000000000000000000000"), + hex!("30af569000000000000000000000000000000000000000000000000000000000"), + hex!("30af650000000000000000000000000000000000000000000000000000000000"), + hex!("30af6f0000000000000000000000000000000000000000000000000000000000"), + hex!("30af8f0000000000000000000000000000000000000000000000000000000000"), + hex!("3100000000000000000000000000000000000000000000000000000000000000"), + ] { + hashed_storage.upsert(hashed_address, StorageEntry { key: H256(key), value }).unwrap(); + hb.add_leaf(Nibbles::unpack(&key), &reth_rlp::encode_fixed_size(&value)); + } + let root = hb.root(); + + drop(hb); + let updates = receiver.iter().collect::>(); + let updates = updates.iter().cloned().collect(); + + (root, updates) + } + + fn extension_node_trie(tx: &mut Transaction<'_, Env>) -> H256 { + let a = Account { + nonce: 0, + balance: U256::from(1u64), + bytecode_hash: Some(H256::random()), + ..Default::default() + }; + let val = encode_account(a, None); + + let mut hashed_accounts = tx.cursor_write::().unwrap(); + let mut hb = HashBuilder::new(None); + + for key in [ + hex!("30af561000000000000000000000000000000000000000000000000000000000"), + hex!("30af569000000000000000000000000000000000000000000000000000000000"), + hex!("30af650000000000000000000000000000000000000000000000000000000000"), + hex!("30af6f0000000000000000000000000000000000000000000000000000000000"), + hex!("30af8f0000000000000000000000000000000000000000000000000000000000"), + hex!("3100000000000000000000000000000000000000000000000000000000000000"), + ] { + hashed_accounts.upsert(H256(key), a).unwrap(); + hb.add_leaf(Nibbles::unpack(&key), &val); + } + + hb.root() + } + + fn assert_trie_updates(account_updates: &BTreeMap) { + assert_eq!(account_updates.len(), 2); + + let node = account_updates.get(&Nibbles::from(vec![0x3])).unwrap(); + let expected = BranchNodeCompact::new(0b0011, 0b0001, 0b0000, vec![], None); + assert_eq!(node, &expected); + + let node = account_updates.get(&Nibbles::from(vec![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); + } +} diff --git a/crates/trie/src/walker.rs b/crates/trie/src/walker.rs new file mode 100644 index 0000000000..8a4dcf0fd5 --- /dev/null +++ b/crates/trie/src/walker.rs @@ -0,0 +1,339 @@ +use crate::{ + cursor::{CursorSubNode, TrieCursor}, + prefix_set::PrefixSet, + Nibbles, +}; +use reth_db::{table::Key, Error}; +use reth_primitives::{trie::BranchNodeCompact, H256}; +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 . +pub struct TrieWalker<'a, K, C> { + /// A mutable reference to a trie cursor instance used for navigating the trie. + pub cursor: &'a mut 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 + /// is determined by whether the current key's prefix is included in the prefix set and if the + /// hash flag is set. + pub can_skip_current_node: bool, + /// A `PrefixSet` representing the changes to be applied to the trie. + pub changes: PrefixSet, + __phantom: PhantomData, +} + +impl<'a, K: Key + From>, C: TrieCursor> TrieWalker<'a, K, C> { + /// Constructs a new TrieWalker, setting up the initial state of the stack and cursor. + pub fn new(cursor: &'a mut C, changes: PrefixSet) -> Self { + // Initialize the walker with a single empty stack element. + let mut this = Self { + cursor, + changes, + can_skip_current_node: false, + stack: vec![CursorSubNode::default()], + __phantom: PhantomData::default(), + }; + + // Set up the root node of the trie in the stack, if it exists. + if let Some((key, value)) = this.node(true).unwrap() { + this.stack[0] = CursorSubNode::new(key, Some(value)); + } + + // Update the skip state for the root node. + this.update_skip_node(); + this + } + + /// Prints the current stack of trie nodes. + pub fn print_stack(&self) { + println!("====================== STACK ======================"); + for node in &self.stack { + println!("{node:?}"); + } + println!("====================== END STACK ======================\n"); + } + + /// Advances the walker to the next trie node and updates the skip node flag. + /// + /// # Returns + /// + /// * `Result, Error>` - The next key in the trie or an error. + pub fn advance(&mut self) -> Result, Error> { + if let Some(last) = self.stack.last() { + if !self.can_skip_current_node && self.children_are_in_trie() { + // If we can't skip the current node and the children are in the trie, + // either consume the next node or move to the next sibling. + match last.nibble { + -1 => self.move_to_next_sibling(true)?, + _ => self.consume_node()?, + } + } else { + // If we can skip the current node, move to the next sibling. + self.move_to_next_sibling(false)?; + } + + // Update the skip node flag based on the new position in the trie. + self.update_skip_node(); + } + + // Return the current key. + Ok(self.key()) + } + + /// Retrieves the current root node from the DB, seeking either the exact node or the next one. + fn node(&mut self, exact: bool) -> Result, Error> { + let key = self.key().expect("key must exist"); + let entry = if exact { + self.cursor.seek_exact(key.hex_data.into())? + } else { + self.cursor.seek(key.hex_data.into())? + }; + + if let Some((_, node)) = &entry { + assert!(!node.state_mask.is_empty()); + } + + Ok(entry.map(|(k, v)| (Nibbles::from(k), v))) + } + + /// Consumes the next node in the trie, updating the stack. + fn consume_node(&mut self) -> Result<(), Error> { + let Some((key, node)) = self.node(false)? else { + // If no next node is found, clear the stack. + self.stack.clear(); + return Ok(()); + }; + + // Overwrite the root node's first nibble + // We need to sync the stack with the trie structure when consuming a new node. This is + // necessary for proper traversal and accurately representing the trie in the stack. + if !key.is_empty() && !self.stack.is_empty() { + self.stack[0].nibble = key[0] as i8; + } + + // Create a new CursorSubNode and push it to the stack. + let subnode = CursorSubNode::new(key, Some(node)); + let nibble = subnode.nibble; + self.stack.push(subnode); + self.update_skip_node(); + + // Delete the current node if it's included in the prefix set or it doesn't contain the root + // hash. + if !self.can_skip_current_node || nibble != -1 { + self.cursor.delete_current()?; + } + + Ok(()) + } + + /// Moves to the next sibling node in the trie, updating the stack. + fn move_to_next_sibling(&mut self, allow_root_to_child_nibble: bool) -> Result<(), Error> { + let Some(subnode) = self.stack.last_mut() else { + return Ok(()); + }; + + // Check if the walker needs to backtrack to the previous level in the trie during its + // traversal. + if subnode.nibble >= 15 || (subnode.nibble < 0 && !allow_root_to_child_nibble) { + self.stack.pop(); + self.move_to_next_sibling(false)?; + return Ok(()) + } + + subnode.nibble += 1; + + if subnode.node.is_none() { + return self.consume_node() + } + + // Find the next sibling with state. + while subnode.nibble < 16 { + if subnode.state_flag() { + return Ok(()) + } + subnode.nibble += 1; + } + + // Pop the current node and move to the next sibling. + self.stack.pop(); + self.move_to_next_sibling(false)?; + + Ok(()) + } + + /// Returns the current key in the trie. + pub fn key(&self) -> Option { + self.stack.last().map(|n| n.full_key()) + } + + /// Returns the current hash in the trie if any. + pub fn hash(&self) -> Option { + self.stack.last().and_then(|n| n.hash()) + } + + /// Indicates whether the children of the current node are present in the trie. + pub fn children_are_in_trie(&self) -> bool { + self.stack.last().map_or(false, |n| n.tree_flag()) + } + + /// Returns the next unprocessed key in the trie. + pub fn next_unprocessed_key(&self) -> Option { + self.key() + .as_ref() + .and_then(|key| { + if self.can_skip_current_node { + key.increment().map(|inc| inc.pack()) + } else { + Some(key.pack()) + } + }) + .map(|mut key| { + key.resize(32, 0); + H256::from_slice(key.as_slice()) + }) + } + + fn update_skip_node(&mut self) { + self.can_skip_current_node = if let Some(key) = self.key() { + let contains_prefix = self.changes.contains(key); + let hash_flag = self.stack.last().unwrap().hash_flag(); + !contains_prefix && hash_flag + } else { + false + }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cursor::{AccountTrieCursor, StorageTrieCursor}; + use reth_db::{mdbx::test_utils::create_test_rw_db, tables, transaction::DbTxMut}; + use reth_provider::Transaction; + + #[test] + fn walk_nodes_with_common_prefix() { + let inputs = vec![ + (vec![0x5u8], BranchNodeCompact::new(0b1_0000_0101, 0b1_0000_0100, 0, vec![], None)), + (vec![0x5u8, 0x2, 0xC], BranchNodeCompact::new(0b1000_0111, 0, 0, vec![], None)), + (vec![0x5u8, 0x8], BranchNodeCompact::new(0b0110, 0b0100, 0, vec![], None)), + ]; + let expected = vec![ + vec![0x5, 0x0], + // The [0x5, 0x2] prefix is shared by the first 2 nodes, however: + // 1. 0x2 for the first node points to the child node path + // 2. 0x2 for the second node is a key. + // So to proceed to add 1 and 3, we need to push the sibling first (0xC). + vec![0x5, 0x2], + vec![0x5, 0x2, 0xC, 0x0], + vec![0x5, 0x2, 0xC, 0x1], + vec![0x5, 0x2, 0xC, 0x2], + vec![0x5, 0x2, 0xC, 0x7], + vec![0x5, 0x8], + vec![0x5, 0x8, 0x1], + vec![0x5, 0x8, 0x2], + ]; + + let db = create_test_rw_db(); + let tx = Transaction::new(db.as_ref()).unwrap(); + let account_trie = + AccountTrieCursor::new(tx.cursor_write::().unwrap()); + test_cursor(account_trie, &inputs, &expected); + + let storage_trie = StorageTrieCursor::new( + tx.cursor_dup_write::().unwrap(), + H256::random(), + ); + test_cursor(storage_trie, &inputs, &expected); + } + + fn test_cursor(mut trie: T, inputs: &[(Vec, BranchNodeCompact)], expected: &[Vec]) + where + K: Key + From>, + T: TrieCursor, + { + for (k, v) in inputs { + trie.upsert(k.clone().into(), v.clone()).unwrap(); + } + + let mut walker = TrieWalker::new(&mut trie, Default::default()); + assert!(walker.key().unwrap().is_empty()); + + // We're traversing the path in lexigraphical order. + for expected in expected { + let got = walker.advance().unwrap(); + assert_eq!(got.unwrap(), Nibbles::from(&expected[..])); + } + + // There should be 8 paths traversed in total from 3 branches. + let got = walker.advance().unwrap(); + assert!(got.is_none()); + } + + #[test] + fn cursor_rootnode_with_changesets() { + let db = create_test_rw_db(); + let tx = Transaction::new(db.as_ref()).unwrap(); + + let mut trie = StorageTrieCursor::new( + tx.cursor_dup_write::().unwrap(), + H256::random(), + ); + + let nodes = vec![ + ( + vec![], + BranchNodeCompact::new( + // 2 and 4 are set + 0b10100, + 0b00100, + 0, + vec![], + Some(H256::random()), + ), + ), + ( + vec![0x2], + BranchNodeCompact::new( + // 1 is set + 0b00010, + 0, + 0b00010, + vec![H256::random()], + None, + ), + ), + ]; + + for (k, v) in nodes { + trie.upsert(k.into(), v).unwrap(); + } + + // No changes + let mut cursor = TrieWalker::new(&mut trie, Default::default()); + assert_eq!(cursor.key(), Some(Nibbles::from(vec![]))); // root + assert!(cursor.can_skip_current_node); // due to root_hash + cursor.advance().unwrap(); // skips to the end of trie + assert_eq!(cursor.key(), None); + + // We insert something that's not part of the existing trie/prefix. + let mut changed = PrefixSet::default(); + changed.insert(&[0xF, 0x1]); + let mut cursor = TrieWalker::new(&mut trie, changed); + + // Root node + assert_eq!(cursor.key(), Some(Nibbles::from(vec![]))); + // Should not be able to skip state due to the changed values + assert!(!cursor.can_skip_current_node); + cursor.advance().unwrap(); + assert_eq!(cursor.key(), Some(Nibbles::from(vec![0x2]))); + cursor.advance().unwrap(); + assert_eq!(cursor.key(), Some(Nibbles::from(vec![0x2, 0x1]))); + cursor.advance().unwrap(); + assert_eq!(cursor.key(), Some(Nibbles::from(vec![0x4]))); + + cursor.advance().unwrap(); + assert_eq!(cursor.key(), None); // the end of trie + } +}