diff --git a/Cargo.lock b/Cargo.lock index 20487d39bc848..fb4888edd125b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6957,6 +6957,7 @@ dependencies = [ "reth-errors", "reth-ethereum-engine-primitives", "reth-evm", + "reth-execution-errors", "reth-exex-types", "reth-metrics", "reth-network-p2p", diff --git a/crates/chain-state/src/in_memory.rs b/crates/chain-state/src/in_memory.rs index 79eb8b11fced9..d95288d83f700 100644 --- a/crates/chain-state/src/in_memory.rs +++ b/crates/chain-state/src/in_memory.rs @@ -828,6 +828,8 @@ impl NewCanonicalChain { #[cfg(test)] mod tests { + use std::collections::HashSet; + use super::*; use crate::test_utils::TestBlockBuilder; use rand::Rng; @@ -839,7 +841,7 @@ mod tests { AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider, }; - use reth_trie::{prefix_set::TriePrefixSetsMut, AccountProof, HashedStorage}; + use reth_trie::{prefix_set::TriePrefixSetsMut, AccountProof, HashedStorage, MultiProof}; fn create_mock_state( test_block_builder: &mut TestBlockBuilder, @@ -959,6 +961,14 @@ mod tests { Ok(AccountProof::new(Address::random())) } + fn multiproof( + &self, + hashed_state: HashedPostState, + targets: HashMap>, + ) -> ProviderResult { + Ok(MultiProof::default()) + } + fn witness( &self, _overlay: HashedPostState, diff --git a/crates/chain-state/src/memory_overlay.rs b/crates/chain-state/src/memory_overlay.rs index 1782627b9142a..bc2f9193f7c9c 100644 --- a/crates/chain-state/src/memory_overlay.rs +++ b/crates/chain-state/src/memory_overlay.rs @@ -9,9 +9,12 @@ use reth_storage_api::{ }; use reth_trie::{ prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState, - HashedStorage, + HashedStorage, MultiProof, +}; +use std::{ + collections::{HashMap, HashSet}, + sync::OnceLock, }; -use std::{collections::HashMap, sync::OnceLock}; /// A state provider that stores references to in-memory blocks along with their state as well as /// the historical state provider for fallback lookups. @@ -168,6 +171,17 @@ impl StateProofProvider for MemoryOverlayStateProvider { self.historical.proof(hashed_state, address, slots) } + // TODO: !!!!! + fn multiproof( + &self, + state: HashedPostState, + targets: HashMap>, + ) -> ProviderResult { + let mut hashed_state = self.trie_state().hashed_state.clone(); + hashed_state.extend(state); + self.historical.multiproof(hashed_state, targets) + } + // TODO: Currently this does not reuse available in-memory trie nodes. fn witness( &self, diff --git a/crates/engine/service/src/service.rs b/crates/engine/service/src/service.rs index 2be1bcca8643c..13682053b604e 100644 --- a/crates/engine/service/src/service.rs +++ b/crates/engine/service/src/service.rs @@ -91,6 +91,7 @@ where blockchain_db, executor_factory, consensus, + pipeline_task_spawner.clone(), payload_validator, persistence_handle, payload_builder, diff --git a/crates/engine/tree/Cargo.toml b/crates/engine/tree/Cargo.toml index 2c8d1922a039b..7f3c7100b8327 100644 --- a/crates/engine/tree/Cargo.toml +++ b/crates/engine/tree/Cargo.toml @@ -20,6 +20,7 @@ reth-consensus.workspace = true reth-engine-primitives.workspace = true reth-errors.workspace = true reth-evm.workspace = true +reth-execution-errors.workspace = true reth-network-p2p.workspace = true reth-payload-builder.workspace = true reth-payload-primitives.workspace = true @@ -34,6 +35,7 @@ reth-tasks.workspace = true reth-node-types.workspace = true reth-trie.workspace = true reth-trie-parallel.workspace = true +alloy-rlp.workspace = true # common futures.workspace = true diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs index 11b2e517003ac..49dc3c5ebd7e2 100644 --- a/crates/engine/tree/src/tree/mod.rs +++ b/crates/engine/tree/src/tree/mod.rs @@ -1,9 +1,10 @@ use crate::{ backfill::{BackfillAction, BackfillSyncState}, chain::FromOrchestrator, - engine::{DownloadRequest, EngineApiEvent, FromEngine}, + engine::{DownloadRequest, EngineApiEvent, EngineApiRequest, FromEngine}, persistence::PersistenceHandle, }; +use alloy_rlp::BufMut; use reth_beacon_consensus::{ BeaconConsensusEngineEvent, BeaconEngineMessage, ForkchoiceStateTracker, InvalidHeaderCache, OnForkChoiceUpdated, MIN_BLOCKS_FOR_PIPELINE_RUN, @@ -28,8 +29,8 @@ use reth_primitives::{ }; use reth_provider::{ providers::ConsistentDbView, BlockReader, DatabaseProviderFactory, ExecutionOutcome, - ProviderError, StateProviderBox, StateProviderFactory, StateReader, StateRootProvider, - TransactionVariant, + ProviderError, StateProvider, StateProviderBox, StateProviderFactory, StateReader, + StateRootProvider, TransactionVariant, }; use reth_revm::database::StateProviderDatabase; use reth_rpc_types::{ @@ -40,7 +41,11 @@ use reth_rpc_types::{ ExecutionPayload, }; use reth_stages_api::ControlFlow; -use reth_trie::{prefix_set::TriePrefixSetsMut, updates::TrieUpdates, HashedPostState}; +use reth_tasks::TaskSpawner; +use reth_trie::{ + prefix_set::TriePrefixSetsMut, updates::TrieUpdates, witness::TrieWitness, HashedPostState, + MultiProof, EMPTY_ROOT_HASH, +}; use reth_trie_parallel::parallel_root::ParallelStateRoot; use std::{ cmp::Ordering, @@ -53,21 +58,25 @@ use std::{ }, time::Instant, }; +use streaming_database::{StateAccess, StreamingDatabase}; use tokio::sync::{ - mpsc::{UnboundedReceiver, UnboundedSender}, - oneshot, - oneshot::error::TryRecvError, + mpsc::{self, UnboundedReceiver, UnboundedSender}, + oneshot::{self, error::TryRecvError}, }; use tracing::*; mod config; -mod invalid_block_hook; -mod metrics; -use crate::{engine::EngineApiRequest, tree::metrics::EngineApiMetrics}; pub use config::TreeConfig; + +mod invalid_block_hook; pub use invalid_block_hook::{InvalidBlockHooks, NoopInvalidBlockHook}; pub use reth_engine_primitives::InvalidBlockHook; +mod metrics; +use metrics::EngineApiMetrics; + +mod streaming_database; + /// Keeps track of the state of the tree. /// /// ## Invariants @@ -461,6 +470,7 @@ pub struct EngineApiTreeHandler { provider: P, executor_provider: E, consensus: Arc, + task_spawner: Box, payload_validator: ExecutionPayloadValidator, /// Keeps track of internals such as executed and buffered blocks. state: EngineApiTreeState, @@ -531,6 +541,7 @@ where provider: P, executor_provider: E, consensus: Arc, + task_spawner: Box, payload_validator: ExecutionPayloadValidator, outgoing: UnboundedSender, state: EngineApiTreeState, @@ -545,6 +556,7 @@ where provider, executor_provider, consensus, + task_spawner, payload_validator, incoming, outgoing, @@ -576,6 +588,7 @@ where provider: P, executor_provider: E, consensus: Arc, + task_spawner: Box, payload_validator: ExecutionPayloadValidator, persistence: PersistenceHandle, payload_builder: PayloadBuilderHandle, @@ -604,6 +617,7 @@ where provider, executor_provider, consensus, + task_spawner, payload_validator, tx, state, @@ -2127,6 +2141,7 @@ where missing_ancestor, })) }; + let proof_provider = self.state_provider(block.parent_hash)?.unwrap(); // now validate against the parent let parent_block = self.sealed_header_by_hash(block.parent_hash)?.ok_or_else(|| { @@ -2139,7 +2154,39 @@ where return Err(e.into()) } - let executor = self.executor_provider.executor(StateProviderDatabase::new(&state_provider)); + let (database, mut rx) = StreamingDatabase::new_with_rx(&state_provider); + let executor = self.executor_provider.executor(StateProviderDatabase::new(database)); + + // Spawn proof gathering task + let (multiproof_tx, multiproof_rx) = oneshot::channel(); + self.task_spawner.spawn(Box::pin(async move { + let mut multiproof = MultiProof::default(); + while let Some(next) = rx.recv().await { + let mut targets = HashMap::from([match next { + StateAccess::Account(address) => (address, HashSet::default()), + StateAccess::StorageSlot(address, slot) => (address, HashSet::from(slot)), + }]); + + 'inner: loop { + let Ok(next) = rx.try_recv() else { + break 'inner; + }; + match next { + StateAccess::Account(address) => { + targets.entry(address).or_default(); + } + StateAccess::StorageSlot(address, slot) => { + targets.entry(address).or_default().insert(slot); + } + } + } + + multiproof + .extend(proof_provider.multiproof(HashPostState::default(), targets).unwrap()); + } + + let _ = multiproof_tx.send((proof_provider, multiproof)); + })); let block_number = block.number; let block_hash = block.hash(); @@ -2169,6 +2216,8 @@ where let hashed_state = HashedPostState::from_bundle_state(&output.state.state); + let (proof_provider, multiproof) = multiproof_rx.blocking_recv().unwrap(); + let root_time = Instant::now(); let mut state_root_result = None; @@ -2249,6 +2298,104 @@ where Ok(InsertPayloadOk2::Inserted(BlockStatus2::Valid)) } + fn compute_state_root_from_proofs( + &self, + state_provider: Box, + state: &HashedPostState, + multiproof: MultiProof, + ) -> ProviderResult<(B256, TrieUpdates)> { + let prefix_sets = state.construct_prefix_sets(); + + let mut account_rlp = Vec::with_capacity(128); + let mut account_trie_nodes = BTreeMap::default(); + for (hashed_address, hashed_slots) in proof_targets { + let storage_multiproof = + multiproof.storages.remove(&hashed_address).unwrap_or_default(); + + // Gather and record account trie nodes. + let account = state + .accounts + .get(&hashed_address) + .ok_or(TrieWitnessError::MissingAccount(hashed_address))?; + let value = if account.is_some() || storage_multiproof.root != EMPTY_ROOT_HASH { + account_rlp.clear(); + TrieAccount::from((account.unwrap_or_default(), storage_multiproof.root)) + .encode(&mut account_rlp as &mut dyn BufMut); + Some(account_rlp.clone()) + } else { + None + }; + let key = Nibbles::unpack(hashed_address); + let proof = multiproof.account_subtree.iter().filter(|e| key.starts_with(e.0)); + account_trie_nodes.extend(self.target_nodes(key.clone(), value, proof)?); + + // Gather and record storage trie nodes for this account. + let mut storage_trie_nodes = BTreeMap::default(); + let storage = state.storages.get(&hashed_address); + for hashed_slot in hashed_slots { + let slot_key = Nibbles::unpack(hashed_slot); + let slot_value = storage + .and_then(|s| s.storage.get(&hashed_slot)) + .filter(|v| !v.is_zero()) + .map(|v| alloy_rlp::encode_fixed_size(v).to_vec()); + let proof = storage_multiproof.subtree.iter().filter(|e| slot_key.starts_with(e.0)); + storage_trie_nodes.extend(self.target_nodes( + slot_key.clone(), + slot_value, + proof, + )?); + } + + let (storage_root, storage_trie_updates) = + TrieWitness::next_root_from_proofs(storage_trie_nodes, true, |key: Nibbles| { + // Right pad the target with 0s. + let mut padded_key = key.pack(); + padded_key.resize(32, 0); + let target = (hashed_address, Vec::from([B256::from_slice(&padded_key)])); + // TODO: proof_provider.hashed_proof + let mut proof = Proof::new( + self.trie_cursor_factory.clone(), + self.hashed_cursor_factory.clone(), + ) + .with_prefix_sets_mut(self.prefix_sets.clone()) + .with_targets(HashMap::from([target])) + .storage_multiproof(hashed_address)?; + + // The subtree only contains the proof for a single target. + let node = proof + .subtree + .remove(&key) + .ok_or(TrieWitnessError::MissingTargetNode(key))?; + Ok(node) + })?; + debug_assert_eq!(storage_multiproof.root, storage_root); + } + + let (state_root, trie_updates) = + TrieWitness::next_root_from_proofs(account_trie_nodes, true, |key: Nibbles| { + // Right pad the target with 0s. + let mut padded_key = key.pack(); + padded_key.resize(32, 0); + // TODO: proof_provider.hashed_proof + let mut proof = Proof::new( + self.trie_cursor_factory.clone(), + self.hashed_cursor_factory.clone(), + ) + .with_prefix_sets_mut(self.prefix_sets.clone()) + .with_targets(HashMap::from([(B256::from_slice(&padded_key), Vec::new())])) + .multiproof()?; + + // The subtree only contains the proof for a single target. + let node = proof + .account_subtree + .remove(&key) + .ok_or(TrieWitnessError::MissingTargetNode(key))?; + Ok(node) + })?; + + todo!() + } + /// Compute state root for the given hashed post state in parallel. /// /// # Returns diff --git a/crates/engine/tree/src/tree/streaming_database.rs b/crates/engine/tree/src/tree/streaming_database.rs new file mode 100644 index 0000000000000..2f9b8ccbb6183 --- /dev/null +++ b/crates/engine/tree/src/tree/streaming_database.rs @@ -0,0 +1,53 @@ +use reth_errors::ProviderResult; +use reth_primitives::{Account, Address, BlockNumber, Bytecode, StorageKey, StorageValue, B256}; +use reth_revm::database::EvmStateProvider; +use tokio::sync::mpsc; + +#[derive(Clone, Copy, Debug)] +pub enum StateAccess { + Account(Address), + StorageSlot(Address, B256), +} + +pub struct StreamingDatabase { + inner: DB, + sender: mpsc::UnboundedSender, +} + +impl StreamingDatabase { + pub fn new(inner: DB, sender: mpsc::UnboundedSender) -> Self { + Self { inner, sender } + } + + pub fn new_with_rx(inner: DB) -> (Self, mpsc::UnboundedReceiver) { + let (tx, rx) = mpsc::unbounded_channel(); + (Self::new(inner, tx), rx) + } +} + +impl EvmStateProvider for StreamingDatabase +where + DB: EvmStateProvider, +{ + fn basic_account(&self, address: Address) -> ProviderResult> { + let _ = self.sender.send(StateAccess::Account(address)); + self.inner.basic_account(address) + } + + fn storage( + &self, + account: Address, + storage_key: StorageKey, + ) -> ProviderResult> { + let _ = self.sender.send(StateAccess::StorageSlot(account, storage_key)); + self.inner.storage(account, storage_key) + } + + fn block_hash(&self, number: BlockNumber) -> ProviderResult> { + self.inner.block_hash(number) + } + + fn bytecode_by_hash(&self, code_hash: B256) -> ProviderResult> { + self.bytecode_by_hash(code_hash) + } +} diff --git a/crates/revm/src/test_utils.rs b/crates/revm/src/test_utils.rs index 92f09e461d96c..067515eb212c7 100644 --- a/crates/revm/src/test_utils.rs +++ b/crates/revm/src/test_utils.rs @@ -1,3 +1,5 @@ +use std::collections::HashSet; + use crate::precompile::HashMap; use reth_primitives::{ keccak256, Account, Address, BlockNumber, Bytecode, Bytes, StorageKey, B256, U256, @@ -9,7 +11,7 @@ use reth_storage_api::{ use reth_storage_errors::provider::ProviderResult; use reth_trie::{ prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState, - HashedStorage, + HashedStorage, MultiProof, }; #[cfg(not(feature = "std"))] @@ -122,6 +124,14 @@ impl StateProofProvider for StateProviderTest { unimplemented!("proof generation is not supported") } + fn multiproof( + &self, + _hashed_state: HashedPostState, + _targets: HashMap>, + ) -> ProviderResult { + unimplemented!("proof generation is not supported") + } + fn witness( &self, _overlay: HashedPostState, diff --git a/crates/rpc/rpc-eth-types/src/cache/db.rs b/crates/rpc/rpc-eth-types/src/cache/db.rs index 60b895c47baec..66720253e6b1b 100644 --- a/crates/rpc/rpc-eth-types/src/cache/db.rs +++ b/crates/rpc/rpc-eth-types/src/cache/db.rs @@ -71,6 +71,14 @@ impl<'a> reth_storage_api::StateProofProvider for StateProviderTraitObjWrapper<' self.0.proof(hashed_state, address, slots) } + fn multiproof( + &self, + hashed_state: reth_trie::HashedPostState, + targets: std::collections::HashMap>, + ) -> ProviderResult { + self.0.multiproof(hashed_state, targets) + } + fn witness( &self, overlay: reth_trie::HashedPostState, diff --git a/crates/storage/provider/src/providers/bundle_state_provider.rs b/crates/storage/provider/src/providers/bundle_state_provider.rs index 50c88a51fcd3c..60b826c43d29c 100644 --- a/crates/storage/provider/src/providers/bundle_state_provider.rs +++ b/crates/storage/provider/src/providers/bundle_state_provider.rs @@ -6,9 +6,9 @@ use reth_storage_api::{StateProofProvider, StorageRootProvider}; use reth_storage_errors::provider::ProviderResult; use reth_trie::{ prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState, - HashedStorage, + HashedStorage, MultiProof, }; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; /// A state provider that resolves to data from either a wrapped [`crate::ExecutionOutcome`] /// or an underlying state provider. @@ -147,6 +147,17 @@ impl StateProofProvider self.state_provider.proof(state, address, slots) } + fn multiproof( + &self, + hashed_state: HashedPostState, + targets: HashMap>, + ) -> ProviderResult { + let bundle_state = self.block_execution_data_provider.execution_outcome().state(); + let mut state = HashedPostState::from_bundle_state(&bundle_state.state); + state.extend(hashed_state); + self.state_provider.multiproof(state, targets) + } + fn witness( &self, overlay: HashedPostState, diff --git a/crates/storage/provider/src/providers/state/historical.rs b/crates/storage/provider/src/providers/state/historical.rs index 19f758bc62fdc..10d828931e937 100644 --- a/crates/storage/provider/src/providers/state/historical.rs +++ b/crates/storage/provider/src/providers/state/historical.rs @@ -17,13 +17,16 @@ use reth_storage_api::{StateProofProvider, StorageRootProvider}; use reth_storage_errors::provider::ProviderResult; use reth_trie::{ prefix_set::TriePrefixSetsMut, proof::Proof, updates::TrieUpdates, witness::TrieWitness, - AccountProof, HashedPostState, HashedStorage, StateRoot, StorageRoot, + AccountProof, HashedPostState, HashedStorage, MultiProof, StateRoot, StorageRoot, }; use reth_trie_db::{ DatabaseHashedPostState, DatabaseHashedStorage, DatabaseProof, DatabaseStateRoot, DatabaseStorageRoot, DatabaseTrieWitness, }; -use std::{collections::HashMap, fmt::Debug}; +use std::{ + collections::{HashMap, HashSet}, + fmt::Debug, +}; /// State provider for a given block number which takes a tx reference. /// @@ -363,6 +366,17 @@ impl<'b, TX: DbTx> StateProofProvider for HistoricalStateProviderRef<'b, TX> { .map_err(Into::::into) } + fn multiproof( + &self, + hashed_state: HashedPostState, + targets: HashMap>, + ) -> ProviderResult { + let mut revert_state = self.revert_state()?; + revert_state.extend(hashed_state); + Proof::overlay_multiproof(self.tx, revert_state, targets) + .map_err(Into::::into) + } + fn witness( &self, overlay: HashedPostState, diff --git a/crates/storage/provider/src/providers/state/latest.rs b/crates/storage/provider/src/providers/state/latest.rs index 72fd5baac0a18..6247c39efb78d 100644 --- a/crates/storage/provider/src/providers/state/latest.rs +++ b/crates/storage/provider/src/providers/state/latest.rs @@ -137,6 +137,15 @@ impl<'b, TX: DbTx> StateProofProvider for LatestStateProviderRef<'b, TX> { .map_err(Into::::into) } + fn multiproof( + &self, + hashed_state: HashedPostState, + targets: HashMap>, + ) -> ProviderResult { + Proof::overlay_multiproof(self.tx, hashed_state, targets) + .map_err(Into::::into) + } + fn witness( &self, overlay: HashedPostState, diff --git a/crates/storage/provider/src/providers/state/macros.rs b/crates/storage/provider/src/providers/state/macros.rs index e3499c96ffe8a..727aabdbc47c8 100644 --- a/crates/storage/provider/src/providers/state/macros.rs +++ b/crates/storage/provider/src/providers/state/macros.rs @@ -52,6 +52,7 @@ macro_rules! delegate_provider_impls { } StateProofProvider $(where [$($generics)*])? { fn proof(&self, state: reth_trie::HashedPostState, address: reth_primitives::Address, slots: &[reth_primitives::B256]) -> reth_storage_errors::provider::ProviderResult; + fn multiproof(&self, state: reth_trie::HashedPostState, targets: std::collections::HashMap>) -> reth_storage_errors::provider::ProviderResult; fn witness(&self, state: reth_trie::HashedPostState, target: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult>; } ); diff --git a/crates/storage/provider/src/test_utils/mock.rs b/crates/storage/provider/src/test_utils/mock.rs index c539c0d2c72ea..329e8e3d4f996 100644 --- a/crates/storage/provider/src/test_utils/mock.rs +++ b/crates/storage/provider/src/test_utils/mock.rs @@ -26,11 +26,11 @@ use reth_storage_api::{ use reth_storage_errors::provider::{ConsistentViewError, ProviderError, ProviderResult}; use reth_trie::{ prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState, - HashedStorage, + HashedStorage, MultiProof, }; use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg}; use std::{ - collections::{BTreeMap, HashMap}, + collections::{BTreeMap, HashMap, HashSet}, ops::{RangeBounds, RangeInclusive}, sync::Arc, }; @@ -634,6 +634,14 @@ impl StateProofProvider for MockEthProvider { Ok(AccountProof::new(address)) } + fn multiproof( + &self, + hashed_state: HashedPostState, + targets: HashMap>, + ) -> ProviderResult { + Ok(MultiProof::default()) + } + fn witness( &self, _overlay: HashedPostState, diff --git a/crates/storage/provider/src/test_utils/noop.rs b/crates/storage/provider/src/test_utils/noop.rs index 338dc6fca17d7..c438e1b4ba903 100644 --- a/crates/storage/provider/src/test_utils/noop.rs +++ b/crates/storage/provider/src/test_utils/noop.rs @@ -1,5 +1,5 @@ use std::{ - collections::HashMap, + collections::{HashMap, HashSet}, ops::{RangeBounds, RangeInclusive}, sync::Arc, }; @@ -24,7 +24,7 @@ use reth_storage_api::{StateProofProvider, StorageRootProvider}; use reth_storage_errors::provider::ProviderResult; use reth_trie::{ prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState, - HashedStorage, + HashedStorage, MultiProof, }; use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg}; use tokio::sync::{broadcast, watch}; @@ -372,6 +372,14 @@ impl StateProofProvider for NoopProvider { Ok(AccountProof::new(address)) } + fn multiproof( + &self, + hashed_state: HashedPostState, + targets: HashMap>, + ) -> ProviderResult { + Ok(MultiProof::default()) + } + fn witness( &self, _overlay: HashedPostState, diff --git a/crates/storage/storage-api/src/trie.rs b/crates/storage/storage-api/src/trie.rs index d4a8c6aab571f..58f352a1332aa 100644 --- a/crates/storage/storage-api/src/trie.rs +++ b/crates/storage/storage-api/src/trie.rs @@ -2,9 +2,9 @@ use alloy_primitives::{Address, Bytes, B256}; use reth_storage_errors::provider::ProviderResult; use reth_trie::{ prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState, - HashedStorage, + HashedStorage, MultiProof, }; -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap, HashSet}; /// A type that can compute the state root of a given post state. #[auto_impl::auto_impl(&, Box, Arc)] @@ -66,6 +66,12 @@ pub trait StateProofProvider: Send + Sync { slots: &[B256], ) -> ProviderResult; + fn multiproof( + &self, + hashed_state: HashedPostState, + targets: HashMap>, + ) -> ProviderResult; + /// Get trie witness for provided state. fn witness( &self, diff --git a/crates/trie/common/Cargo.toml b/crates/trie/common/Cargo.toml index 0bd28140f4471..07c2229bdfe64 100644 --- a/crates/trie/common/Cargo.toml +++ b/crates/trie/common/Cargo.toml @@ -17,7 +17,7 @@ reth-codecs.workspace = true alloy-primitives.workspace = true alloy-rlp = { workspace = true, features = ["arrayvec"] } -alloy-trie = { workspace = true, features = ["serde"] } +alloy-trie = { workspace = true, features = ["std", "serde"] } alloy-consensus.workspace = true alloy-genesis.workspace = true revm-primitives.workspace = true diff --git a/crates/trie/common/src/proofs.rs b/crates/trie/common/src/proofs.rs index 87f9b4e4b6c22..fcc3d1622127d 100644 --- a/crates/trie/common/src/proofs.rs +++ b/crates/trie/common/src/proofs.rs @@ -10,7 +10,7 @@ use alloy_trie::{ }; use reth_primitives_traits::{constants::KECCAK_EMPTY, Account}; use serde::{Deserialize, Serialize}; -use std::collections::{BTreeMap, HashMap}; +use std::collections::{hash_map, BTreeMap, HashMap}; /// The state multiproof of target accounts and multiproofs of their storage tries. #[derive(Clone, Default, Debug)] @@ -72,6 +72,21 @@ impl MultiProof { } Ok(AccountProof { address, info, proof, storage_root, storage_proofs }) } + + pub fn extend(&mut self, other: Self) { + self.account_subtree.extend(other.account_subtree); + for (key, storage) in other.storages { + match self.storages.entry(key) { + hash_map::Entry::Occupied(mut e) => { + assert_eq!(e.get().root, storage.root); + e.get_mut().subtree.extend(storage.subtree); + } + hash_map::Entry::Vacant(e) => { + e.insert(storage); + } + } + } + } } /// The merkle multiproof of storage trie. diff --git a/crates/trie/db/src/proof.rs b/crates/trie/db/src/proof.rs index 8739dfc51880e..1f6864b712854 100644 --- a/crates/trie/db/src/proof.rs +++ b/crates/trie/db/src/proof.rs @@ -1,8 +1,12 @@ +use std::collections::{HashMap, HashSet}; + use crate::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory}; use reth_db_api::transaction::DbTx; use reth_execution_errors::StateProofError; -use reth_primitives::{Address, B256}; -use reth_trie::{hashed_cursor::HashedPostStateCursorFactory, proof::Proof, HashedPostState}; +use reth_primitives::{keccak256, Address, B256}; +use reth_trie::{ + hashed_cursor::HashedPostStateCursorFactory, proof::Proof, HashedPostState, MultiProof, +}; use reth_trie_common::AccountProof; /// Extends [`Proof`] with operations specific for working with a database transaction. @@ -17,6 +21,12 @@ pub trait DatabaseProof<'a, TX> { address: Address, slots: &[B256], ) -> Result; + + fn overlay_multiproof( + tx: &'a TX, + post_state: HashedPostState, + targets: HashMap>, + ) -> Result; } impl<'a, TX: DbTx> DatabaseProof<'a, TX> @@ -42,4 +52,27 @@ impl<'a, TX: DbTx> DatabaseProof<'a, TX> .with_prefix_sets_mut(prefix_sets) .account_proof(address, slots) } + + fn overlay_multiproof( + tx: &'a TX, + post_state: HashedPostState, + targets: HashMap>, + ) -> Result { + let prefix_sets = post_state.construct_prefix_sets(); + let sorted = post_state.into_sorted(); + let hashed_cursor_factory = + HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &sorted); + Self::from_tx(tx) + .with_hashed_cursor_factory(hashed_cursor_factory) + .with_prefix_sets_mut(prefix_sets) + .with_targets( + targets + .into_iter() + .map(|(address, slots)| { + (keccak256(address), slots.into_iter().map(keccak256).collect()) + }) + .collect(), + ) + .multiproof() + } } diff --git a/crates/trie/trie/src/witness.rs b/crates/trie/trie/src/witness.rs index fe6db2754d7e4..b8ab89361c490 100644 --- a/crates/trie/trie/src/witness.rs +++ b/crates/trie/trie/src/witness.rs @@ -1,6 +1,6 @@ use crate::{ hashed_cursor::HashedCursorFactory, prefix_set::TriePrefixSetsMut, proof::Proof, - trie_cursor::TrieCursorFactory, HashedPostState, + trie_cursor::TrieCursorFactory, updates::TrieUpdates, HashedPostState, }; use alloy_primitives::{keccak256, Bytes, B256}; use alloy_rlp::{BufMut, Decodable, Encodable}; @@ -8,7 +8,7 @@ use itertools::Either; use reth_execution_errors::{StateProofError, TrieWitnessError}; use reth_primitives::constants::EMPTY_ROOT_HASH; use reth_trie_common::{ - BranchNode, HashBuilder, Nibbles, TrieAccount, TrieNode, CHILD_INDEX_RANGE, + BranchNode, BranchNodeCompact, HashBuilder, Nibbles, TrieAccount, TrieNode, CHILD_INDEX_RANGE, }; use std::collections::{BTreeMap, HashMap, HashSet}; @@ -123,29 +123,32 @@ where )?); } - let storage_root = Self::next_root_from_proofs(storage_trie_nodes, |key: Nibbles| { - // Right pad the target with 0s. - let mut padded_key = key.pack(); - padded_key.resize(32, 0); - let target = (hashed_address, Vec::from([B256::from_slice(&padded_key)])); - let mut proof = Proof::new( - self.trie_cursor_factory.clone(), - self.hashed_cursor_factory.clone(), - ) - .with_prefix_sets_mut(self.prefix_sets.clone()) - .with_targets(HashMap::from([target])) - .storage_multiproof(hashed_address)?; + let (storage_root, _) = + Self::next_root_from_proofs(storage_trie_nodes, false, |key: Nibbles| { + // Right pad the target with 0s. + let mut padded_key = key.pack(); + padded_key.resize(32, 0); + let target = (hashed_address, Vec::from([B256::from_slice(&padded_key)])); + let mut proof = Proof::new( + self.trie_cursor_factory.clone(), + self.hashed_cursor_factory.clone(), + ) + .with_prefix_sets_mut(self.prefix_sets.clone()) + .with_targets(HashMap::from([target])) + .storage_multiproof(hashed_address)?; - // The subtree only contains the proof for a single target. - let node = - proof.subtree.remove(&key).ok_or(TrieWitnessError::MissingTargetNode(key))?; - self.witness.insert(keccak256(node.as_ref()), node.clone()); // record in witness - Ok(node) - })?; + // The subtree only contains the proof for a single target. + let node = proof + .subtree + .remove(&key) + .ok_or(TrieWitnessError::MissingTargetNode(key))?; + self.witness.insert(keccak256(node.as_ref()), node.clone()); // record in witness + Ok(node) + })?; debug_assert_eq!(storage_multiproof.root, storage_root); } - Self::next_root_from_proofs(account_trie_nodes, |key: Nibbles| { + Self::next_root_from_proofs(account_trie_nodes, false, |key: Nibbles| { // Right pad the target with 0s. let mut padded_key = key.pack(); padded_key.resize(32, 0); @@ -210,10 +213,12 @@ where Ok(trie_nodes) } - fn next_root_from_proofs( + pub fn next_root_from_proofs( trie_nodes: BTreeMap>>, + retain_updates: bool, mut trie_node_provider: impl FnMut(Nibbles) -> Result, - ) -> Result { + ) -> Result<(B256, revm::primitives::HashMap), TrieWitnessError> + { // Ignore branch child hashes in the path of leaves or lower child hashes. let mut keys = trie_nodes.keys().peekable(); let mut ignored = HashSet::::default(); @@ -223,7 +228,7 @@ where } } - let mut hash_builder = HashBuilder::default(); + let mut hash_builder = HashBuilder::default().with_updates(retain_updates); let mut trie_nodes = trie_nodes.into_iter().filter(|e| !ignored.contains(&e.0)).peekable(); while let Some((path, value)) = trie_nodes.next() { match value { @@ -266,7 +271,10 @@ where } } } - Ok(hash_builder.root()) + + let root = hash_builder.root(); + let (_, updates) = hash_builder.split(); + Ok((root, updates)) } }