diff --git a/crates/chain-state/src/in_memory.rs b/crates/chain-state/src/in_memory.rs index a6c85538107..cf1ddf413bf 100644 --- a/crates/chain-state/src/in_memory.rs +++ b/crates/chain-state/src/in_memory.rs @@ -17,7 +17,7 @@ use reth_primitives_traits::{ SignedTransaction, }; use reth_storage_api::StateProviderBox; -use reth_trie::{updates::TrieUpdates, HashedPostState}; +use reth_trie::{updates::TrieUpdatesSorted, HashedPostStateSorted}; use std::{collections::BTreeMap, sync::Arc, time::Instant}; use tokio::sync::{broadcast, watch}; @@ -725,10 +725,10 @@ pub struct ExecutedBlock { pub recovered_block: Arc>, /// Block's execution outcome. pub execution_output: Arc>, - /// Block's hashed state. - pub hashed_state: Arc, - /// Trie updates that result from calculating the state root for the block. - pub trie_updates: Arc, + /// Block's sorted hashed state. + pub hashed_state: Arc, + /// Sorted trie updates that result from calculating the state root for the block. + pub trie_updates: Arc, } impl Default for ExecutedBlock { @@ -763,13 +763,13 @@ impl ExecutedBlock { /// Returns a reference to the hashed state result of the execution outcome #[inline] - pub fn hashed_state(&self) -> &HashedPostState { + pub fn hashed_state(&self) -> &HashedPostStateSorted { &self.hashed_state } /// Returns a reference to the trie updates resulting from the execution outcome #[inline] - pub fn trie_updates(&self) -> &TrieUpdates { + pub fn trie_updates(&self) -> &TrieUpdatesSorted { &self.trie_updates } @@ -875,8 +875,8 @@ mod tests { StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider, }; use reth_trie::{ - AccountProof, HashedStorage, MultiProof, MultiProofTargets, StorageMultiProof, - StorageProof, TrieInput, + updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, + MultiProofTargets, StorageMultiProof, StorageProof, TrieInput, }; fn create_mock_state( diff --git a/crates/chain-state/src/memory_overlay.rs b/crates/chain-state/src/memory_overlay.rs index 254edb248b4..b248bc2e731 100644 --- a/crates/chain-state/src/memory_overlay.rs +++ b/crates/chain-state/src/memory_overlay.rs @@ -53,7 +53,7 @@ impl<'a, N: NodePrimitives> MemoryOverlayStateProviderRef<'a, N> { /// Return lazy-loaded trie state aggregated from in-memory blocks. fn trie_input(&self) -> &TrieInput { self.trie_input.get_or_init(|| { - TrieInput::from_blocks( + TrieInput::from_blocks_sorted( self.in_memory .iter() .rev() diff --git a/crates/chain-state/src/test_utils.rs b/crates/chain-state/src/test_utils.rs index 5d318aca56c..867f59607f1 100644 --- a/crates/chain-state/src/test_utils.rs +++ b/crates/chain-state/src/test_utils.rs @@ -23,7 +23,7 @@ use reth_primitives_traits::{ SignedTransaction, }; use reth_storage_api::NodePrimitivesProvider; -use reth_trie::{root::state_root_unhashed, updates::TrieUpdates, HashedPostState}; +use reth_trie::{root::state_root_unhashed, updates::TrieUpdatesSorted, HashedPostStateSorted}; use revm_database::BundleState; use revm_state::AccountInfo; use std::{ @@ -216,8 +216,8 @@ impl TestBlockBuilder { block_number, vec![Requests::default()], )), - hashed_state: Arc::new(HashedPostState::default()), - trie_updates: Arc::new(TrieUpdates::default()), + hashed_state: Arc::new(HashedPostStateSorted::default()), + trie_updates: Arc::new(TrieUpdatesSorted::default()), } } diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs index 5e2ed1c513c..cbdb8a9494b 100644 --- a/crates/engine/tree/src/tree/mod.rs +++ b/crates/engine/tree/src/tree/mod.rs @@ -1759,8 +1759,8 @@ where Ok(Some(ExecutedBlock { recovered_block: Arc::new(RecoveredBlock::new_sealed(block, senders)), execution_output: Arc::new(execution_output), - hashed_state: Arc::new(hashed_state), - trie_updates: Arc::new(trie_updates.into()), + hashed_state: Arc::new(hashed_state.into_sorted()), + trie_updates: Arc::new(trie_updates), })) } diff --git a/crates/engine/tree/src/tree/payload_processor/multiproof.rs b/crates/engine/tree/src/tree/payload_processor/multiproof.rs index ca3bd380d4d..3ff97ef9b38 100644 --- a/crates/engine/tree/src/tree/payload_processor/multiproof.rs +++ b/crates/engine/tree/src/tree/payload_processor/multiproof.rs @@ -13,9 +13,8 @@ use metrics::{Gauge, Histogram}; use reth_metrics::Metrics; use reth_revm::state::EvmState; use reth_trie::{ - added_removed_keys::MultiAddedRemovedKeys, prefix_set::TriePrefixSetsMut, - updates::TrieUpdatesSorted, DecodedMultiProof, HashedPostState, HashedPostStateSorted, - HashedStorage, MultiProofTargets, TrieInput, + added_removed_keys::MultiAddedRemovedKeys, DecodedMultiProof, HashedPostState, HashedStorage, + MultiProofTargets, }; use reth_trie_parallel::{ proof::ParallelProof, @@ -56,35 +55,6 @@ impl SparseTrieUpdate { } } -/// Common configuration for multi proof tasks -#[derive(Debug, Clone, Default)] -pub(crate) struct MultiProofConfig { - /// The sorted collection of cached in-memory intermediate trie nodes that - /// can be reused for computation. - pub nodes_sorted: Arc, - /// The sorted in-memory overlay hashed state. - pub state_sorted: Arc, - /// The collection of prefix sets for the computation. Since the prefix sets _always_ - /// invalidate the in-memory nodes, not all keys from `state_sorted` might be present here, - /// if we have cached nodes for them. - pub prefix_sets: Arc, -} - -impl MultiProofConfig { - /// Creates a new state root config from the trie input. - /// - /// This returns a cleared [`TrieInput`] so that we can reuse any allocated space in the - /// [`TrieInput`]. - pub(crate) fn from_input(mut input: TrieInput) -> (TrieInput, Self) { - let config = Self { - nodes_sorted: Arc::new(input.nodes.drain_into_sorted()), - state_sorted: Arc::new(input.state.drain_into_sorted()), - prefix_sets: Arc::new(input.prefix_sets.clone()), - }; - (input.cleared(), config) - } -} - /// Messages used internally by the multi proof task. #[derive(Debug)] pub(super) enum MultiProofMessage { diff --git a/crates/engine/tree/src/tree/payload_validator.rs b/crates/engine/tree/src/tree/payload_validator.rs index fdd6b30a6e8..74ca6e47a6d 100644 --- a/crates/engine/tree/src/tree/payload_validator.rs +++ b/crates/engine/tree/src/tree/payload_validator.rs @@ -5,7 +5,7 @@ use crate::tree::{ error::{InsertBlockError, InsertBlockErrorKind, InsertPayloadError}, executor::WorkloadExecutor, instrumented_state::InstrumentedStateProvider, - payload_processor::{multiproof::MultiProofConfig, PayloadProcessor}, + payload_processor::PayloadProcessor, precompile_cache::{CachedPrecompile, CachedPrecompileMetrics, PrecompileCacheMap}, sparse_trie::StateRootComputeOutcome, EngineApiMetrics, EngineApiTreeState, ExecutionEnv, PayloadHandle, StateProviderBuilder, @@ -38,7 +38,7 @@ use reth_provider::{ StateRootProvider, TrieReader, }; use reth_revm::db::State; -use reth_trie::{updates::TrieUpdates, HashedPostState, TrieInput}; +use reth_trie::{updates::TrieUpdates, HashedPostState, TrieInputSorted}; use reth_trie_parallel::root::{ParallelStateRoot, ParallelStateRootError}; use std::{collections::HashMap, sync::Arc, time::Instant}; use tracing::{debug, debug_span, error, info, instrument, trace, warn}; @@ -121,8 +121,6 @@ where metrics: EngineApiMetrics, /// Validator for the payload. validator: V, - /// A cleared trie input, kept around to be reused so allocations can be minimized. - trie_input: Option, } impl BasicEngineValidator @@ -166,7 +164,6 @@ where invalid_block_hook, metrics: EngineApiMetrics::default(), validator, - trie_input: Default::default(), } } @@ -531,8 +528,8 @@ where Ok(ExecutedBlock { recovered_block: Arc::new(block), execution_output: Arc::new(ExecutionOutcome::from((output, block_num_hash.number))), - hashed_state: Arc::new(hashed_state), - trie_updates: Arc::new(trie_output), + hashed_state: Arc::new(hashed_state.into_sorted()), + trie_updates: Arc::new(trie_output.into_sorted()), }) } @@ -642,26 +639,24 @@ where hashed_state: &HashedPostState, state: &EngineApiTreeState, ) -> Result<(B256, TrieUpdates), ParallelStateRootError> { - let (mut input, block_hash) = self.compute_trie_input(parent_hash, state, None)?; + let (mut input, block_hash) = self.compute_trie_input(parent_hash, state)?; - // Extend with block we are validating root for. - input.append_ref(hashed_state); + // Extend state overlay with current block's sorted state. + input.prefix_sets.extend(hashed_state.construct_prefix_sets()); + let sorted_hashed_state = hashed_state.clone().into_sorted(); + Arc::make_mut(&mut input.state).extend_ref(&sorted_hashed_state); - // Convert the TrieInput into a MultProofConfig, since everything uses the sorted - // forms of the state/trie fields. - let (_, multiproof_config) = MultiProofConfig::from_input(input); + let TrieInputSorted { nodes, state, prefix_sets: prefix_sets_mut } = input; let factory = OverlayStateProviderFactory::new(self.provider.clone()) .with_block_hash(Some(block_hash)) - .with_trie_overlay(Some(multiproof_config.nodes_sorted)) - .with_hashed_state_overlay(Some(multiproof_config.state_sorted)); + .with_trie_overlay(Some(nodes)) + .with_hashed_state_overlay(Some(state)); // The `hashed_state` argument is already taken into account as part of the overlay, but we // need to use the prefix sets which were generated from it to indicate to the // ParallelStateRoot which parts of the trie need to be recomputed. - let prefix_sets = Arc::into_inner(multiproof_config.prefix_sets) - .expect("MultiProofConfig was never cloned") - .freeze(); + let prefix_sets = prefix_sets_mut.freeze(); ParallelStateRoot::new(factory, prefix_sets).incremental_root_with_updates() } @@ -760,31 +755,23 @@ where > { match strategy { StateRootStrategy::StateRootTask => { - // get allocated trie input if it exists - let allocated_trie_input = self.trie_input.take(); - // Compute trie input let trie_input_start = Instant::now(); - let (trie_input, block_hash) = - self.compute_trie_input(parent_hash, state, allocated_trie_input)?; + let (trie_input, block_hash) = self.compute_trie_input(parent_hash, state)?; self.metrics .block_validation .trie_input_duration .record(trie_input_start.elapsed().as_secs_f64()); - // Convert the TrieInput into a MultProofConfig, since everything uses the sorted - // forms of the state/trie fields. - let (trie_input, multiproof_config) = MultiProofConfig::from_input(trie_input); - self.trie_input.replace(trie_input); + // Create OverlayStateProviderFactory with sorted trie data for multiproofs + let TrieInputSorted { nodes, state, .. } = trie_input; - // Create OverlayStateProviderFactory with the multiproof config, for use with - // multiproofs. let multiproof_provider_factory = OverlayStateProviderFactory::new(self.provider.clone()) .with_block_hash(Some(block_hash)) - .with_trie_overlay(Some(multiproof_config.nodes_sorted)) - .with_hashed_state_overlay(Some(multiproof_config.state_sorted)); + .with_trie_overlay(Some(nodes)) + .with_hashed_state_overlay(Some(state)); // Use state root task only if prefix sets are empty, otherwise proof generation is // too expensive because it requires walking all paths in every proof. @@ -912,14 +899,14 @@ where /// Computes the trie input at the provided parent hash, as well as the block number of the /// highest persisted ancestor. /// - /// The goal of this function is to take in-memory blocks and generate a [`TrieInput`] that - /// serves as an overlay to the database blocks. + /// The goal of this function is to take in-memory blocks and generate a [`TrieInputSorted`] + /// that serves as an overlay to the database blocks. /// /// It works as follows: /// 1. Collect in-memory blocks that are descendants of the provided parent hash using /// [`crate::tree::TreeState::blocks_by_hash`]. This returns the highest persisted ancestor /// hash (`block_hash`) and the list of in-memory descendant blocks. - /// 2. Extend the `TrieInput` with the contents of these in-memory blocks (from oldest to + /// 2. Extend the `TrieInputSorted` with the contents of these in-memory blocks (from oldest to /// newest) to build the overlay state and trie updates that sit on top of the database view /// anchored at `block_hash`. #[instrument( @@ -932,11 +919,7 @@ where &self, parent_hash: B256, state: &EngineApiTreeState, - allocated_trie_input: Option, - ) -> ProviderResult<(TrieInput, B256)> { - // get allocated trie input or use a default trie input - let mut input = allocated_trie_input.unwrap_or_default(); - + ) -> ProviderResult<(TrieInputSorted, B256)> { let (block_hash, blocks) = state.tree_state.blocks_by_hash(parent_hash).unwrap_or_else(|| (parent_hash, vec![])); @@ -946,10 +929,24 @@ where debug!(target: "engine::tree::payload_validator", historical = ?block_hash, blocks = blocks.len(), "Parent found in memory"); } - // Extend with contents of parent in-memory blocks. - input.extend_with_blocks( - blocks.iter().rev().map(|block| (block.hashed_state(), block.trie_updates())), - ); + // Extend with contents of parent in-memory blocks directly in sorted form. + let mut input = TrieInputSorted::default(); + let mut blocks_iter = blocks.iter().rev(); + + if let Some(first) = blocks_iter.next() { + input.state = Arc::clone(&first.hashed_state); + input.nodes = Arc::clone(&first.trie_updates); + + // Only clone and mutate if there are multiple in-memory blocks. + if blocks.len() > 1 { + let state_mut = Arc::make_mut(&mut input.state); + let nodes_mut = Arc::make_mut(&mut input.nodes); + for block in blocks_iter { + state_mut.extend_ref(block.hashed_state()); + nodes_mut.extend_ref(block.trie_updates()); + } + } + } Ok((input, block_hash)) } diff --git a/crates/engine/tree/src/tree/tests.rs b/crates/engine/tree/src/tree/tests.rs index 7fbae4cac5c..a8a74a0e26a 100644 --- a/crates/engine/tree/src/tree/tests.rs +++ b/crates/engine/tree/src/tree/tests.rs @@ -829,8 +829,8 @@ fn test_tree_state_on_new_head_deep_fork() { test_harness.tree.state.tree_state.insert_executed(ExecutedBlock { recovered_block: Arc::new(block.clone()), execution_output: Arc::new(ExecutionOutcome::default()), - hashed_state: Arc::new(HashedPostState::default()), - trie_updates: Arc::new(TrieUpdates::default()), + hashed_state: Arc::new(HashedPostState::default().into_sorted()), + trie_updates: Arc::new(TrieUpdates::default().into_sorted()), }); } test_harness.tree.state.tree_state.set_canonical_head(chain_a.last().unwrap().num_hash()); @@ -839,8 +839,8 @@ fn test_tree_state_on_new_head_deep_fork() { test_harness.tree.state.tree_state.insert_executed(ExecutedBlock { recovered_block: Arc::new(block.clone()), execution_output: Arc::new(ExecutionOutcome::default()), - hashed_state: Arc::new(HashedPostState::default()), - trie_updates: Arc::new(TrieUpdates::default()), + hashed_state: Arc::new(HashedPostState::default().into_sorted()), + trie_updates: Arc::new(TrieUpdates::default().into_sorted()), }); } diff --git a/crates/optimism/flashblocks/src/worker.rs b/crates/optimism/flashblocks/src/worker.rs index 8cf7777f6a6..a9a533fb00a 100644 --- a/crates/optimism/flashblocks/src/worker.rs +++ b/crates/optimism/flashblocks/src/worker.rs @@ -123,7 +123,7 @@ where ExecutedBlock { recovered_block: block.into(), execution_output: Arc::new(execution_outcome), - hashed_state: Arc::new(hashed_state), + hashed_state: Arc::new(hashed_state.into_sorted()), trie_updates: Arc::default(), }, ); diff --git a/crates/optimism/payload/src/builder.rs b/crates/optimism/payload/src/builder.rs index 05f33d3b699..ad51abf9e6a 100644 --- a/crates/optimism/payload/src/builder.rs +++ b/crates/optimism/payload/src/builder.rs @@ -389,8 +389,8 @@ impl OpBuilder<'_, Txs> { let executed: ExecutedBlock = ExecutedBlock { recovered_block: Arc::new(block), execution_output: Arc::new(execution_outcome), - hashed_state: Arc::new(hashed_state), - trie_updates: Arc::new(trie_updates), + hashed_state: Arc::new(hashed_state.into_sorted()), + trie_updates: Arc::new(trie_updates.into_sorted()), }; let no_tx_pool = ctx.attributes().no_tx_pool(); diff --git a/crates/ress/provider/src/lib.rs b/crates/ress/provider/src/lib.rs index d986eb9e953..ecfdb3c4d8a 100644 --- a/crates/ress/provider/src/lib.rs +++ b/crates/ress/provider/src/lib.rs @@ -156,11 +156,12 @@ where // NOTE: there might be a race condition where target ancestor hash gets evicted from the // database. let witness_state_provider = self.provider.state_by_block_hash(ancestor_hash)?; - let mut trie_input = TrieInput::default(); - for block in executed_ancestors.into_iter().rev() { - let trie_updates = block.trie_updates.as_ref(); - trie_input.append_cached_ref(trie_updates, &block.hashed_state); - } + let trie_input = TrieInput::from_blocks_sorted( + executed_ancestors + .iter() + .rev() + .map(|block| (block.hashed_state.as_ref(), block.trie_updates.as_ref())), + ); let mut hashed_state = db.into_state(); hashed_state.extend(record.hashed_state); diff --git a/crates/rpc/rpc-eth-api/src/helpers/pending_block.rs b/crates/rpc/rpc-eth-api/src/helpers/pending_block.rs index 1dda44d090e..6d6c147dce0 100644 --- a/crates/rpc/rpc-eth-api/src/helpers/pending_block.rs +++ b/crates/rpc/rpc-eth-api/src/helpers/pending_block.rs @@ -372,8 +372,8 @@ pub trait LoadPendingBlock: Ok(ExecutedBlock { recovered_block: block.into(), execution_output: Arc::new(execution_outcome), - hashed_state: Arc::new(hashed_state), - trie_updates: Arc::new(trie_updates), + hashed_state: Arc::new(hashed_state.into_sorted()), + trie_updates: Arc::new(trie_updates.into_sorted()), }) } } diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs index ece6ef56c85..e7fee73f374 100644 --- a/crates/storage/provider/src/providers/database/provider.rs +++ b/crates/storage/provider/src/providers/database/provider.rs @@ -296,12 +296,10 @@ impl DatabaseProvider HashedPostStateSorted { - let mut updated_accounts = Vec::new(); - let mut destroyed_accounts = HashSet::default(); - for (hashed_address, info) in self.accounts { - if let Some(info) = info { - updated_accounts.push((hashed_address, info)); - } else { - destroyed_accounts.insert(hashed_address); - } + /// Extend this hashed post state with sorted data, converting directly into the unsorted + /// `HashMap` representation. This is more efficient than first converting to `HashedPostState` + /// and then extending, as it avoids creating intermediate `HashMap` allocations. + pub fn extend_from_sorted(&mut self, sorted: &HashedPostStateSorted) { + // Reserve capacity for accounts + self.accounts + .reserve(sorted.accounts.accounts.len() + sorted.accounts.destroyed_accounts.len()); + + // Insert updated accounts + for (address, account) in &sorted.accounts.accounts { + self.accounts.insert(*address, Some(*account)); } - updated_accounts.sort_unstable_by_key(|(address, _)| *address); - let accounts = HashedAccountsSorted { accounts: updated_accounts, destroyed_accounts }; - let storages = self - .storages - .into_iter() - .map(|(hashed_address, storage)| (hashed_address, storage.into_sorted())) - .collect(); + // Insert destroyed accounts + for address in &sorted.accounts.destroyed_accounts { + self.accounts.insert(*address, None); + } - HashedPostStateSorted { accounts, storages } + // Reserve capacity for storages + self.storages.reserve(sorted.storages.len()); + + // Extend storages + for (hashed_address, sorted_storage) in &sorted.storages { + match self.storages.entry(*hashed_address) { + hash_map::Entry::Vacant(entry) => { + let mut new_storage = HashedStorage::new(false); + new_storage.extend_from_sorted(sorted_storage); + entry.insert(new_storage); + } + hash_map::Entry::Occupied(mut entry) => { + entry.get_mut().extend_from_sorted(sorted_storage); + } + } + } } - /// Converts hashed post state into [`HashedPostStateSorted`], but keeping the maps allocated by - /// draining. - /// - /// This effectively clears all the fields in the [`HashedPostStateSorted`]. - /// - /// This allows us to reuse the allocated space. This allocates new space for the sorted hashed - /// post state, like `into_sorted`. - pub fn drain_into_sorted(&mut self) -> HashedPostStateSorted { + /// Converts hashed post state into [`HashedPostStateSorted`]. + pub fn into_sorted(self) -> HashedPostStateSorted { let mut updated_accounts = Vec::new(); let mut destroyed_accounts = HashSet::default(); - for (hashed_address, info) in self.accounts.drain() { + for (hashed_address, info) in self.accounts { if let Some(info) = info { updated_accounts.push((hashed_address, info)); } else { @@ -367,7 +374,7 @@ impl HashedPostState { let storages = self .storages - .drain() + .into_iter() .map(|(hashed_address, storage)| (hashed_address, storage.into_sorted())) .collect(); @@ -441,6 +448,29 @@ impl HashedStorage { self.storage.extend(other.storage.iter().map(|(&k, &v)| (k, v))); } + /// Extend hashed storage with sorted data, converting directly into the unsorted `HashMap` + /// representation. This is more efficient than first converting to `HashedStorage` and + /// then extending, as it avoids creating intermediate `HashMap` allocations. + pub fn extend_from_sorted(&mut self, sorted: &HashedStorageSorted) { + if sorted.wiped { + self.wiped = true; + self.storage.clear(); + } + + // Reserve capacity for all slots + self.storage.reserve(sorted.non_zero_valued_slots.len() + sorted.zero_valued_slots.len()); + + // Insert non-zero valued slots + for (slot, value) in &sorted.non_zero_valued_slots { + self.storage.insert(*slot, *value); + } + + // Insert zero-valued slots + for slot in &sorted.zero_valued_slots { + self.storage.insert(*slot, U256::ZERO); + } + } + /// Converts hashed storage into [`HashedStorageSorted`]. pub fn into_sorted(self) -> HashedStorageSorted { let mut non_zero_valued_slots = Vec::new(); @@ -514,6 +544,13 @@ impl HashedPostStateSorted { .or_insert_with(|| other_storage.clone()); } } + + /// Clears all accounts and storage data. + pub fn clear(&mut self) { + self.accounts.accounts.clear(); + self.accounts.destroyed_accounts.clear(); + self.storages.clear(); + } } impl AsRef for HashedPostStateSorted { @@ -1246,4 +1283,88 @@ mod tests { assert_eq!(storage3.zero_valued_slots.len(), 1); assert!(storage3.zero_valued_slots.contains(&B256::from([4; 32]))); } + + /// Test extending with sorted accounts merges correctly into `HashMap` + #[test] + fn test_hashed_post_state_extend_from_sorted_with_accounts() { + let addr1 = B256::random(); + let addr2 = B256::random(); + + let mut state = HashedPostState::default(); + state.accounts.insert(addr1, Some(Default::default())); + + let mut sorted_state = HashedPostStateSorted::default(); + sorted_state.accounts.accounts.push((addr2, Default::default())); + + state.extend_from_sorted(&sorted_state); + + assert_eq!(state.accounts.len(), 2); + assert!(state.accounts.contains_key(&addr1)); + assert!(state.accounts.contains_key(&addr2)); + } + + /// Test destroyed accounts (None values) are inserted correctly + #[test] + fn test_hashed_post_state_extend_from_sorted_with_destroyed_accounts() { + let addr1 = B256::random(); + + let mut state = HashedPostState::default(); + + let mut sorted_state = HashedPostStateSorted::default(); + sorted_state.accounts.destroyed_accounts.insert(addr1); + + state.extend_from_sorted(&sorted_state); + + assert!(state.accounts.contains_key(&addr1)); + assert_eq!(state.accounts.get(&addr1), Some(&None)); + } + + /// Test non-wiped storage merges both zero and non-zero valued slots + #[test] + fn test_hashed_storage_extend_from_sorted_non_wiped() { + let slot1 = B256::random(); + let slot2 = B256::random(); + let slot3 = B256::random(); + + let mut storage = HashedStorage::from_iter(false, [(slot1, U256::from(100))]); + + let mut zero_valued_slots = B256Set::default(); + zero_valued_slots.insert(slot3); + + let sorted = HashedStorageSorted { + non_zero_valued_slots: vec![(slot2, U256::from(200))], + zero_valued_slots, + wiped: false, + }; + + storage.extend_from_sorted(&sorted); + + assert!(!storage.wiped); + assert_eq!(storage.storage.len(), 3); + assert_eq!(storage.storage.get(&slot1), Some(&U256::from(100))); + assert_eq!(storage.storage.get(&slot2), Some(&U256::from(200))); + assert_eq!(storage.storage.get(&slot3), Some(&U256::ZERO)); + } + + /// Test wiped=true clears existing storage and only keeps new slots (critical edge case) + #[test] + fn test_hashed_storage_extend_from_sorted_wiped() { + let slot1 = B256::random(); + let slot2 = B256::random(); + + let mut storage = HashedStorage::from_iter(false, [(slot1, U256::from(100))]); + + let sorted = HashedStorageSorted { + non_zero_valued_slots: vec![(slot2, U256::from(200))], + zero_valued_slots: Default::default(), + wiped: true, + }; + + storage.extend_from_sorted(&sorted); + + assert!(storage.wiped); + // After wipe, old storage should be cleared and only new storage remains + assert_eq!(storage.storage.len(), 1); + assert_eq!(storage.storage.get(&slot2), Some(&U256::from(200))); + } } diff --git a/crates/trie/common/src/input.rs b/crates/trie/common/src/input.rs index 522cfa9ed41..89f7811a3a2 100644 --- a/crates/trie/common/src/input.rs +++ b/crates/trie/common/src/input.rs @@ -1,4 +1,9 @@ -use crate::{prefix_set::TriePrefixSetsMut, updates::TrieUpdates, HashedPostState}; +use crate::{ + prefix_set::TriePrefixSetsMut, + updates::{TrieUpdates, TrieUpdatesSorted}, + HashedPostState, HashedPostStateSorted, +}; +use alloc::sync::Arc; /// Inputs for trie-related computations. #[derive(Default, Debug, Clone)] @@ -41,6 +46,20 @@ impl TrieInput { input } + /// Create new trie input from the provided sorted blocks, from oldest to newest. + /// Converts sorted types to unsorted for aggregation. + pub fn from_blocks_sorted<'a>( + blocks: impl IntoIterator, + ) -> Self { + let mut input = Self::default(); + for (hashed_state, trie_updates) in blocks { + // Extend directly from sorted types, avoiding intermediate HashMap allocations + input.nodes.extend_from_sorted(trie_updates); + input.state.extend_from_sorted(hashed_state); + } + input + } + /// Extend the trie input with the provided blocks, from oldest to newest. /// /// For blocks with missing trie updates, the trie input will be extended with prefix sets @@ -119,3 +138,40 @@ impl TrieInput { self } } + +/// Sorted variant of [`TrieInput`] for efficient proof generation. +/// +/// This type holds sorted versions of trie data structures, which eliminates the need +/// for expensive sorting operations during multiproof generation. +#[derive(Default, Debug, Clone)] +pub struct TrieInputSorted { + /// Sorted cached in-memory intermediate trie nodes. + pub nodes: Arc, + /// Sorted in-memory overlay hashed state. + pub state: Arc, + /// Prefix sets for computation. + pub prefix_sets: TriePrefixSetsMut, +} + +impl TrieInputSorted { + /// Create new sorted trie input. + pub const fn new( + nodes: Arc, + state: Arc, + prefix_sets: TriePrefixSetsMut, + ) -> Self { + Self { nodes, state, prefix_sets } + } + + /// Create from unsorted [`TrieInput`] by sorting. + pub fn from_unsorted(input: TrieInput) -> Self { + Self { + nodes: Arc::new(input.nodes.into_sorted()), + state: Arc::new(input.state.into_sorted()), + prefix_sets: input.prefix_sets, + } + } +} + +#[cfg(test)] +mod tests {} diff --git a/crates/trie/common/src/lib.rs b/crates/trie/common/src/lib.rs index e4292a52016..6f7453c0105 100644 --- a/crates/trie/common/src/lib.rs +++ b/crates/trie/common/src/lib.rs @@ -17,7 +17,7 @@ pub use hashed_state::*; /// Input for trie computation. mod input; -pub use input::TrieInput; +pub use input::{TrieInput, TrieInputSorted}; /// The implementation of hash builder. pub mod hash_builder; diff --git a/crates/trie/common/src/updates.rs b/crates/trie/common/src/updates.rs index b0d178cd1d0..7432b710352 100644 --- a/crates/trie/common/src/updates.rs +++ b/crates/trie/common/src/updates.rs @@ -73,6 +73,48 @@ impl TrieUpdates { self.account_nodes.retain(|nibbles, _| !other.removed_nodes.contains(nibbles)); } + /// Extend trie updates with sorted data, converting directly into the unsorted `HashMap` + /// representation. This is more efficient than first converting to `TrieUpdates` and + /// then extending, as it avoids creating intermediate `HashMap` allocations. + /// + /// This top-level helper merges account nodes and delegates each account's storage trie to + /// [`StorageTrieUpdates::extend_from_sorted`]. + pub fn extend_from_sorted(&mut self, sorted: &TrieUpdatesSorted) { + // Reserve capacity for account nodes + let new_nodes_count = sorted + .account_nodes + .iter() + .filter(|(nibbles, node)| !nibbles.is_empty() && node.is_some()) + .count(); + self.account_nodes.reserve(new_nodes_count); + + // Insert account nodes from sorted (only non-None entries) + for (nibbles, maybe_node) in &sorted.account_nodes { + if nibbles.is_empty() { + continue; + } + match maybe_node { + Some(node) => { + self.removed_nodes.remove(nibbles); + self.account_nodes.insert(*nibbles, node.clone()); + } + None => { + self.account_nodes.remove(nibbles); + self.removed_nodes.insert(*nibbles); + } + } + } + + // Extend storage tries + self.storage_tries.reserve(sorted.storage_tries.len()); + for (hashed_address, sorted_storage) in &sorted.storage_tries { + self.storage_tries + .entry(*hashed_address) + .or_default() + .extend_from_sorted(sorted_storage); + } + } + /// Insert storage updates for a given hashed address. pub fn insert_storage_updates( &mut self, @@ -108,17 +150,6 @@ impl TrieUpdates { /// Converts trie updates into [`TrieUpdatesSorted`]. pub fn into_sorted(mut self) -> TrieUpdatesSorted { - self.drain_into_sorted() - } - - /// Converts trie updates into [`TrieUpdatesSorted`], but keeping the maps allocated by - /// draining. - /// - /// This effectively clears all the fields in the [`TrieUpdatesSorted`]. - /// - /// This allows us to reuse the allocated space. This allocates new space for the sorted - /// updates, like `into_sorted`. - pub fn drain_into_sorted(&mut self) -> TrieUpdatesSorted { let mut account_nodes = self .account_nodes .drain() @@ -253,6 +284,42 @@ impl StorageTrieUpdates { self.storage_nodes.retain(|nibbles, _| !other.removed_nodes.contains(nibbles)); } + /// Extend storage trie updates with sorted data, converting directly into the unsorted + /// `HashMap` representation. This is more efficient than first converting to + /// `StorageTrieUpdates` and then extending, as it avoids creating intermediate `HashMap` + /// allocations. + /// + /// This is invoked from [`TrieUpdates::extend_from_sorted`] for each account. + pub fn extend_from_sorted(&mut self, sorted: &StorageTrieUpdatesSorted) { + if sorted.is_deleted { + self.storage_nodes.clear(); + self.removed_nodes.clear(); + } + self.is_deleted |= sorted.is_deleted; + + // Reserve capacity for storage nodes + let new_nodes_count = sorted + .storage_nodes + .iter() + .filter(|(nibbles, node)| !nibbles.is_empty() && node.is_some()) + .count(); + self.storage_nodes.reserve(new_nodes_count); + + // Remove nodes marked as removed and insert new nodes + for (nibbles, maybe_node) in &sorted.storage_nodes { + if nibbles.is_empty() { + continue; + } + if let Some(node) = maybe_node { + self.removed_nodes.remove(nibbles); + self.storage_nodes.insert(*nibbles, node.clone()); + } else { + self.storage_nodes.remove(nibbles); + self.removed_nodes.insert(*nibbles); + } + } + } + /// Finalize storage trie updates for by taking updates from walker and hash builder. pub fn finalize(&mut self, hash_builder: HashBuilder, removed_keys: HashSet) { // Retrieve updated nodes from hash builder. @@ -499,6 +566,12 @@ impl TrieUpdatesSorted { .or_insert_with(|| storage_trie.clone()); } } + + /// Clears all account nodes and storage tries. + pub fn clear(&mut self) { + self.account_nodes.clear(); + self.storage_tries.clear(); + } } impl AsRef for TrieUpdatesSorted { @@ -749,6 +822,151 @@ mod tests { assert_eq!(storage3.storage_nodes[0].0, Nibbles::from_nibbles_unchecked([0x06])); assert_eq!(storage3.storage_nodes[1].0, Nibbles::from_nibbles_unchecked([0x07])); } + + /// Test extending with storage tries adds both nodes and removed nodes correctly + #[test] + fn test_trie_updates_extend_from_sorted_with_storage_tries() { + let hashed_address = B256::from([1; 32]); + + let mut updates = TrieUpdates::default(); + + let storage_trie = StorageTrieUpdatesSorted { + is_deleted: false, + storage_nodes: vec![ + (Nibbles::from_nibbles_unchecked([0x0a]), Some(BranchNodeCompact::default())), + (Nibbles::from_nibbles_unchecked([0x0b]), None), + ], + }; + + let sorted = TrieUpdatesSorted { + account_nodes: vec![], + storage_tries: B256Map::from_iter([(hashed_address, storage_trie)]), + }; + + updates.extend_from_sorted(&sorted); + + assert_eq!(updates.storage_tries.len(), 1); + let storage = updates.storage_tries.get(&hashed_address).unwrap(); + assert!(!storage.is_deleted); + assert_eq!(storage.storage_nodes.len(), 1); + assert!(storage.removed_nodes.contains(&Nibbles::from_nibbles_unchecked([0x0b]))); + } + + /// Test deleted=true clears old storage nodes before adding new ones (critical edge case) + #[test] + fn test_trie_updates_extend_from_sorted_with_deleted_storage() { + let hashed_address = B256::from([1; 32]); + + let mut updates = TrieUpdates::default(); + updates.storage_tries.insert( + hashed_address, + StorageTrieUpdates { + is_deleted: false, + storage_nodes: HashMap::from_iter([( + Nibbles::from_nibbles_unchecked([0x01]), + BranchNodeCompact::default(), + )]), + removed_nodes: Default::default(), + }, + ); + + let storage_trie = StorageTrieUpdatesSorted { + is_deleted: true, + storage_nodes: vec![( + Nibbles::from_nibbles_unchecked([0x0a]), + Some(BranchNodeCompact::default()), + )], + }; + + let sorted = TrieUpdatesSorted { + account_nodes: vec![], + storage_tries: B256Map::from_iter([(hashed_address, storage_trie)]), + }; + + updates.extend_from_sorted(&sorted); + + let storage = updates.storage_tries.get(&hashed_address).unwrap(); + assert!(storage.is_deleted); + // After deletion, old nodes should be cleared + assert_eq!(storage.storage_nodes.len(), 1); + assert!(storage.storage_nodes.contains_key(&Nibbles::from_nibbles_unchecked([0x0a]))); + } + + /// Test non-deleted storage merges nodes and tracks removed nodes + #[test] + fn test_storage_trie_updates_extend_from_sorted_non_deleted() { + let mut storage = StorageTrieUpdates { + is_deleted: false, + storage_nodes: HashMap::from_iter([( + Nibbles::from_nibbles_unchecked([0x01]), + BranchNodeCompact::default(), + )]), + removed_nodes: Default::default(), + }; + + let sorted = StorageTrieUpdatesSorted { + is_deleted: false, + storage_nodes: vec![ + (Nibbles::from_nibbles_unchecked([0x02]), Some(BranchNodeCompact::default())), + (Nibbles::from_nibbles_unchecked([0x03]), None), + ], + }; + + storage.extend_from_sorted(&sorted); + + assert!(!storage.is_deleted); + assert_eq!(storage.storage_nodes.len(), 2); + assert!(storage.removed_nodes.contains(&Nibbles::from_nibbles_unchecked([0x03]))); + } + + /// Test deleted=true clears old nodes before extending (edge case) + #[test] + fn test_storage_trie_updates_extend_from_sorted_deleted() { + let mut storage = StorageTrieUpdates { + is_deleted: false, + storage_nodes: HashMap::from_iter([( + Nibbles::from_nibbles_unchecked([0x01]), + BranchNodeCompact::default(), + )]), + removed_nodes: Default::default(), + }; + + let sorted = StorageTrieUpdatesSorted { + is_deleted: true, + storage_nodes: vec![( + Nibbles::from_nibbles_unchecked([0x0a]), + Some(BranchNodeCompact::default()), + )], + }; + + storage.extend_from_sorted(&sorted); + + assert!(storage.is_deleted); + // Old nodes should be cleared when deleted + assert_eq!(storage.storage_nodes.len(), 1); + assert!(storage.storage_nodes.contains_key(&Nibbles::from_nibbles_unchecked([0x0a]))); + } + + /// Test empty nibbles are filtered out during conversion (edge case bug) + #[test] + fn test_trie_updates_extend_from_sorted_filters_empty_nibbles() { + let mut updates = TrieUpdates::default(); + + let sorted = TrieUpdatesSorted { + account_nodes: vec![ + (Nibbles::default(), Some(BranchNodeCompact::default())), // Empty nibbles + (Nibbles::from_nibbles_unchecked([0x01]), Some(BranchNodeCompact::default())), + ], + storage_tries: B256Map::default(), + }; + + updates.extend_from_sorted(&sorted); + + // Empty nibbles should be filtered out + assert_eq!(updates.account_nodes.len(), 1); + assert!(updates.account_nodes.contains_key(&Nibbles::from_nibbles_unchecked([0x01]))); + assert!(!updates.account_nodes.contains_key(&Nibbles::default())); + } } /// Bincode-compatible trie updates type serde implementations.