diff --git a/Cargo.lock b/Cargo.lock index 47d0a328b00..bd4d89ae6ac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10463,7 +10463,6 @@ dependencies = [ "alloy-trie", "arbitrary", "assert_matches", - "auto_impl", "itertools 0.14.0", "metrics", "pretty_assertions", diff --git a/crates/engine/tree/src/tree/payload_processor/mod.rs b/crates/engine/tree/src/tree/payload_processor/mod.rs index 92ff14890b5..bc274525a81 100644 --- a/crates/engine/tree/src/tree/payload_processor/mod.rs +++ b/crates/engine/tree/src/tree/payload_processor/mod.rs @@ -597,6 +597,7 @@ where proof_worker_handle, trie_metrics.clone(), sparse_state_trie, + parent_state_root, chunk_size, ); diff --git a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs b/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs index 46d9b2c3f5c..668002a0824 100644 --- a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs +++ b/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs @@ -27,8 +27,9 @@ use reth_trie_parallel::{ root::ParallelStateRootError, }; use reth_trie_sparse::{ - errors::SparseTrieResult, ConfigurableSparseTrie, DeferredDrops, LeafUpdate, - RevealableSparseTrie, SparseStateTrie, SparseTrie, + errors::{SparseStateTrieErrorKind, SparseTrieErrorKind, SparseTrieResult}, + ConfigurableSparseTrie, DeferredDrops, LeafUpdate, RevealableSparseTrie, SparseStateTrie, + SparseTrie, }; use revm_primitives::{hash_map::Entry, B256Map}; use tracing::{debug, debug_span, error, instrument, trace_span}; @@ -46,6 +47,8 @@ pub(super) struct SparseTrieCacheTask, /// `SparseStateTrie` used for computing the state root. trie: SparseStateTrie, + /// The parent block's state root. + parent_state_root: B256, /// Handle to the proof worker pools (storage and account). proof_worker_handle: ProofWorkerHandle, @@ -120,6 +123,7 @@ where proof_worker_handle: ProofWorkerHandle, metrics: MultiProofTaskMetrics, trie: SparseStateTrie, + parent_state_root: B256, chunk_size: usize, ) -> Self { let (proof_result_tx, proof_result_rx) = crossbeam_channel::unbounded(); @@ -138,6 +142,7 @@ where updates: hashed_state_rx, proof_worker_handle, trie, + parent_state_root, chunk_size, max_targets_for_chunking: DEFAULT_MAX_TARGETS_FOR_CHUNKING, account_updates: Default::default(), @@ -359,10 +364,25 @@ where debug!(target: "engine::root", "All proofs processed, ending calculation"); let start = Instant::now(); - let (state_root, trie_updates) = - self.trie.root_with_updates(&self.proof_worker_handle).map_err(|e| { - ParallelStateRootError::Other(format!("could not calculate state root: {e:?}")) - })?; + let (state_root, trie_updates) = match self.trie.root_with_updates() { + Ok(result) => result, + Err(err) + if matches!( + err.kind(), + SparseStateTrieErrorKind::Sparse(SparseTrieErrorKind::Blind) + ) => + { + // A still-blind account trie means this block never changed state, so preserve + // the cached parent root instead of fetching and revealing + // the unchanged root node. + (self.parent_state_root, TrieUpdates::default()) + } + Err(err) => { + return Err(ParallelStateRootError::Other(format!( + "could not calculate state root: {err:?}" + ))) + } + }; #[cfg(feature = "trie-debug")] let debug_recorders = self.trie.take_debug_recorders(); @@ -873,6 +893,11 @@ enum SparseTrieTaskMessage { mod tests { use super::*; use alloy_primitives::{keccak256, Address, B256, U256}; + use reth_provider::{ + providers::OverlayStateProviderFactory, test_utils::create_test_provider_factory, + }; + use reth_trie_db::ChangesetCache; + use reth_trie_parallel::proof_task::ProofTaskCtx; use reth_trie_sparse::ArenaParallelSparseTrie; #[test] @@ -953,4 +978,43 @@ mod tests { assert_eq!(decoded.storage_root, storage_root); assert_eq!(account_rlp_buf, encoded); } + + #[test] + fn run_returns_parent_root_without_revealing_blind_trie_when_no_state_updates() { + let runtime = reth_tasks::Runtime::test(); + let provider_factory = create_test_provider_factory(); + let overlay_factory = + OverlayStateProviderFactory::new(provider_factory, ChangesetCache::new()); + let proof_worker_handle = + ProofWorkerHandle::new(&runtime, ProofTaskCtx::new(overlay_factory), false); + + let default_trie = RevealableSparseTrie::blind_from(ConfigurableSparseTrie::Arena( + ArenaParallelSparseTrie::default(), + )); + let trie = SparseStateTrie::default() + .with_accounts_trie(default_trie.clone()) + .with_default_storage_trie(default_trie) + .with_updates(true); + + let parent_state_root = B256::from([0x55; 32]); + let (updates_tx, updates_rx) = crossbeam_channel::unbounded(); + let mut task = SparseTrieCacheTask::new_with_trie( + &runtime, + updates_rx, + proof_worker_handle, + MultiProofTaskMetrics::default(), + trie, + parent_state_root, + 1, + ); + + updates_tx.send(StateRootMessage::FinishedStateUpdates).unwrap(); + drop(updates_tx); + + let outcome = task.run().expect("state root computation should succeed"); + + assert_eq!(outcome.state_root, parent_state_root); + assert!(outcome.trie_updates.is_empty()); + assert!(task.trie.state_trie_ref().is_none(), "blind trie should not be revealed"); + } } diff --git a/crates/trie/parallel/Cargo.toml b/crates/trie/parallel/Cargo.toml index 637e98d3c09..4e82f55165d 100644 --- a/crates/trie/parallel/Cargo.toml +++ b/crates/trie/parallel/Cargo.toml @@ -17,7 +17,6 @@ reth-primitives-traits = { workspace = true, features = ["dashmap", "std"] } reth-execution-errors.workspace = true reth-provider.workspace = true reth-storage-errors.workspace = true -reth-trie-sparse = { workspace = true, features = ["std"] } reth-tasks = { workspace = true, features = ["rayon"] } reth-trie.workspace = true @@ -46,6 +45,7 @@ reth-metrics = { workspace = true, optional = true } metrics = { workspace = true, optional = true } # `trie-debug` feature +reth-trie-sparse = { workspace = true, optional = true, features = ["trie-debug"] } rand = { workspace = true, optional = true } [dev-dependencies] @@ -63,13 +63,17 @@ tokio = { workspace = true, features = ["rt", "rt-multi-thread", "macros"] } [features] default = ["metrics"] -metrics = ["reth-metrics", "dep:metrics", "reth-trie/metrics", "reth-trie-sparse/metrics"] -trie-debug = ["dep:rand", "reth-trie-sparse/trie-debug"] +metrics = ["reth-metrics", "dep:metrics", "reth-trie/metrics"] +trie-debug = [ + "dep:rand", + "dep:reth-trie-sparse", + "reth-trie-sparse?/trie-debug", +] test-utils = [ "reth-primitives-traits/test-utils", "reth-provider/test-utils", "reth-trie-db/test-utils", - "reth-trie-sparse/test-utils", "reth-trie/test-utils", "reth-tasks/test-utils", + "reth-trie-sparse?/test-utils", ] diff --git a/crates/trie/parallel/src/proof_task.rs b/crates/trie/parallel/src/proof_task.rs index a7546c9cd3e..2e48bbbaff1 100644 --- a/crates/trie/parallel/src/proof_task.rs +++ b/crates/trie/parallel/src/proof_task.rs @@ -4,8 +4,8 @@ //! # Architecture //! //! - **Worker Pools**: Pre-spawned workers with dedicated database transactions -//! - Storage pool: Handles storage proofs and blinded storage node requests -//! - Account pool: Handles account multiproofs and blinded account node requests +//! - Storage pool: Handles storage proofs +//! - Account pool: Handles account multiproofs //! - **Direct Channel Access**: `ProofWorkerHandle` provides type-safe queue methods with direct //! access to worker channels, eliminating routing overhead //! - **Automatic Shutdown**: Workers terminate gracefully when all handles are dropped @@ -38,26 +38,22 @@ use alloy_primitives::{ B256, U256, }; use crossbeam_channel::{unbounded, Receiver as CrossbeamReceiver, Sender as CrossbeamSender}; -use reth_execution_errors::{SparseTrieError, SparseTrieErrorKind, StateProofError}; +use reth_execution_errors::StateProofError; use reth_primitives_traits::{dashmap::DashMap, FastInstant as Instant}; use reth_provider::{DatabaseProviderROFactory, ProviderError, ProviderResult}; use reth_storage_errors::db::DatabaseError; use reth_tasks::Runtime; use reth_trie::{ hashed_cursor::{HashedCursorFactory, HashedStorageCursor, InstrumentedHashedCursor}, - proof::{ProofBlindedAccountProvider, ProofBlindedStorageProvider}, proof_v2, trie_cursor::{InstrumentedTrieCursor, TrieCursorFactory, TrieStorageCursor}, - DecodedMultiProofV2, HashedPostState, MultiProofTargetsV2, Nibbles, ProofTrieNodeV2, - ProofV2Target, + DecodedMultiProofV2, HashedPostState, MultiProofTargetsV2, ProofTrieNodeV2, ProofV2Target, }; -use reth_trie_sparse::provider::{RevealedNode, TrieNodeProvider, TrieNodeProviderFactory}; use std::{ cell::RefCell, rc::Rc, sync::{ atomic::{AtomicBool, AtomicUsize, Ordering}, - mpsc::{channel, Receiver, Sender}, Arc, }, time::Duration, @@ -69,8 +65,6 @@ use crate::proof_task_metrics::{ ProofTaskCursorMetrics, ProofTaskCursorMetricsCache, ProofTaskTrieMetrics, }; -type TrieNodeProviderResult = Result, SparseTrieError>; - /// Type alias for the V2 account proof calculator with instrumented cursors. type V2AccountProofCalculator<'a, Provider> = proof_v2::ProofCalculator< InstrumentedTrieCursor<'a, ::AccountTrieCursor<'a>>, @@ -338,15 +332,13 @@ impl ProofWorkerHandle { self.storage_work_tx .send(StorageWorkerJob::StorageProof { input, proof_result_sender }) .map_err(|err| { - if let StorageWorkerJob::StorageProof { proof_result_sender, .. } = err.0 { - let _ = proof_result_sender.send(StorageProofResultMessage { - hashed_address, - result: Err(DatabaseError::Other( - "storage workers unavailable".to_string(), - ) - .into()), - }); - } + let StorageWorkerJob::StorageProof { proof_result_sender, .. } = err.0; + let _ = proof_result_sender.send(StorageProofResultMessage { + hashed_address, + result: Err( + DatabaseError::Other("storage workers unavailable".to_string()).into() + ), + }); ProviderError::other(std::io::Error::other("storage workers unavailable")) }) @@ -365,51 +357,19 @@ impl ProofWorkerHandle { let error = ProviderError::other(std::io::Error::other("account workers unavailable")); - if let AccountWorkerJob::AccountMultiproof { input } = err.0 { - let ProofResultContext { sender: result_tx, state, start_time: start } = - input.into_proof_result_sender(); + let AccountWorkerJob::AccountMultiproof { input } = err.0; + let ProofResultContext { sender: result_tx, state, start_time: start } = + input.into_proof_result_sender(); - let _ = result_tx.send(ProofResultMessage { - result: Err(ParallelStateRootError::Provider(error.clone())), - elapsed: start.elapsed(), - state, - }); - } + let _ = result_tx.send(ProofResultMessage { + result: Err(ParallelStateRootError::Provider(error.clone())), + elapsed: start.elapsed(), + state, + }); error }) } - - /// Dispatch blinded storage node request to storage worker pool - pub(crate) fn dispatch_blinded_storage_node( - &self, - account: B256, - path: Nibbles, - ) -> Result, ProviderError> { - let (tx, rx) = channel(); - self.storage_work_tx - .send(StorageWorkerJob::BlindedStorageNode { account, path, result_sender: tx }) - .map_err(|_| { - ProviderError::other(std::io::Error::other("storage workers unavailable")) - })?; - - Ok(rx) - } - - /// Dispatch blinded account node request to account worker pool - pub(crate) fn dispatch_blinded_account_node( - &self, - path: Nibbles, - ) -> Result, ProviderError> { - let (tx, rx) = channel(); - self.account_work_tx - .send(AccountWorkerJob::BlindedAccountNode { path, result_sender: tx }) - .map_err(|_| { - ProviderError::other(std::io::Error::other("account workers unavailable")) - })?; - - Ok(rx) - } } /// Data used for initializing cursor factories that is shared across all proof worker instances. @@ -502,67 +462,6 @@ where Ok(StorageProofResult { proof, root }) } - - /// Process a blinded storage node request. - /// - /// Used by storage workers to retrieve blinded storage trie nodes for proof construction. - fn process_blinded_storage_node( - &self, - account: B256, - path: &Nibbles, - ) -> TrieNodeProviderResult { - let storage_node_provider = - ProofBlindedStorageProvider::new(&self.provider, &self.provider, account); - storage_node_provider.trie_node(path) - } -} -impl TrieNodeProviderFactory for ProofWorkerHandle { - type AccountNodeProvider = ProofTaskTrieNodeProvider; - type StorageNodeProvider = ProofTaskTrieNodeProvider; - - fn account_node_provider(&self) -> Self::AccountNodeProvider { - ProofTaskTrieNodeProvider::AccountNode { handle: self.clone() } - } - - fn storage_node_provider(&self, account: B256) -> Self::StorageNodeProvider { - ProofTaskTrieNodeProvider::StorageNode { account, handle: self.clone() } - } -} - -/// Trie node provider for retrieving trie nodes by path. -#[derive(Debug)] -pub enum ProofTaskTrieNodeProvider { - /// Blinded account trie node provider. - AccountNode { - /// Handle to the proof worker pools. - handle: ProofWorkerHandle, - }, - /// Blinded storage trie node provider. - StorageNode { - /// Target account. - account: B256, - /// Handle to the proof worker pools. - handle: ProofWorkerHandle, - }, -} - -impl TrieNodeProvider for ProofTaskTrieNodeProvider { - fn trie_node(&self, path: &Nibbles) -> Result, SparseTrieError> { - match self { - Self::AccountNode { handle } => { - let rx = handle - .dispatch_blinded_account_node(*path) - .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?; - rx.recv().map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))? - } - Self::StorageNode { handle, account } => { - let rx = handle - .dispatch_blinded_storage_node(*account, *path) - .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?; - rx.recv().map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))? - } - } - } } /// Channel used by worker threads to deliver `ProofResultMessage` items back to @@ -647,21 +546,12 @@ pub(crate) enum StorageWorkerJob { /// Context for sending the proof result. proof_result_sender: CrossbeamSender, }, - /// Blinded storage node retrieval request - BlindedStorageNode { - /// Target account - account: B256, - /// Path to the storage node - path: Nibbles, - /// Channel to send result back to original caller - result_sender: Sender, - }, } /// Worker for storage trie operations. /// /// Each worker maintains a dedicated database transaction and processes -/// storage proof requests and blinded node lookups. +/// storage proof requests. struct StorageProofWorker { /// Shared task context with database factory and prefix sets task_ctx: ProofTaskCtx, @@ -737,7 +627,6 @@ where ); let mut storage_proofs_processed = 0u64; - let mut storage_nodes_processed = 0u64; let mut cursor_metrics_cache = ProofTaskCursorMetricsCache::default(); let trie_cursor = proof_tx.provider.storage_trie_cursor(B256::ZERO)?; let hashed_cursor = proof_tx.provider.hashed_storage_cursor(B256::ZERO)?; @@ -787,17 +676,6 @@ where &mut storage_proofs_processed, ); } - - StorageWorkerJob::BlindedStorageNode { account, path, result_sender } => { - Self::process_blinded_node( - self.worker_id, - &proof_tx, - account, - path, - result_sender, - &mut storage_nodes_processed, - ); - } } // Mark worker as available again. @@ -813,14 +691,12 @@ where target: "trie::proof_task", worker_id = self.worker_id, storage_proofs_processed, - storage_nodes_processed, total_idle_time_us = total_idle_time.as_micros(), "Storage worker shutting down" ); #[cfg(feature = "metrics")] { - self.metrics.record_storage_nodes(storage_nodes_processed as usize); self.metrics.record_storage_worker_idle_time(total_idle_time); self.cursor_metrics.record(&mut cursor_metrics_cache); } @@ -883,59 +759,12 @@ where "Storage proof completed" ); } - - /// Processes a blinded storage node lookup request. - fn process_blinded_node( - worker_id: usize, - proof_tx: &ProofTaskTx, - account: B256, - path: Nibbles, - result_sender: Sender, - storage_nodes_processed: &mut u64, - ) where - Provider: TrieCursorFactory + HashedCursorFactory, - { - trace!( - target: "trie::proof_task", - worker_id, - ?account, - ?path, - "Processing blinded storage node" - ); - - let start = Instant::now(); - let result = proof_tx.process_blinded_storage_node(account, &path); - let elapsed = start.elapsed(); - - *storage_nodes_processed += 1; - - if result_sender.send(result).is_err() { - trace!( - target: "trie::proof_task", - worker_id, - ?account, - ?path, - storage_nodes_processed, - "Blinded storage node receiver dropped, discarding result" - ); - } - - trace!( - target: "trie::proof_task", - worker_id, - ?account, - ?path, - elapsed_us = elapsed.as_micros(), - total_processed = storage_nodes_processed, - "Blinded storage node completed" - ); - } } /// Worker for account trie operations. /// /// Each worker maintains a dedicated database transaction and processes -/// account multiproof requests and blinded node lookups. +/// account multiproof requests. struct AccountProofWorker { /// Shared task context with database factory and prefix sets task_ctx: ProofTaskCtx, @@ -1014,7 +843,6 @@ where ); let mut account_proofs_processed = 0u64; - let mut account_nodes_processed = 0u64; let mut cursor_metrics_cache = ProofTaskCursorMetricsCache::default(); // Create both account and storage calculators for V2 proofs. @@ -1100,16 +928,6 @@ where total_idle_time += value_encoder_stats.storage_wait_time; value_encoder_stats_cache.extend(&value_encoder_stats); } - - AccountWorkerJob::BlindedAccountNode { path, result_sender } => { - Self::process_blinded_node( - self.worker_id, - &provider, - path, - result_sender, - &mut account_nodes_processed, - ); - } } // Mark worker as available again. @@ -1126,14 +944,12 @@ where target: "trie::proof_task", worker_id=self.worker_id, account_proofs_processed, - account_nodes_processed, total_idle_time_us = total_idle_time.as_micros(), "Account worker shutting down" ); #[cfg(feature = "metrics")] { - self.metrics.record_account_nodes(account_nodes_processed as usize); self.metrics.record_account_worker_idle_time(total_idle_time); self.cursor_metrics.record(&mut cursor_metrics_cache); self.metrics.record_value_encoder_stats(&value_encoder_stats_cache); @@ -1234,53 +1050,6 @@ where value_encoder_stats } - - /// Processes a blinded account node lookup request. - fn process_blinded_node( - worker_id: usize, - provider: &Provider, - path: Nibbles, - result_sender: Sender, - account_nodes_processed: &mut u64, - ) where - Provider: TrieCursorFactory + HashedCursorFactory, - { - let span = debug_span!( - target: "trie::proof_task", - "Blinded account node calculation", - ?path, - ); - let _span_guard = span.enter(); - - trace!( - target: "trie::proof_task", - "Processing blinded account node" - ); - - let start = Instant::now(); - let account_node_provider = ProofBlindedAccountProvider::new(provider, provider); - let result = account_node_provider.trie_node(&path); - let elapsed = start.elapsed(); - - *account_nodes_processed += 1; - - if result_sender.send(result).is_err() { - trace!( - target: "trie::proof_task", - worker_id, - ?path, - account_nodes_processed, - "Blinded account node receiver dropped, discarding result" - ); - } - - trace!( - target: "trie::proof_task", - node_time_us = elapsed.as_micros(), - total_processed = account_nodes_processed, - "Blinded account node completed" - ); - } } /// Queues V2 storage proofs for all accounts in the targets and returns receivers. @@ -1377,13 +1146,6 @@ enum AccountWorkerJob { /// Account multiproof input parameters input: Box, }, - /// Blinded account node retrieval request - BlindedAccountNode { - /// Path to the account node - path: Nibbles, - /// Channel to send result back to original caller - result_sender: Sender, - }, } #[cfg(test)] diff --git a/crates/trie/parallel/src/proof_task_metrics.rs b/crates/trie/parallel/src/proof_task_metrics.rs index af4cbb1db88..a915b3bd4ac 100644 --- a/crates/trie/parallel/src/proof_task_metrics.rs +++ b/crates/trie/parallel/src/proof_task_metrics.rs @@ -11,10 +11,6 @@ use std::time::Duration; #[derive(Clone, Metrics)] #[metrics(scope = "trie.proof_task")] pub struct ProofTaskTrieMetrics { - /// A histogram for the number of blinded account nodes fetched. - blinded_account_nodes: Histogram, - /// A histogram for the number of blinded storage nodes fetched. - blinded_storage_nodes: Histogram, /// Histogram for storage worker idle time in seconds (waiting for proof jobs). storage_worker_idle_time_seconds: Histogram, /// Histogram for account worker idle time in seconds (waiting for proof jobs + storage @@ -35,16 +31,6 @@ pub struct ProofTaskTrieMetrics { } impl ProofTaskTrieMetrics { - /// Record account nodes fetched. - pub fn record_account_nodes(&self, count: usize) { - self.blinded_account_nodes.record(count as f64); - } - - /// Record storage nodes fetched. - pub fn record_storage_nodes(&self, count: usize) { - self.blinded_storage_nodes.record(count as f64); - } - /// Record storage worker idle time. pub fn record_storage_worker_idle_time(&self, duration: Duration) { self.storage_worker_idle_time_seconds.record(duration.as_secs_f64()); diff --git a/crates/trie/sparse/Cargo.toml b/crates/trie/sparse/Cargo.toml index 5d2c388a84e..035947987ca 100644 --- a/crates/trie/sparse/Cargo.toml +++ b/crates/trie/sparse/Cargo.toml @@ -24,7 +24,6 @@ alloy-primitives.workspace = true alloy-rlp.workspace = true # misc -auto_impl.workspace = true rayon = { workspace = true, optional = true } serde = { workspace = true, features = ["derive"], optional = true } serde_json = { workspace = true, optional = true } diff --git a/crates/trie/sparse/src/arena/mod.rs b/crates/trie/sparse/src/arena/mod.rs index 29643e80e5e..7558e59b3b4 100644 --- a/crates/trie/sparse/src/arena/mod.rs +++ b/crates/trie/sparse/src/arena/mod.rs @@ -592,8 +592,6 @@ impl Default for ArenaParallelismThresholds { /// the upper trie to route each update to the correct subtrie, then processes subtries in /// parallel when the update count exceeds [`ArenaParallelismThresholds::min_updates`]. /// -/// [`SparseTrie::update_leaf`] and [`SparseTrie::remove_leaf`] are not yet implemented. -/// /// After updates, structural changes (branch collapse, subtrie unwrapping) are handled by /// propagating dirty state back up through the upper trie. /// @@ -2428,23 +2426,6 @@ impl SparseTrie for ArenaParallelSparseTrie { Ok(()) } - fn update_leaf( - &mut self, - _full_path: Nibbles, - _value: Vec, - _provider: P, - ) -> SparseTrieResult<()> { - unimplemented!("ArenaParallelSparseTrie uses update_leaves for batch leaf updates") - } - - fn remove_leaf( - &mut self, - _full_path: &Nibbles, - _provider: P, - ) -> SparseTrieResult<()> { - unimplemented!("ArenaParallelSparseTrie uses update_leaves for batch leaf removals") - } - #[instrument(level = "trace", target = TRACE_TARGET, skip_all, ret)] fn root(&mut self) -> B256 { #[cfg(feature = "trie-debug")] diff --git a/crates/trie/sparse/src/lib.rs b/crates/trie/sparse/src/lib.rs index c52c059afe4..78058872dfa 100644 --- a/crates/trie/sparse/src/lib.rs +++ b/crates/trie/sparse/src/lib.rs @@ -26,8 +26,6 @@ pub use parallel::*; mod lower; -pub mod provider; - #[cfg(feature = "metrics")] mod metrics; diff --git a/crates/trie/sparse/src/parallel.rs b/crates/trie/sparse/src/parallel.rs index fa97e308a15..6e17e1d01f2 100644 --- a/crates/trie/sparse/src/parallel.rs +++ b/crates/trie/sparse/src/parallel.rs @@ -1,8 +1,8 @@ #[cfg(feature = "trie-debug")] use crate::debug_recorder::{LeafUpdateRecord, ProofTrieNodeRecord, RecordedOp, TrieDebugRecorder}; use crate::{ - lower::LowerSparseSubtrie, provider::TrieNodeProvider, LeafLookup, LeafLookupError, - RlpNodeStackItem, SparseNode, SparseNodeState, SparseNodeType, SparseTrie, SparseTrieUpdates, + lower::LowerSparseSubtrie, LeafLookup, LeafLookupError, RlpNodeStackItem, SparseNode, + SparseNodeState, SparseNodeType, SparseTrie, SparseTrieUpdates, }; use alloc::{borrow::Cow, boxed::Box, vec, vec::Vec}; use alloy_primitives::{ @@ -432,1043 +432,954 @@ impl SparseTrie for ParallelSparseTrie { } } - fn update_leaf( - &mut self, - full_path: Nibbles, - value: Vec, - _provider: P, - ) -> SparseTrieResult<()> { - debug_assert_eq!( - full_path.len(), - B256::len_bytes() * 2, - "update_leaf full_path must be 64 nibbles (32 bytes), got {} nibbles", - full_path.len() - ); + #[instrument(level = "trace", target = "trie::sparse::parallel", skip(self))] + fn root(&mut self) -> B256 { + trace!(target: "trie::parallel_sparse", "Calculating trie root hash"); - trace!( - target: "trie::parallel_sparse", - ?full_path, - value_len = value.len(), - "Updating leaf", - ); + #[cfg(feature = "trie-debug")] + self.debug_recorder.record(RecordedOp::Root); - // Check if the value already exists - if so, just update it (no structural changes needed) - if self.upper_subtrie.inner.values.contains_key(&full_path) { - self.prefix_set.insert(full_path); - self.upper_subtrie.inner.values.insert(full_path, value); - return Ok(()); - } - // Also check lower subtries for existing value - if let Some(subtrie) = self.lower_subtrie_for_path(&full_path) && - subtrie.inner.values.contains_key(&full_path) + if self.prefix_set.is_empty() && + let Some(rlp_node) = self + .upper_subtrie + .nodes + .get(&Nibbles::default()) + .and_then(|node| node.cached_rlp_node()) { - self.prefix_set.insert(full_path); - self.lower_subtrie_for_path_mut(&full_path) - .expect("subtrie exists") - .inner - .values - .insert(full_path, value); - return Ok(()); + return rlp_node + .as_hash() + .expect("RLP-encoding of the root node cannot be less than 32 bytes") } - // Insert value into upper subtrie temporarily. We'll move it to the correct subtrie - // during traversal, or clean it up if we error. - self.upper_subtrie.inner.values.insert(full_path, value.clone()); + // Update all lower subtrie hashes + self.update_subtrie_hashes(); - // Start at the root, traversing until we find either the node to update or a subtrie to - // update. - // - // We first traverse the upper subtrie for two levels, and moving any created nodes to a - // lower subtrie if necessary. - // - // We use `next` to keep track of the next node that we need to traverse to, and - // `new_nodes` to keep track of any nodes that were created during the traversal. - let mut new_nodes = Vec::new(); - let mut next = Some(Nibbles::default()); + // Update hashes for the upper subtrie using our specialized function + // that can access both upper and lower subtrie nodes + let mut prefix_set = core::mem::take(&mut self.prefix_set).freeze(); + let root_rlp = self.update_upper_subtrie_hashes(&mut prefix_set); - // Traverse the upper subtrie to find the node to update or the subtrie to update. - // - // We stop when the next node to traverse would be in a lower subtrie, or if there are no - // more nodes to traverse. - while let Some(current) = - next.as_mut().filter(|next| SparseSubtrieType::path_len_is_upper(next.len())) - { - // Traverse the next node, keeping track of any changed nodes and the next step in the - // trie. If traversal fails, clean up the value we inserted and propagate the error. - let step_result = self.upper_subtrie.update_next_node(current, &full_path); + // Return the root hash + root_rlp.as_hash().unwrap_or(EMPTY_ROOT_HASH) + } - if step_result.is_err() { - self.upper_subtrie.inner.values.remove(&full_path); - return step_result.map(|_| ()); - } + fn is_root_cached(&self) -> bool { + self.prefix_set.is_empty() && + self.upper_subtrie + .nodes + .get(&Nibbles::default()) + .is_some_and(|node| node.cached_rlp_node().is_some()) + } - match step_result? { - LeafUpdateStep::Continue => {} - LeafUpdateStep::Complete { inserted_nodes } => { - new_nodes.extend(inserted_nodes); - next = None; - } - LeafUpdateStep::NodeNotFound => { - next = None; - } - } - } + #[instrument(level = "trace", target = "trie::sparse::parallel", skip(self))] + fn update_subtrie_hashes(&mut self) { + trace!(target: "trie::parallel_sparse", "Updating subtrie hashes"); - // Move nodes from upper subtrie to lower subtries - for node_path in &new_nodes { - // Skip nodes that belong in the upper subtrie - if SparseSubtrieType::path_len_is_upper(node_path.len()) { - continue - } + #[cfg(feature = "trie-debug")] + self.debug_recorder.record(RecordedOp::UpdateSubtrieHashes); - let node = - self.upper_subtrie.nodes.remove(node_path).expect("node belongs to upper subtrie"); + // Take changed subtries according to the prefix set + let mut prefix_set = core::mem::take(&mut self.prefix_set).freeze(); + let num_changed_keys = prefix_set.len(); + let (mut changed_subtries, unchanged_prefix_set) = + self.take_changed_lower_subtries(&mut prefix_set); - // If it's a leaf node, extract its value before getting mutable reference to subtrie. - let leaf_value = if let SparseNode::Leaf { key, .. } = &node { - let mut leaf_full_path = *node_path; - leaf_full_path.extend(key); - Some(( - leaf_full_path, - self.upper_subtrie - .inner - .values - .remove(&leaf_full_path) - .expect("leaf nodes have associated values entries"), - )) - } else { - None - }; + // update metrics + #[cfg(feature = "metrics")] + self.metrics.subtries_updated.record(changed_subtries.len() as f64); - // Get or create the subtrie with the exact node path (not truncated to 2 nibbles). - let subtrie = self.subtrie_for_path_mut(node_path); + // Update the prefix set with the keys that didn't have matching subtries + self.prefix_set = unchanged_prefix_set; - // Insert the leaf value if we have one - if let Some((leaf_full_path, value)) = leaf_value { - subtrie.inner.values.insert(leaf_full_path, value); + // Update subtrie hashes serially parallelism is not enabled + if !self.is_update_parallelism_enabled(num_changed_keys) { + for changed_subtrie in &mut changed_subtries { + changed_subtrie.subtrie.update_hashes( + &mut changed_subtrie.prefix_set, + &mut changed_subtrie.update_actions_buf, + &self.branch_node_masks, + ); } - // Insert the node into the lower subtrie - subtrie.nodes.insert(*node_path, node); + self.insert_changed_subtries(changed_subtries); + return } - // If we reached the max depth of the upper trie, we may have had more nodes to insert. - if let Some(next_path) = next.filter(|n| !SparseSubtrieType::path_len_is_upper(n.len())) { - // The value was inserted into the upper subtrie's `values` at the top of this method. - // At this point we know the value is not in the upper subtrie, and the call to - // `update_leaf` below will insert it into the lower subtrie. So remove it from the - // upper subtrie. - self.upper_subtrie.inner.values.remove(&full_path); + #[cfg(not(feature = "std"))] + unreachable!("nostd is checked by is_update_parallelism_enabled"); - // Use subtrie_for_path to ensure the subtrie has the correct path. - // - // The next_path here represents where we need to continue traversal, which may - // be longer than 2 nibbles if we're following an extension node. - let subtrie = self.subtrie_for_path_mut(&next_path); + #[cfg(feature = "std")] + // Update subtrie hashes in parallel + { + use rayon::prelude::*; - // Create an empty root at the subtrie path if the subtrie is empty - if subtrie.nodes.is_empty() { - subtrie.nodes.insert(subtrie.path, SparseNode::Empty); - } + changed_subtries.par_iter_mut().for_each(|changed_subtrie| { + #[cfg(feature = "metrics")] + let start = Instant::now(); + changed_subtrie.subtrie.update_hashes( + &mut changed_subtrie.prefix_set, + &mut changed_subtrie.update_actions_buf, + &self.branch_node_masks, + ); + #[cfg(feature = "metrics")] + self.metrics.subtrie_hash_update_latency.record(start.elapsed()); + }); - // If we didn't update the target leaf, we need to call update_leaf on the subtrie - // to ensure that the leaf is updated correctly. - if let Err(e) = subtrie.update_leaf(full_path, value) { - // Clean up: remove the value from lower subtrie if it was inserted - if let Some(lower) = self.lower_subtrie_for_path_mut(&full_path) { - lower.inner.values.remove(&full_path); - } - return Err(e); - } + self.insert_changed_subtries(changed_subtries); } - - // Insert into prefix_set only after all operations succeed - self.prefix_set.insert(full_path); - - Ok(()) } - fn remove_leaf( - &mut self, - full_path: &Nibbles, - _provider: P, - ) -> SparseTrieResult<()> { - debug_assert_eq!( - full_path.len(), - B256::len_bytes() * 2, - "remove_leaf full_path must be 64 nibbles (32 bytes), got {} nibbles", - full_path.len() - ); - - trace!( - target: "trie::parallel_sparse", - ?full_path, - "Removing leaf", - ); - - // When removing a leaf node it's possibly necessary to modify its parent node, and possibly - // the parent's parent node. It is not ever necessary to descend further than that; once an - // extension node is hit it must terminate in a branch or the root, which won't need further - // updates. So the situation with maximum updates is: - // - // - Leaf - // - Branch with 2 children, one being this leaf - // - Extension - // - // ...which will result in just a leaf or extension, depending on what the branch's other - // child is. - // - // Therefore, first traverse the trie in order to find the leaf node and at most its parent - // and grandparent. - - let leaf_path; - let leaf_subtrie_type; + fn get_leaf_value(&self, full_path: &Nibbles) -> Option<&Vec> { + // `subtrie_for_path` is intended for a node path, but here we are using a full key path. So + // we need to check if the subtrie that the key might belong to has any nodes; if not then + // the key's portion of the trie doesn't have enough depth to reach into the subtrie, and + // the key will be in the upper subtrie + if let Some(subtrie) = self.subtrie_for_path(full_path) && + !subtrie.is_empty() + { + return subtrie.inner.values.get(full_path); + } - let mut branch_parent_path: Option = None; - let mut branch_parent_node: Option = None; + self.upper_subtrie.inner.values.get(full_path) + } - let mut ext_grandparent_path: Option = None; - let mut ext_grandparent_node: Option = None; + fn updates_ref(&self) -> Cow<'_, SparseTrieUpdates> { + self.updates.as_ref().map_or(Cow::Owned(SparseTrieUpdates::default()), Cow::Borrowed) + } - let mut curr_path = Nibbles::new(); // start traversal from root - let mut curr_subtrie_type = SparseSubtrieType::Upper; + fn take_updates(&mut self) -> SparseTrieUpdates { + match self.updates.take() { + Some(updates) => { + // NOTE: we need to preserve Some case + self.updates = Some(SparseTrieUpdates::with_capacity( + updates.updated_nodes.len(), + updates.removed_nodes.len(), + )); + updates + } + None => SparseTrieUpdates::default(), + } + } - // List of node paths which need to be marked dirty - let mut paths_to_mark_dirty = Vec::new(); + fn wipe(&mut self) { + self.upper_subtrie.wipe(); + for trie in &mut *self.lower_subtries { + trie.wipe(); + } + self.prefix_set = PrefixSetMut::all(); + self.updates = self.updates.is_some().then(SparseTrieUpdates::wiped); + } - loop { - let curr_subtrie = match curr_subtrie_type { - SparseSubtrieType::Upper => &mut self.upper_subtrie, - SparseSubtrieType::Lower(idx) => { - self.lower_subtries[idx].as_revealed_mut().expect("lower subtrie is revealed") - } - }; - let curr_node = curr_subtrie.nodes.get_mut(&curr_path).unwrap(); + fn clear(&mut self) { + self.upper_subtrie.clear(); + self.upper_subtrie.nodes.insert(Nibbles::default(), SparseNode::Empty); + for subtrie in &mut *self.lower_subtries { + subtrie.clear(); + } + self.prefix_set.clear(); + self.updates = None; + self.branch_node_masks.clear(); + #[cfg(feature = "trie-debug")] + self.debug_recorder.reset(); + // `update_actions_buffers` doesn't need to be cleared; we want to reuse the Vecs it has + // buffered, and all of those are already inherently cleared when they get used. + } - match Self::find_next_to_leaf(&curr_path, curr_node, full_path) { - FindNextToLeafOutcome::NotFound => return Ok(()), // leaf isn't in the trie - FindNextToLeafOutcome::BlindedNode(path) => { - return Err(SparseTrieErrorKind::BlindedNode(path).into()) - } - FindNextToLeafOutcome::Found => { - // this node is the target leaf - leaf_path = curr_path; - leaf_subtrie_type = curr_subtrie_type; - break; - } - FindNextToLeafOutcome::ContinueFrom(next_path) => { - // Any branches/extensions along the path to the leaf will have their `hash` - // field unset, as it will no longer be valid once the leaf is removed. - match curr_node { - SparseNode::Branch { .. } => { - paths_to_mark_dirty - .push((SparseSubtrieType::from_path(&curr_path), curr_path)); + fn find_leaf( + &self, + full_path: &Nibbles, + expected_value: Option<&Vec>, + ) -> Result { + // Inclusion proof + // + // First, do a quick check if the value exists in either the upper or lower subtrie's values + // map. We assume that if there exists a leaf node, then its value will be in the `values` + // map. + if let Some(actual_value) = core::iter::once(self.upper_subtrie.as_ref()) + .chain(self.lower_subtrie_for_path(full_path)) + .filter_map(|subtrie| subtrie.inner.values.get(full_path)) + .next() + { + // We found the leaf, check if the value matches (if expected value was provided) + return expected_value + .is_none_or(|v| v == actual_value) + .then_some(LeafLookup::Exists) + .ok_or_else(|| LeafLookupError::ValueMismatch { + path: *full_path, + expected: expected_value.cloned(), + actual: actual_value.clone(), + }) + } - // If there is already an extension leading into a branch, then that - // extension is no longer relevant. - match (&branch_parent_path, &ext_grandparent_path) { - (Some(branch), Some(ext)) if branch.len() > ext.len() => { - ext_grandparent_path = None; - ext_grandparent_node = None; - } - _ => (), - }; - branch_parent_path = Some(curr_path); - branch_parent_node = Some(curr_node.clone()); - } - SparseNode::Extension { .. } => { - paths_to_mark_dirty - .push((SparseSubtrieType::from_path(&curr_path), curr_path)); + // If the value does not exist in the `values` map, then this means that the leaf either: + // - Does not exist in the trie + // - Is missing from the witness + // We traverse the trie to find the location where this leaf would have been, showing + // that it is not in the trie. Or we find a blinded node, showing that the witness is + // not complete. + let mut curr_path = Nibbles::new(); // start traversal from root + let mut curr_subtrie = self.upper_subtrie.as_ref(); + let mut curr_subtrie_is_upper = true; - // We can assume a new branch node will be found after the extension, so - // there's no need to modify branch_parent_path/node even if it's - // already set. - ext_grandparent_path = Some(curr_path); - ext_grandparent_node = Some(curr_node.clone()); - } - SparseNode::Empty | SparseNode::Leaf { .. } => { - unreachable!( - "find_next_to_leaf only continues to a branch or extension" - ) - } + loop { + match curr_subtrie.nodes.get(&curr_path).unwrap() { + SparseNode::Empty => return Ok(LeafLookup::NonExistent), + SparseNode::Leaf { key, .. } => { + let mut found_full_path = curr_path; + found_full_path.extend(key); + assert!(&found_full_path != full_path, "target leaf {full_path:?} found, even though value wasn't in values hashmap"); + return Ok(LeafLookup::NonExistent) + } + SparseNode::Extension { key, .. } => { + if full_path.len() == curr_path.len() { + return Ok(LeafLookup::NonExistent) } - - curr_path = next_path; - - // Update subtrie type if we're crossing into the lower trie. - let next_subtrie_type = SparseSubtrieType::from_path(&curr_path); - if matches!(curr_subtrie_type, SparseSubtrieType::Upper) && - matches!(next_subtrie_type, SparseSubtrieType::Lower(_)) - { - curr_subtrie_type = next_subtrie_type; + curr_path.extend(key); + if !full_path.starts_with(&curr_path) { + return Ok(LeafLookup::NonExistent) } } - }; + SparseNode::Branch { state_mask, blinded_mask, blinded_hashes, .. } => { + if full_path.len() == curr_path.len() { + return Ok(LeafLookup::NonExistent) + } + let nibble = full_path.get_unchecked(curr_path.len()); + if !state_mask.is_bit_set(nibble) { + return Ok(LeafLookup::NonExistent) + } + curr_path.push_unchecked(nibble); + if blinded_mask.is_bit_set(nibble) { + return Err(LeafLookupError::BlindedNode { + path: curr_path, + hash: blinded_hashes[nibble as usize], + }) + } + } + } + + // If we were previously looking at the upper trie, and the new path is in the + // lower trie, we need to pull out a ref to the lower trie. + if curr_subtrie_is_upper && + let Some(lower_subtrie) = self.lower_subtrie_for_path(&curr_path) + { + curr_subtrie = lower_subtrie; + curr_subtrie_is_upper = false; + } } + } - // Before mutating, check if branch collapse would require revealing a blinded node. - // This ensures remove_leaf is atomic: if it errors, the trie is unchanged. - if let (Some(branch_path), Some(SparseNode::Branch { state_mask, blinded_mask, .. })) = - (&branch_parent_path, &branch_parent_node) - { - let mut check_mask = *state_mask; - let child_nibble = leaf_path.get_unchecked(branch_path.len()); - check_mask.unset_bit(child_nibble); + fn shrink_nodes_to(&mut self, size: usize) { + // Distribute the capacity across upper and lower subtries + // + // Always include upper subtrie, plus any lower subtries + let total_subtries = 1 + NUM_LOWER_SUBTRIES; + let size_per_subtrie = size / total_subtries; - if check_mask.count_bits() == 1 { - let remaining_nibble = - check_mask.first_set_bit_index().expect("state mask is not empty"); + // Shrink the upper subtrie + self.upper_subtrie.shrink_nodes_to(size_per_subtrie); - if blinded_mask.is_bit_set(remaining_nibble) { - let mut path = *branch_path; - path.push_unchecked(remaining_nibble); - return Err(SparseTrieErrorKind::BlindedNode(path).into()); - } - } + // Shrink lower subtries (works for both revealed and blind with allocation) + for subtrie in &mut *self.lower_subtries { + subtrie.shrink_nodes_to(size_per_subtrie); } - // We've traversed to the leaf and collected its ancestors as necessary. Remove the leaf - // from its SparseSubtrie and reset the hashes of the nodes along the path. - self.prefix_set.insert(*full_path); - let leaf_subtrie = match leaf_subtrie_type { - SparseSubtrieType::Upper => &mut self.upper_subtrie, - SparseSubtrieType::Lower(idx) => { - self.lower_subtries[idx].as_revealed_mut().expect("lower subtrie is revealed") - } - }; - leaf_subtrie.inner.values.remove(full_path); - for (subtrie_type, path) in paths_to_mark_dirty { - let node = match subtrie_type { - SparseSubtrieType::Upper => self.upper_subtrie.nodes.get_mut(&path), - SparseSubtrieType::Lower(idx) => self.lower_subtries[idx] - .as_revealed_mut() - .expect("lower subtrie is revealed") - .nodes - .get_mut(&path), - } - .expect("node exists"); + // shrink masks map + self.branch_node_masks.shrink_to(size); + } - match node { - SparseNode::Extension { state, .. } | SparseNode::Branch { state, .. } => { - *state = SparseNodeState::Dirty - } - SparseNode::Empty | SparseNode::Leaf { .. } => { - unreachable!( - "only branch and extension nodes can be marked dirty when removing a leaf" - ) - } - } - } - self.remove_node(&leaf_path); + fn shrink_values_to(&mut self, size: usize) { + // Distribute the capacity across upper and lower subtries + // + // Always include upper subtrie, plus any lower subtries + let total_subtries = 1 + NUM_LOWER_SUBTRIES; + let size_per_subtrie = size / total_subtries; - // If the leaf was at the root replace its node with the empty value. We can stop execution - // here, all remaining logic is related to the ancestors of the leaf. - if leaf_path.is_empty() { - self.upper_subtrie.nodes.insert(leaf_path, SparseNode::Empty); - return Ok(()) + // Shrink the upper subtrie + self.upper_subtrie.shrink_values_to(size_per_subtrie); + + // Shrink lower subtries (works for both revealed and blind with allocation) + for subtrie in &mut *self.lower_subtries { + subtrie.shrink_values_to(size_per_subtrie); } + } - // If there is a parent branch node (very likely, unless the leaf is at the root) execute - // any required changes for that node, relative to the removed leaf. - if let ( - Some(branch_path), - &Some(SparseNode::Branch { mut state_mask, blinded_mask, ref blinded_hashes, .. }), - ) = (&branch_parent_path, &branch_parent_node) - { - let child_nibble = leaf_path.get_unchecked(branch_path.len()); - state_mask.unset_bit(child_nibble); + /// O(1) size hint based on total node count (including hash stubs). + fn size_hint(&self) -> usize { + let upper_count = self.upper_subtrie.nodes.len(); + let lower_count: usize = self + .lower_subtries + .iter() + .filter_map(|s| s.as_revealed_ref()) + .map(|s| s.nodes.len()) + .sum(); + upper_count + lower_count + } - let new_branch_node = if state_mask.count_bits() == 1 { - // If only one child is left set in the branch node, we need to collapse it. Get - // full path of the only child node left. - let remaining_child_nibble = - state_mask.first_set_bit_index().expect("state mask is not empty"); - let mut remaining_child_path = *branch_path; - remaining_child_path.push_unchecked(remaining_child_nibble); + fn memory_size(&self) -> usize { + self.memory_size() + } - trace!( - target: "trie::parallel_sparse", - ?leaf_path, - ?branch_path, - ?remaining_child_path, - "Branch node has only one child", - ); + fn prune(&mut self, retained_leaves: &[Nibbles]) -> usize { + #[cfg(feature = "trie-debug")] + self.debug_recorder.reset(); - // If the remaining child node is not yet revealed then we have to reveal it here, - // otherwise it's not possible to know how to collapse the branch. - if blinded_mask.is_bit_set(remaining_child_nibble) { - return Err(SparseTrieErrorKind::BlindedNode(remaining_child_path).into()); - } + let mut retained_leaves = retained_leaves.to_vec(); + retained_leaves.sort_unstable(); - let remaining_child_node = self - .subtrie_for_path_mut(&remaining_child_path) - .nodes - .get(&remaining_child_path) - .unwrap(); + let mut effective_pruned_roots = Vec::::new(); + let mut stack: SmallVec<[Nibbles; 32]> = SmallVec::new(); + stack.push(Nibbles::default()); - let (new_branch_node, remove_child) = Self::branch_changes_on_leaf_removal( - branch_path, - &remaining_child_path, - remaining_child_node, - ); + while let Some(path) = stack.pop() { + let Some(node) = + self.subtrie_for_path(&path).and_then(|subtrie| subtrie.nodes.get(&path).cloned()) + else { + continue; + }; - if remove_child { - self.move_value_on_leaf_removal( - branch_path, - &new_branch_node, - &remaining_child_path, - ); - self.remove_node(&remaining_child_path); - } + match node { + SparseNode::Empty | SparseNode::Leaf { .. } => {} + SparseNode::Extension { key, state, .. } => { + let mut child = path; + child.extend(&key); - if let Some(updates) = self.updates.as_mut() { - updates.updated_nodes.remove(branch_path); - updates.removed_nodes.insert(*branch_path); - } + if has_retained_descendant(&retained_leaves, &child) { + stack.push(child); + continue; + } - new_branch_node - } else { - // If more than one child is left set in the branch, we just re-insert it with the - // updated state_mask. - SparseNode::Branch { - state_mask, - blinded_mask, - blinded_hashes: blinded_hashes.clone(), - state: SparseNodeState::Dirty, + // Root extension has no parent branch edge to blind; keep it as-is. + if path.is_empty() { + continue; + } + + let Some(hash) = state.cached_hash() else { continue }; + self.subtrie_for_path_mut_untracked(&path) + .expect("node subtrie exists") + .nodes + .remove(&path); + + let parent_path = path.slice(0..path.len() - 1); + // Parent can live in a different subtrie when `path` is the root of a lower + // subtrie, so resolve it by `parent_path` rather than reusing `path`'s subtrie. + let SparseNode::Branch { blinded_mask, blinded_hashes, .. } = self + .subtrie_for_path_mut_untracked(&parent_path) + .expect("parent subtrie exists") + .nodes + .get_mut(&parent_path) + .expect("expected parent branch node") + else { + panic!("expected branch node at path {parent_path:?}"); + }; + + let nibble = path.last().unwrap(); + blinded_mask.set_bit(nibble); + blinded_hashes[nibble as usize] = hash; + effective_pruned_roots.push(path); } - }; + SparseNode::Branch { state_mask, blinded_mask, blinded_hashes, .. } => { + let mut blinded_mask = blinded_mask; + let mut blinded_hashes = blinded_hashes; + for nibble in state_mask.iter() { + if blinded_mask.is_bit_set(nibble) { + continue; + } - let branch_subtrie = self.subtrie_for_path_mut(branch_path); - branch_subtrie.nodes.insert(*branch_path, new_branch_node.clone()); - branch_parent_node = Some(new_branch_node); - }; + let mut child = path; + child.push_unchecked(nibble); + if has_retained_descendant(&retained_leaves, &child) { + stack.push(child); + continue; + } - // If there is a grandparent extension node then there will necessarily be a parent branch - // node. Execute any required changes for the extension node, relative to the (possibly now - // replaced with a leaf or extension) branch node. - if let (Some(ext_path), Some(SparseNode::Extension { key: shortkey, .. })) = - (ext_grandparent_path, &ext_grandparent_node) - { - let ext_subtrie = self.subtrie_for_path_mut(&ext_path); - let branch_path = branch_parent_path.as_ref().unwrap(); + let Entry::Occupied(entry) = + self.subtrie_for_path_mut_untracked(&child).unwrap().nodes.entry(child) + else { + panic!("expected node at path {child:?}"); + }; + + let Some(hash) = entry.get().cached_hash() else { + continue; + }; + entry.remove(); + blinded_mask.set_bit(nibble); + blinded_hashes[nibble as usize] = hash; + effective_pruned_roots.push(child); + } - if let Some(new_ext_node) = Self::extension_changes_on_leaf_removal( - &ext_path, - shortkey, - branch_path, - branch_parent_node.as_ref().unwrap(), - ) { - ext_subtrie.nodes.insert(ext_path, new_ext_node.clone()); - self.move_value_on_leaf_removal(&ext_path, &new_ext_node, branch_path); - self.remove_node(branch_path); + let SparseNode::Branch { + blinded_mask: old_blinded_mask, + blinded_hashes: old_blinded_hashes, + .. + } = self + .subtrie_for_path_mut_untracked(&path) + .unwrap() + .nodes + .get_mut(&path) + .unwrap() + else { + unreachable!("expected branch node at path {path:?}"); + }; + *old_blinded_mask = blinded_mask; + *old_blinded_hashes = blinded_hashes; + } } } - Ok(()) + Self::finalize_pruned_roots(self, effective_pruned_roots) } - #[instrument(level = "trace", target = "trie::sparse::parallel", skip(self))] - fn root(&mut self) -> B256 { - trace!(target: "trie::parallel_sparse", "Calculating trie root hash"); + fn update_leaves( + &mut self, + updates: &mut alloy_primitives::map::B256Map, + mut proof_required_fn: impl FnMut(B256, u8), + ) -> SparseTrieResult<()> { + use crate::LeafUpdate; #[cfg(feature = "trie-debug")] - self.debug_recorder.record(RecordedOp::Root); - - if self.prefix_set.is_empty() && - let Some(rlp_node) = self - .upper_subtrie - .nodes - .get(&Nibbles::default()) - .and_then(|node| node.cached_rlp_node()) - { - return rlp_node - .as_hash() - .expect("RLP-encoding of the root node cannot be less than 32 bytes") - } - - // Update all lower subtrie hashes - self.update_subtrie_hashes(); - - // Update hashes for the upper subtrie using our specialized function - // that can access both upper and lower subtrie nodes - let mut prefix_set = core::mem::take(&mut self.prefix_set).freeze(); - let root_rlp = self.update_upper_subtrie_hashes(&mut prefix_set); + let recorded_updates: Vec<_> = + updates.iter().map(|(k, v)| (*k, LeafUpdateRecord::from(v))).collect(); + #[cfg(feature = "trie-debug")] + let mut recorded_proof_targets: Vec<(B256, u8)> = Vec::new(); - // Return the root hash - root_rlp.as_hash().unwrap_or(EMPTY_ROOT_HASH) - } + // Drain updates to avoid cloning keys while preserving the map's allocation. + // On success, entries remain removed; on blinded node failure, they're re-inserted. + let drained: Vec<_> = updates.drain().collect(); - fn is_root_cached(&self) -> bool { - self.prefix_set.is_empty() && - self.upper_subtrie - .nodes - .get(&Nibbles::default()) - .is_some_and(|node| node.cached_rlp_node().is_some()) - } + for (key, update) in drained { + let full_path = Nibbles::unpack(key); - #[instrument(level = "trace", target = "trie::sparse::parallel", skip(self))] - fn update_subtrie_hashes(&mut self) { - trace!(target: "trie::parallel_sparse", "Updating subtrie hashes"); + match update { + LeafUpdate::Changed(value) => { + if value.is_empty() { + // Removal is atomic - returns a retriable error before any mutations (via + // pre_validate_reveal_chain). + match self.remove_leaf(&full_path) { + Ok(()) => {} + Err(e) => { + if let Some(path) = Self::get_retriable_path(&e) { + let (target_key, min_len) = + Self::proof_target_for_path(key, &full_path, &path); + proof_required_fn(target_key, min_len); + #[cfg(feature = "trie-debug")] + recorded_proof_targets.push((target_key, min_len)); + updates.insert(key, LeafUpdate::Changed(value)); + } else { + return Err(e); + } + } + } + } else { + // Update/insert: update_leaf is atomic - cleans up on error. + if let Err(e) = self.update_leaf(full_path, value.clone()) { + if let Some(path) = Self::get_retriable_path(&e) { + let (target_key, min_len) = + Self::proof_target_for_path(key, &full_path, &path); + proof_required_fn(target_key, min_len); + #[cfg(feature = "trie-debug")] + recorded_proof_targets.push((target_key, min_len)); + updates.insert(key, LeafUpdate::Changed(value)); + } else { + return Err(e); + } + } + } + } + LeafUpdate::Touched => { + // Touched is read-only: check if path is accessible, request proof if blinded. + match self.find_leaf(&full_path, None) { + Err(LeafLookupError::BlindedNode { path, .. }) => { + let (target_key, min_len) = + Self::proof_target_for_path(key, &full_path, &path); + proof_required_fn(target_key, min_len); + #[cfg(feature = "trie-debug")] + recorded_proof_targets.push((target_key, min_len)); + updates.insert(key, LeafUpdate::Touched); + } + // Path is fully revealed (exists or proven non-existent), no action needed. + Ok(_) | Err(LeafLookupError::ValueMismatch { .. }) => {} + } + } + } + } #[cfg(feature = "trie-debug")] - self.debug_recorder.record(RecordedOp::UpdateSubtrieHashes); + self.debug_recorder.record(RecordedOp::UpdateLeaves { + updates: recorded_updates, + proof_targets: recorded_proof_targets, + }); - // Take changed subtries according to the prefix set - let mut prefix_set = core::mem::take(&mut self.prefix_set).freeze(); - let num_changed_keys = prefix_set.len(); - let (mut changed_subtries, unchanged_prefix_set) = - self.take_changed_lower_subtries(&mut prefix_set); + Ok(()) + } - // update metrics - #[cfg(feature = "metrics")] - self.metrics.subtries_updated.record(changed_subtries.len() as f64); + #[cfg(feature = "trie-debug")] + fn take_debug_recorder(&mut self) -> TrieDebugRecorder { + core::mem::take(&mut self.debug_recorder) + } - // Update the prefix set with the keys that didn't have matching subtries - self.prefix_set = unchanged_prefix_set; + fn commit_updates( + &mut self, + updated: &HashMap, + removed: &HashSet, + ) { + // Sync branch_node_masks with what's being committed to DB. + // This ensures that on subsequent root() calls, the masks reflect the actual + // DB state, which is needed for correct removal detection. + self.branch_node_masks.reserve(updated.len()); + for (path, node) in updated { + self.branch_node_masks.insert( + *path, + BranchNodeMasks { tree_mask: node.tree_mask, hash_mask: node.hash_mask }, + ); + } + for path in removed { + self.branch_node_masks.remove(path); + } + } +} - // Update subtrie hashes serially parallelism is not enabled - if !self.is_update_parallelism_enabled(num_changed_keys) { - for changed_subtrie in &mut changed_subtries { - changed_subtrie.subtrie.update_hashes( - &mut changed_subtrie.prefix_set, - &mut changed_subtrie.update_actions_buf, - &self.branch_node_masks, - ); - } +impl ParallelSparseTrie { + /// Sets the thresholds that control when parallelism is used during operations. + pub const fn with_parallelism_thresholds(mut self, thresholds: ParallelismThresholds) -> Self { + self.parallelism_thresholds = thresholds; + self + } - self.insert_changed_subtries(changed_subtries); - return - } + /// Returns true if retaining updates is enabled for the overall trie. + const fn updates_enabled(&self) -> bool { + self.updates.is_some() + } + /// Returns true if parallelism should be enabled for revealing the given number of nodes. + /// Will always return false in nostd builds. + const fn is_reveal_parallelism_enabled(&self, num_nodes: usize) -> bool { #[cfg(not(feature = "std"))] - unreachable!("nostd is checked by is_update_parallelism_enabled"); + { + let _ = num_nodes; + return false; + } #[cfg(feature = "std")] - // Update subtrie hashes in parallel { - use rayon::prelude::*; - - changed_subtries.par_iter_mut().for_each(|changed_subtrie| { - #[cfg(feature = "metrics")] - let start = Instant::now(); - changed_subtrie.subtrie.update_hashes( - &mut changed_subtrie.prefix_set, - &mut changed_subtrie.update_actions_buf, - &self.branch_node_masks, - ); - #[cfg(feature = "metrics")] - self.metrics.subtrie_hash_update_latency.record(start.elapsed()); - }); - - self.insert_changed_subtries(changed_subtries); + num_nodes >= self.parallelism_thresholds.min_revealed_nodes } } - fn get_leaf_value(&self, full_path: &Nibbles) -> Option<&Vec> { - // `subtrie_for_path` is intended for a node path, but here we are using a full key path. So - // we need to check if the subtrie that the key might belong to has any nodes; if not then - // the key's portion of the trie doesn't have enough depth to reach into the subtrie, and - // the key will be in the upper subtrie - if let Some(subtrie) = self.subtrie_for_path(full_path) && - !subtrie.is_empty() + /// Returns true if parallelism should be enabled for updating hashes with the given number + /// of changed keys. Will always return false in nostd builds. + const fn is_update_parallelism_enabled(&self, num_changed_keys: usize) -> bool { + #[cfg(not(feature = "std"))] { - return subtrie.inner.values.get(full_path); + let _ = num_changed_keys; + return false; } - self.upper_subtrie.inner.values.get(full_path) + #[cfg(feature = "std")] + { + num_changed_keys >= self.parallelism_thresholds.min_updated_nodes + } } - fn updates_ref(&self) -> Cow<'_, SparseTrieUpdates> { - self.updates.as_ref().map_or(Cow::Owned(SparseTrieUpdates::default()), Cow::Borrowed) + /// Checks if an error is retriable (`BlindedNode` or `NodeNotFoundInProvider`) and extracts + /// the path if so. + /// + /// Both error types indicate that a node needs to be revealed before the operation can + /// succeed. `BlindedNode` occurs when traversing to a Hash node, while `NodeNotFoundInProvider` + /// occurs when `retain_updates` is enabled and an extension node's child needs revealing. + const fn get_retriable_path(e: &SparseTrieError) -> Option { + match e.kind() { + SparseTrieErrorKind::BlindedNode(path) | + SparseTrieErrorKind::NodeNotFoundInProvider { path } => Some(*path), + _ => None, + } } - fn take_updates(&mut self) -> SparseTrieUpdates { - match self.updates.take() { - Some(updates) => { - // NOTE: we need to preserve Some case - self.updates = Some(SparseTrieUpdates::with_capacity( - updates.updated_nodes.len(), - updates.removed_nodes.len(), - )); - updates - } - None => SparseTrieUpdates::default(), - } + /// Converts a nibbles path to a B256, right-padding with zeros to 64 nibbles. + fn nibbles_to_padded_b256(path: &Nibbles) -> B256 { + let mut bytes = [0u8; 32]; + path.pack_to(&mut bytes); + B256::from(bytes) } - fn wipe(&mut self) { - self.upper_subtrie.wipe(); - for trie in &mut *self.lower_subtries { - trie.wipe(); - } - self.prefix_set = PrefixSetMut::all(); - self.updates = self.updates.is_some().then(SparseTrieUpdates::wiped); + /// Computes the proof target key and `min_len` for a blinded node error. + /// + /// Returns `(target_key, min_len)` where: + /// - `target_key` is `full_key` if `path` is a prefix of `full_path`, otherwise the padded path + /// - `min_len` is always based on `path.len()` + fn proof_target_for_path(full_key: B256, full_path: &Nibbles, path: &Nibbles) -> (B256, u8) { + let min_len = (path.len() as u8).min(64); + let target_key = + if full_path.starts_with(path) { full_key } else { Self::nibbles_to_padded_b256(path) }; + (target_key, min_len) } - fn clear(&mut self) { - self.upper_subtrie.clear(); - self.upper_subtrie.nodes.insert(Nibbles::default(), SparseNode::Empty); - for subtrie in &mut *self.lower_subtries { - subtrie.clear(); - } - self.prefix_set.clear(); - self.updates = None; - self.branch_node_masks.clear(); - #[cfg(feature = "trie-debug")] - self.debug_recorder.reset(); - // `update_actions_buffers` doesn't need to be cleared; we want to reuse the Vecs it has - // buffered, and all of those are already inherently cleared when they get used. + /// Creates a new revealed sparse trie from the given root node. + /// + /// This function initializes the internal structures and then reveals the root. + /// It is a convenient method to create a trie when you already have the root node available. + /// + /// # Arguments + /// + /// * `root` - The root node of the trie + /// * `masks` - Trie masks for root branch node + /// * `retain_updates` - Whether to track updates + /// + /// # Returns + /// + /// Self if successful, or an error if revealing fails. + pub fn from_root( + root: TrieNodeV2, + masks: Option, + retain_updates: bool, + ) -> SparseTrieResult { + Self::default().with_root(root, masks, retain_updates) } - fn find_leaf( - &self, - full_path: &Nibbles, - expected_value: Option<&Vec>, - ) -> Result { - // Inclusion proof - // - // First, do a quick check if the value exists in either the upper or lower subtrie's values - // map. We assume that if there exists a leaf node, then its value will be in the `values` - // map. - if let Some(actual_value) = core::iter::once(self.upper_subtrie.as_ref()) - .chain(self.lower_subtrie_for_path(full_path)) - .filter_map(|subtrie| subtrie.inner.values.get(full_path)) - .next() + /// Updates the value of a leaf node at the specified path. + pub fn update_leaf(&mut self, full_path: Nibbles, value: Vec) -> SparseTrieResult<()> { + debug_assert_eq!( + full_path.len(), + B256::len_bytes() * 2, + "update_leaf full_path must be 64 nibbles (32 bytes), got {} nibbles", + full_path.len() + ); + + trace!( + target: "trie::parallel_sparse", + ?full_path, + value_len = value.len(), + "Updating leaf", + ); + + if self.upper_subtrie.inner.values.contains_key(&full_path) { + self.prefix_set.insert(full_path); + self.upper_subtrie.inner.values.insert(full_path, value); + return Ok(()); + } + if let Some(subtrie) = self.lower_subtrie_for_path(&full_path) && + subtrie.inner.values.contains_key(&full_path) { - // We found the leaf, check if the value matches (if expected value was provided) - return expected_value - .is_none_or(|v| v == actual_value) - .then_some(LeafLookup::Exists) - .ok_or_else(|| LeafLookupError::ValueMismatch { - path: *full_path, - expected: expected_value.cloned(), - actual: actual_value.clone(), - }) + self.prefix_set.insert(full_path); + self.lower_subtrie_for_path_mut(&full_path) + .expect("subtrie exists") + .inner + .values + .insert(full_path, value); + return Ok(()); } - // If the value does not exist in the `values` map, then this means that the leaf either: - // - Does not exist in the trie - // - Is missing from the witness - // We traverse the trie to find the location where this leaf would have been, showing - // that it is not in the trie. Or we find a blinded node, showing that the witness is - // not complete. - let mut curr_path = Nibbles::new(); // start traversal from root - let mut curr_subtrie = self.upper_subtrie.as_ref(); - let mut curr_subtrie_is_upper = true; + self.upper_subtrie.inner.values.insert(full_path, value.clone()); - loop { - match curr_subtrie.nodes.get(&curr_path).unwrap() { - SparseNode::Empty => return Ok(LeafLookup::NonExistent), - SparseNode::Leaf { key, .. } => { - let mut found_full_path = curr_path; - found_full_path.extend(key); - assert!(&found_full_path != full_path, "target leaf {full_path:?} found, even though value wasn't in values hashmap"); - return Ok(LeafLookup::NonExistent) - } - SparseNode::Extension { key, .. } => { - if full_path.len() == curr_path.len() { - return Ok(LeafLookup::NonExistent) - } - curr_path.extend(key); - if !full_path.starts_with(&curr_path) { - return Ok(LeafLookup::NonExistent) - } + let mut new_nodes = Vec::new(); + let mut next = Some(Nibbles::default()); + + while let Some(current) = + next.as_mut().filter(|next| SparseSubtrieType::path_len_is_upper(next.len())) + { + let step_result = self.upper_subtrie.update_next_node(current, &full_path); + + if step_result.is_err() { + self.upper_subtrie.inner.values.remove(&full_path); + return step_result.map(|_| ()); + } + + match step_result? { + LeafUpdateStep::Continue => {} + LeafUpdateStep::Complete { inserted_nodes } => { + new_nodes.extend(inserted_nodes); + next = None; } - SparseNode::Branch { state_mask, blinded_mask, blinded_hashes, .. } => { - if full_path.len() == curr_path.len() { - return Ok(LeafLookup::NonExistent) - } - let nibble = full_path.get_unchecked(curr_path.len()); - if !state_mask.is_bit_set(nibble) { - return Ok(LeafLookup::NonExistent) - } - curr_path.push_unchecked(nibble); - if blinded_mask.is_bit_set(nibble) { - return Err(LeafLookupError::BlindedNode { - path: curr_path, - hash: blinded_hashes[nibble as usize], - }) - } + LeafUpdateStep::NodeNotFound => { + next = None; } } + } - // If we were previously looking at the upper trie, and the new path is in the - // lower trie, we need to pull out a ref to the lower trie. - if curr_subtrie_is_upper && - let Some(lower_subtrie) = self.lower_subtrie_for_path(&curr_path) - { - curr_subtrie = lower_subtrie; - curr_subtrie_is_upper = false; + for node_path in &new_nodes { + if SparseSubtrieType::path_len_is_upper(node_path.len()) { + continue; } - } - } - fn shrink_nodes_to(&mut self, size: usize) { - // Distribute the capacity across upper and lower subtries - // - // Always include upper subtrie, plus any lower subtries - let total_subtries = 1 + NUM_LOWER_SUBTRIES; - let size_per_subtrie = size / total_subtries; + let node = + self.upper_subtrie.nodes.remove(node_path).expect("node belongs to upper subtrie"); - // Shrink the upper subtrie - self.upper_subtrie.shrink_nodes_to(size_per_subtrie); + let leaf_value = if let SparseNode::Leaf { key, .. } = &node { + let mut leaf_full_path = *node_path; + leaf_full_path.extend(key); + Some(( + leaf_full_path, + self.upper_subtrie + .inner + .values + .remove(&leaf_full_path) + .expect("leaf nodes have associated values entries"), + )) + } else { + None + }; - // Shrink lower subtries (works for both revealed and blind with allocation) - for subtrie in &mut *self.lower_subtries { - subtrie.shrink_nodes_to(size_per_subtrie); + let subtrie = self.subtrie_for_path_mut(node_path); + + if let Some((leaf_full_path, value)) = leaf_value { + subtrie.inner.values.insert(leaf_full_path, value); + } + + subtrie.nodes.insert(*node_path, node); } - // shrink masks map - self.branch_node_masks.shrink_to(size); - } + if let Some(next_path) = next.filter(|n| !SparseSubtrieType::path_len_is_upper(n.len())) { + self.upper_subtrie.inner.values.remove(&full_path); - fn shrink_values_to(&mut self, size: usize) { - // Distribute the capacity across upper and lower subtries - // - // Always include upper subtrie, plus any lower subtries - let total_subtries = 1 + NUM_LOWER_SUBTRIES; - let size_per_subtrie = size / total_subtries; + let subtrie = self.subtrie_for_path_mut(&next_path); - // Shrink the upper subtrie - self.upper_subtrie.shrink_values_to(size_per_subtrie); + if subtrie.nodes.is_empty() { + subtrie.nodes.insert(subtrie.path, SparseNode::Empty); + } - // Shrink lower subtries (works for both revealed and blind with allocation) - for subtrie in &mut *self.lower_subtries { - subtrie.shrink_values_to(size_per_subtrie); + if let Err(e) = subtrie.update_leaf(full_path, value) { + if let Some(lower) = self.lower_subtrie_for_path_mut(&full_path) { + lower.inner.values.remove(&full_path); + } + return Err(e); + } } - } - /// O(1) size hint based on total node count (including hash stubs). - fn size_hint(&self) -> usize { - let upper_count = self.upper_subtrie.nodes.len(); - let lower_count: usize = self - .lower_subtries - .iter() - .filter_map(|s| s.as_revealed_ref()) - .map(|s| s.nodes.len()) - .sum(); - upper_count + lower_count - } + self.prefix_set.insert(full_path); - fn memory_size(&self) -> usize { - self.memory_size() + Ok(()) } - fn prune(&mut self, retained_leaves: &[Nibbles]) -> usize { - #[cfg(feature = "trie-debug")] - self.debug_recorder.reset(); - - let mut retained_leaves = retained_leaves.to_vec(); - retained_leaves.sort_unstable(); + /// Removes a leaf node at the specified path. + pub fn remove_leaf(&mut self, full_path: &Nibbles) -> SparseTrieResult<()> { + debug_assert_eq!( + full_path.len(), + B256::len_bytes() * 2, + "remove_leaf full_path must be 64 nibbles (32 bytes), got {} nibbles", + full_path.len() + ); - let mut effective_pruned_roots = Vec::::new(); - let mut stack: SmallVec<[Nibbles; 32]> = SmallVec::new(); - stack.push(Nibbles::default()); + trace!( + target: "trie::parallel_sparse", + ?full_path, + "Removing leaf", + ); - while let Some(path) = stack.pop() { - let Some(node) = - self.subtrie_for_path(&path).and_then(|subtrie| subtrie.nodes.get(&path).cloned()) - else { - continue; - }; + let leaf_path; + let leaf_subtrie_type; - match node { - SparseNode::Empty | SparseNode::Leaf { .. } => {} - SparseNode::Extension { key, state, .. } => { - let mut child = path; - child.extend(&key); + let mut branch_parent_path: Option = None; + let mut branch_parent_node: Option = None; - if has_retained_descendant(&retained_leaves, &child) { - stack.push(child); - continue; - } + let mut ext_grandparent_path: Option = None; + let mut ext_grandparent_node: Option = None; - // Root extension has no parent branch edge to blind; keep it as-is. - if path.is_empty() { - continue; - } + let mut curr_path = Nibbles::new(); + let mut curr_subtrie_type = SparseSubtrieType::Upper; - let Some(hash) = state.cached_hash() else { continue }; - self.subtrie_for_path_mut_untracked(&path) - .expect("node subtrie exists") - .nodes - .remove(&path); + let mut paths_to_mark_dirty = Vec::new(); - let parent_path = path.slice(0..path.len() - 1); - // Parent can live in a different subtrie when `path` is the root of a lower - // subtrie, so resolve it by `parent_path` rather than reusing `path`'s subtrie. - let SparseNode::Branch { blinded_mask, blinded_hashes, .. } = self - .subtrie_for_path_mut_untracked(&parent_path) - .expect("parent subtrie exists") - .nodes - .get_mut(&parent_path) - .expect("expected parent branch node") - else { - panic!("expected branch node at path {parent_path:?}"); - }; + loop { + let curr_subtrie = match curr_subtrie_type { + SparseSubtrieType::Upper => &mut self.upper_subtrie, + SparseSubtrieType::Lower(idx) => { + self.lower_subtries[idx].as_revealed_mut().expect("lower subtrie is revealed") + } + }; + let curr_node = curr_subtrie.nodes.get_mut(&curr_path).unwrap(); - let nibble = path.last().unwrap(); - blinded_mask.set_bit(nibble); - blinded_hashes[nibble as usize] = hash; - effective_pruned_roots.push(path); + match Self::find_next_to_leaf(&curr_path, curr_node, full_path) { + FindNextToLeafOutcome::NotFound => return Ok(()), + FindNextToLeafOutcome::BlindedNode(path) => { + return Err(SparseTrieErrorKind::BlindedNode(path).into()) } - SparseNode::Branch { state_mask, blinded_mask, blinded_hashes, .. } => { - let mut blinded_mask = blinded_mask; - let mut blinded_hashes = blinded_hashes; - for nibble in state_mask.iter() { - if blinded_mask.is_bit_set(nibble) { - continue; - } + FindNextToLeafOutcome::Found => { + leaf_path = curr_path; + leaf_subtrie_type = curr_subtrie_type; + break; + } + FindNextToLeafOutcome::ContinueFrom(next_path) => { + match curr_node { + SparseNode::Branch { .. } => { + paths_to_mark_dirty + .push((SparseSubtrieType::from_path(&curr_path), curr_path)); - let mut child = path; - child.push_unchecked(nibble); - if has_retained_descendant(&retained_leaves, &child) { - stack.push(child); - continue; + match (&branch_parent_path, &ext_grandparent_path) { + (Some(branch), Some(ext)) if branch.len() > ext.len() => { + ext_grandparent_path = None; + ext_grandparent_node = None; + } + _ => (), + }; + branch_parent_path = Some(curr_path); + branch_parent_node = Some(curr_node.clone()); + } + SparseNode::Extension { .. } => { + paths_to_mark_dirty + .push((SparseSubtrieType::from_path(&curr_path), curr_path)); + ext_grandparent_path = Some(curr_path); + ext_grandparent_node = Some(curr_node.clone()); + } + SparseNode::Empty | SparseNode::Leaf { .. } => { + unreachable!( + "find_next_to_leaf only continues to a branch or extension" + ) } + } - let Entry::Occupied(entry) = - self.subtrie_for_path_mut_untracked(&child).unwrap().nodes.entry(child) - else { - panic!("expected node at path {child:?}"); - }; + curr_path = next_path; - let Some(hash) = entry.get().cached_hash() else { - continue; - }; - entry.remove(); - blinded_mask.set_bit(nibble); - blinded_hashes[nibble as usize] = hash; - effective_pruned_roots.push(child); + let next_subtrie_type = SparseSubtrieType::from_path(&curr_path); + if matches!(curr_subtrie_type, SparseSubtrieType::Upper) && + matches!(next_subtrie_type, SparseSubtrieType::Lower(_)) + { + curr_subtrie_type = next_subtrie_type; } - - let SparseNode::Branch { - blinded_mask: old_blinded_mask, - blinded_hashes: old_blinded_hashes, - .. - } = self - .subtrie_for_path_mut_untracked(&path) - .unwrap() - .nodes - .get_mut(&path) - .unwrap() - else { - unreachable!("expected branch node at path {path:?}"); - }; - *old_blinded_mask = blinded_mask; - *old_blinded_hashes = blinded_hashes; } - } + }; } - Self::finalize_pruned_roots(self, effective_pruned_roots) - } - - fn update_leaves( - &mut self, - updates: &mut alloy_primitives::map::B256Map, - mut proof_required_fn: impl FnMut(B256, u8), - ) -> SparseTrieResult<()> { - use crate::{provider::NoRevealProvider, LeafUpdate}; + if let (Some(branch_path), Some(SparseNode::Branch { state_mask, blinded_mask, .. })) = + (&branch_parent_path, &branch_parent_node) + { + let mut check_mask = *state_mask; + let child_nibble = leaf_path.get_unchecked(branch_path.len()); + check_mask.unset_bit(child_nibble); - #[cfg(feature = "trie-debug")] - let recorded_updates: Vec<_> = - updates.iter().map(|(k, v)| (*k, LeafUpdateRecord::from(v))).collect(); - #[cfg(feature = "trie-debug")] - let mut recorded_proof_targets: Vec<(B256, u8)> = Vec::new(); + if check_mask.count_bits() == 1 { + let remaining_nibble = + check_mask.first_set_bit_index().expect("state mask is not empty"); - // Drain updates to avoid cloning keys while preserving the map's allocation. - // On success, entries remain removed; on blinded node failure, they're re-inserted. - let drained: Vec<_> = updates.drain().collect(); + if blinded_mask.is_bit_set(remaining_nibble) { + let mut path = *branch_path; + path.push_unchecked(remaining_nibble); + return Err(SparseTrieErrorKind::BlindedNode(path).into()); + } + } + } - for (key, update) in drained { - let full_path = Nibbles::unpack(key); + self.prefix_set.insert(*full_path); + let leaf_subtrie = match leaf_subtrie_type { + SparseSubtrieType::Upper => &mut self.upper_subtrie, + SparseSubtrieType::Lower(idx) => { + self.lower_subtries[idx].as_revealed_mut().expect("lower subtrie is revealed") + } + }; + leaf_subtrie.inner.values.remove(full_path); + for (subtrie_type, path) in paths_to_mark_dirty { + let node = match subtrie_type { + SparseSubtrieType::Upper => self.upper_subtrie.nodes.get_mut(&path), + SparseSubtrieType::Lower(idx) => self.lower_subtries[idx] + .as_revealed_mut() + .expect("lower subtrie is revealed") + .nodes + .get_mut(&path), + } + .expect("node exists"); - match update { - LeafUpdate::Changed(value) => { - if value.is_empty() { - // Removal: remove_leaf with NoRevealProvider is atomic - returns a - // retriable error before any mutations (via pre_validate_reveal_chain). - match self.remove_leaf(&full_path, NoRevealProvider) { - Ok(()) => {} - Err(e) => { - if let Some(path) = Self::get_retriable_path(&e) { - let (target_key, min_len) = - Self::proof_target_for_path(key, &full_path, &path); - proof_required_fn(target_key, min_len); - #[cfg(feature = "trie-debug")] - recorded_proof_targets.push((target_key, min_len)); - updates.insert(key, LeafUpdate::Changed(value)); - } else { - return Err(e); - } - } - } - } else { - // Update/insert: update_leaf is atomic - cleans up on error. - if let Err(e) = self.update_leaf(full_path, value.clone(), NoRevealProvider) - { - if let Some(path) = Self::get_retriable_path(&e) { - let (target_key, min_len) = - Self::proof_target_for_path(key, &full_path, &path); - proof_required_fn(target_key, min_len); - #[cfg(feature = "trie-debug")] - recorded_proof_targets.push((target_key, min_len)); - updates.insert(key, LeafUpdate::Changed(value)); - } else { - return Err(e); - } - } - } + match node { + SparseNode::Extension { state, .. } | SparseNode::Branch { state, .. } => { + *state = SparseNodeState::Dirty } - LeafUpdate::Touched => { - // Touched is read-only: check if path is accessible, request proof if blinded. - match self.find_leaf(&full_path, None) { - Err(LeafLookupError::BlindedNode { path, .. }) => { - let (target_key, min_len) = - Self::proof_target_for_path(key, &full_path, &path); - proof_required_fn(target_key, min_len); - #[cfg(feature = "trie-debug")] - recorded_proof_targets.push((target_key, min_len)); - updates.insert(key, LeafUpdate::Touched); - } - // Path is fully revealed (exists or proven non-existent), no action needed. - Ok(_) | Err(LeafLookupError::ValueMismatch { .. }) => {} - } + SparseNode::Empty | SparseNode::Leaf { .. } => { + unreachable!( + "only branch and extension nodes can be marked dirty when removing a leaf" + ) } } } + self.remove_node(&leaf_path); - #[cfg(feature = "trie-debug")] - self.debug_recorder.record(RecordedOp::UpdateLeaves { - updates: recorded_updates, - proof_targets: recorded_proof_targets, - }); + if leaf_path.is_empty() { + self.upper_subtrie.nodes.insert(leaf_path, SparseNode::Empty); + return Ok(()); + } - Ok(()) - } + if let ( + Some(branch_path), + &Some(SparseNode::Branch { mut state_mask, blinded_mask, ref blinded_hashes, .. }), + ) = (&branch_parent_path, &branch_parent_node) + { + let child_nibble = leaf_path.get_unchecked(branch_path.len()); + state_mask.unset_bit(child_nibble); - #[cfg(feature = "trie-debug")] - fn take_debug_recorder(&mut self) -> TrieDebugRecorder { - core::mem::take(&mut self.debug_recorder) - } + let new_branch_node = if state_mask.count_bits() == 1 { + let remaining_child_nibble = + state_mask.first_set_bit_index().expect("state mask is not empty"); + let mut remaining_child_path = *branch_path; + remaining_child_path.push_unchecked(remaining_child_nibble); - fn commit_updates( - &mut self, - updated: &HashMap, - removed: &HashSet, - ) { - // Sync branch_node_masks with what's being committed to DB. - // This ensures that on subsequent root() calls, the masks reflect the actual - // DB state, which is needed for correct removal detection. - self.branch_node_masks.reserve(updated.len()); - for (path, node) in updated { - self.branch_node_masks.insert( - *path, - BranchNodeMasks { tree_mask: node.tree_mask, hash_mask: node.hash_mask }, - ); - } - for path in removed { - self.branch_node_masks.remove(path); - } - } -} + trace!( + target: "trie::parallel_sparse", + ?leaf_path, + ?branch_path, + ?remaining_child_path, + "Branch node has only one child", + ); -impl ParallelSparseTrie { - /// Sets the thresholds that control when parallelism is used during operations. - pub const fn with_parallelism_thresholds(mut self, thresholds: ParallelismThresholds) -> Self { - self.parallelism_thresholds = thresholds; - self - } + if blinded_mask.is_bit_set(remaining_child_nibble) { + return Err(SparseTrieErrorKind::BlindedNode(remaining_child_path).into()); + } - /// Returns true if retaining updates is enabled for the overall trie. - const fn updates_enabled(&self) -> bool { - self.updates.is_some() - } + let remaining_child_node = self + .subtrie_for_path_mut(&remaining_child_path) + .nodes + .get(&remaining_child_path) + .unwrap(); - /// Returns true if parallelism should be enabled for revealing the given number of nodes. - /// Will always return false in nostd builds. - const fn is_reveal_parallelism_enabled(&self, num_nodes: usize) -> bool { - #[cfg(not(feature = "std"))] - { - let _ = num_nodes; - return false; - } + let (new_branch_node, remove_child) = Self::branch_changes_on_leaf_removal( + branch_path, + &remaining_child_path, + remaining_child_node, + ); - #[cfg(feature = "std")] - { - num_nodes >= self.parallelism_thresholds.min_revealed_nodes - } - } + if remove_child { + self.move_value_on_leaf_removal( + branch_path, + &new_branch_node, + &remaining_child_path, + ); + self.remove_node(&remaining_child_path); + } - /// Returns true if parallelism should be enabled for updating hashes with the given number - /// of changed keys. Will always return false in nostd builds. - const fn is_update_parallelism_enabled(&self, num_changed_keys: usize) -> bool { - #[cfg(not(feature = "std"))] - { - let _ = num_changed_keys; - return false; - } + if let Some(updates) = self.updates.as_mut() { + updates.updated_nodes.remove(branch_path); + updates.removed_nodes.insert(*branch_path); + } - #[cfg(feature = "std")] - { - num_changed_keys >= self.parallelism_thresholds.min_updated_nodes - } - } + new_branch_node + } else { + SparseNode::Branch { + state_mask, + blinded_mask, + blinded_hashes: blinded_hashes.clone(), + state: SparseNodeState::Dirty, + } + }; - /// Checks if an error is retriable (`BlindedNode` or `NodeNotFoundInProvider`) and extracts - /// the path if so. - /// - /// Both error types indicate that a node needs to be revealed before the operation can - /// succeed. `BlindedNode` occurs when traversing to a Hash node, while `NodeNotFoundInProvider` - /// occurs when `retain_updates` is enabled and an extension node's child needs revealing. - const fn get_retriable_path(e: &SparseTrieError) -> Option { - match e.kind() { - SparseTrieErrorKind::BlindedNode(path) | - SparseTrieErrorKind::NodeNotFoundInProvider { path } => Some(*path), - _ => None, - } - } + let branch_subtrie = self.subtrie_for_path_mut(branch_path); + branch_subtrie.nodes.insert(*branch_path, new_branch_node.clone()); + branch_parent_node = Some(new_branch_node); + }; - /// Converts a nibbles path to a B256, right-padding with zeros to 64 nibbles. - fn nibbles_to_padded_b256(path: &Nibbles) -> B256 { - let mut bytes = [0u8; 32]; - path.pack_to(&mut bytes); - B256::from(bytes) - } + if let (Some(ext_path), Some(SparseNode::Extension { key: shortkey, .. })) = + (ext_grandparent_path, &ext_grandparent_node) + { + let ext_subtrie = self.subtrie_for_path_mut(&ext_path); + let branch_path = branch_parent_path.as_ref().unwrap(); - /// Computes the proof target key and `min_len` for a blinded node error. - /// - /// Returns `(target_key, min_len)` where: - /// - `target_key` is `full_key` if `path` is a prefix of `full_path`, otherwise the padded path - /// - `min_len` is always based on `path.len()` - fn proof_target_for_path(full_key: B256, full_path: &Nibbles, path: &Nibbles) -> (B256, u8) { - let min_len = (path.len() as u8).min(64); - let target_key = - if full_path.starts_with(path) { full_key } else { Self::nibbles_to_padded_b256(path) }; - (target_key, min_len) - } + if let Some(new_ext_node) = Self::extension_changes_on_leaf_removal( + &ext_path, + shortkey, + branch_path, + branch_parent_node.as_ref().unwrap(), + ) { + ext_subtrie.nodes.insert(ext_path, new_ext_node.clone()); + self.move_value_on_leaf_removal(&ext_path, &new_ext_node, branch_path); + self.remove_node(branch_path); + } + } - /// Creates a new revealed sparse trie from the given root node. - /// - /// This function initializes the internal structures and then reveals the root. - /// It is a convenient method to create a trie when you already have the root node available. - /// - /// # Arguments - /// - /// * `root` - The root node of the trie - /// * `masks` - Trie masks for root branch node - /// * `retain_updates` - Whether to track updates - /// - /// # Returns - /// - /// Self if successful, or an error if revealing fails. - pub fn from_root( - root: TrieNodeV2, - masks: Option, - retain_updates: bool, - ) -> SparseTrieResult { - Self::default().with_root(root, masks, retain_updates) + Ok(()) } fn finalize_pruned_roots(&mut self, mut effective_pruned_roots: Vec) -> usize { @@ -2421,11 +2332,6 @@ impl SparseSubtrie { /// /// # Returns /// - /// Returns the path and masks of any blinded node revealed as a result of updating the leaf. - /// - /// If an update requires revealing a blinded node, an error is returned if the blinded - /// provider returns an error. - /// /// This method is atomic: if an error occurs during structural changes, all modifications /// are rolled back and the trie state is unchanged. pub fn update_leaf(&mut self, full_path: Nibbles, value: Vec) -> SparseTrieResult<()> { @@ -3500,10 +3406,8 @@ mod tests { SparseSubtrieType, }; use crate::{ - parallel::ChangedSubtrie, - provider::{DefaultTrieNodeProvider, NoRevealProvider}, - trie::SparseNodeState, - LeafLookup, LeafLookupError, SparseNode, SparseTrie, SparseTrieUpdates, + parallel::ChangedSubtrie, trie::SparseNodeState, LeafLookup, LeafLookupError, SparseNode, + SparseTrie, SparseTrieUpdates, }; use alloy_primitives::{ b256, hex, @@ -3654,7 +3558,7 @@ mod tests { leaves: impl IntoIterator)>, ) { for (path, value) in leaves { - trie.update_leaf(path, value, DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(path, value).unwrap(); } } @@ -4337,8 +4241,7 @@ mod tests { // Insert leaf_3 via update_leaf. This modifies the branch at [0x0] to add child // 0x2 and creates a fresh leaf node with hash: None in the lower subtrie. - let provider = NoRevealProvider; - trie.update_leaf(leaf_3_full_path, encode_account_value(3), provider).unwrap(); + trie.update_leaf(leaf_3_full_path, encode_account_value(3)).unwrap(); // Calculate subtrie indexes let subtrie_1_index = SparseSubtrieType::from_path(&leaf_1_path).lower_index().unwrap(); @@ -4525,11 +4428,9 @@ mod tests { .into_iter(), ); - let provider = NoRevealProvider; - // Remove the leaf with a full path of 0x537 let leaf_full_path = pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x3, 0x7])); - trie.remove_leaf(&leaf_full_path, provider).unwrap(); + trie.remove_leaf(&leaf_full_path).unwrap(); let upper_subtrie = &trie.upper_subtrie; let lower_subtrie_50 = trie.lower_subtries[0x50].as_revealed_ref().unwrap(); @@ -4579,11 +4480,9 @@ mod tests { .insert(Nibbles::default(), BranchNodeCompact::new(0b11, 0, 0, vec![], None)); } - let provider = NoRevealProvider; - // Remove the leaf with a full path of 0x012 let leaf_full_path = pad_nibbles_right(Nibbles::from_nibbles([0x0, 0x1, 0x2])); - trie.remove_leaf(&leaf_full_path, provider).unwrap(); + trie.remove_leaf(&leaf_full_path).unwrap(); let upper_subtrie = &trie.upper_subtrie; @@ -4632,11 +4531,9 @@ mod tests { .into_iter(), ); - let provider = NoRevealProvider; - // Remove the leaf with a full path of 0x5012 let leaf_full_path = pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x0, 0x1, 0x2])); - trie.remove_leaf(&leaf_full_path, provider).unwrap(); + trie.remove_leaf(&leaf_full_path).unwrap(); let upper_subtrie = &trie.upper_subtrie; @@ -4683,11 +4580,9 @@ mod tests { .into_iter(), ); - let provider = NoRevealProvider; - // Remove the leaf with a full path of 0x2034 let leaf_full_path = pad_nibbles_right(Nibbles::from_nibbles([0x2, 0x0, 0x3, 0x4])); - trie.remove_leaf(&leaf_full_path, provider).unwrap(); + trie.remove_leaf(&leaf_full_path).unwrap(); let upper_subtrie = &trie.upper_subtrie; @@ -4762,8 +4657,6 @@ mod tests { .into_iter(), ); - let provider = NoRevealProvider; - // Verify initial state - the lower subtrie's path should be 0x123 let lower_subtrie_root_path = Nibbles::from_nibbles([0x1, 0x2, 0x3]); assert_matches!( @@ -4774,7 +4667,7 @@ mod tests { // Remove the leaf at 0x1233 let leaf_full_path = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x3])); - trie.remove_leaf(&leaf_full_path, provider).unwrap(); + trie.remove_leaf(&leaf_full_path).unwrap(); // After removal: // 1. The branch at 0x123 should become an extension to 0x12345 @@ -4821,7 +4714,7 @@ mod tests { .into_iter(), ); - // Create a mock provider that will reveal the blinded leaf + // Create the revealed leaf used by the test setup. let revealed_leaf = create_leaf_node(leaf_key([0x3, 0x4], 63).to_vec(), 42); let mut encoded = Vec::new(); revealed_leaf.encode(&mut encoded); @@ -4829,7 +4722,7 @@ mod tests { // Try removing the leaf with a full path of 0x012, this should fail because the leaf is // blinded let leaf_full_path = pad_nibbles_right(Nibbles::from_nibbles([0x0, 0x1, 0x2])); - let Err(err) = trie.remove_leaf(&leaf_full_path, NoRevealProvider) else { + let Err(err) = trie.remove_leaf(&leaf_full_path) else { panic!("expected error"); }; assert_matches!(err.kind(), SparseTrieErrorKind::BlindedNode(path) if *path == Nibbles::from_nibbles([0x1])); @@ -4841,7 +4734,7 @@ mod tests { masks: None, }]) .unwrap(); - trie.remove_leaf(&leaf_full_path, NoRevealProvider).unwrap(); + trie.remove_leaf(&leaf_full_path).unwrap(); let upper_subtrie = &trie.upper_subtrie; @@ -4871,11 +4764,9 @@ mod tests { SparseNode::new_leaf(pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3]))), ))); - let provider = NoRevealProvider; - // Remove the leaf with a full key of 0x123 let leaf_full_path = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3])); - trie.remove_leaf(&leaf_full_path, provider).unwrap(); + trie.remove_leaf(&leaf_full_path).unwrap(); let upper_subtrie = &trie.upper_subtrie; @@ -4965,21 +4856,16 @@ mod tests { .into_iter(), ); - let provider = NoRevealProvider; - // Remove a leaf which does not exist; this should have no effect. - trie.remove_leaf( - &pad_nibbles_right(Nibbles::from_nibbles([0x0, 0x1, 0x2, 0x3, 0x4, 0xF])), - provider, - ) - .unwrap(); + trie.remove_leaf(&pad_nibbles_right(Nibbles::from_nibbles([0x0, 0x1, 0x2, 0x3, 0x4, 0xF]))) + .unwrap(); for (path, node) in trie.all_nodes() { assert!(node.cached_hash().is_some(), "path {path:?} should still have a hash"); } // Remove the leaf at path 0x01234 let leaf_full_path = pad_nibbles_right(Nibbles::from_nibbles([0x0, 0x1, 0x2, 0x3, 0x4])); - trie.remove_leaf(&leaf_full_path, provider).unwrap(); + trie.remove_leaf(&leaf_full_path).unwrap(); let upper_subtrie = &trie.upper_subtrie; let lower_subtrie_10 = trie.lower_subtries[0x01].as_revealed_ref().unwrap(); @@ -5194,11 +5080,9 @@ mod tests { Default::default(), paths.clone(), ); - - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::default().with_updates(true); for path in &paths { - sparse.update_leaf(*path, value_encoded(), &provider).unwrap(); + sparse.update_leaf(*path, value_encoded()).unwrap(); } let sparse_root = sparse.root(); let sparse_updates = sparse.take_updates(); @@ -5310,7 +5194,6 @@ mod tests { #[test] fn sparse_trie_remove_leaf() { let ctx = ParallelSparseTrieTestContext; - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::default(); let value = alloy_rlp::encode_fixed_size(&U256::ZERO).to_vec(); @@ -5403,10 +5286,7 @@ mod tests { ); sparse - .remove_leaf( - &pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3])), - &provider, - ) + .remove_leaf(&pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3]))) .unwrap(); // Extension (Key = 5) @@ -5465,10 +5345,7 @@ mod tests { ); sparse - .remove_leaf( - &pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1])), - &provider, - ) + .remove_leaf(&pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]))) .unwrap(); // Extension (Key = 5) @@ -5512,10 +5389,7 @@ mod tests { ); sparse - .remove_leaf( - &pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2])), - &provider, - ) + .remove_leaf(&pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2]))) .unwrap(); // Extension (Key = 5) @@ -5556,10 +5430,7 @@ mod tests { ); sparse - .remove_leaf( - &pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0])), - &provider, - ) + .remove_leaf(&pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0]))) .unwrap(); // Extension (Key = 5) @@ -5586,10 +5457,7 @@ mod tests { ); sparse - .remove_leaf( - &pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3])), - &provider, - ) + .remove_leaf(&pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]))) .unwrap(); // Leaf (Path = 53302...) @@ -5607,10 +5475,7 @@ mod tests { ); sparse - .remove_leaf( - &pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2])), - &provider, - ) + .remove_leaf(&pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2]))) .unwrap(); // Empty @@ -5638,8 +5503,6 @@ mod tests { TrieMask::new(0b11), None, )); - - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::from_root( branch.clone(), Some(BranchNodeMasks { @@ -5675,7 +5538,7 @@ mod tests { // Removing a blinded leaf should result in an error assert_matches!( - sparse.remove_leaf(&pad_nibbles_right(Nibbles::from_nibbles([0x0])), &provider).map_err(|e| e.into_kind()), + sparse.remove_leaf(&pad_nibbles_right(Nibbles::from_nibbles([0x0]))).map_err(|e| e.into_kind()), Err(SparseTrieErrorKind::BlindedNode(path)) if path == Nibbles::from_nibbles([0x0]) ); } @@ -5695,8 +5558,6 @@ mod tests { TrieMask::new(0b11), None, )); - - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::from_root( branch.clone(), Some(BranchNodeMasks { @@ -5733,7 +5594,7 @@ mod tests { // Removing a non-existent leaf should be a noop let sparse_old = sparse.clone(); assert_matches!( - sparse.remove_leaf(&pad_nibbles_right(Nibbles::from_nibbles([0x2])), &provider), + sparse.remove_leaf(&pad_nibbles_right(Nibbles::from_nibbles([0x2]))), Ok(()) ); assert_eq!(sparse, sparse_old); @@ -5749,7 +5610,6 @@ mod tests { fn test(updates: Vec<(BTreeMap, BTreeSet)>) { { let mut state = BTreeMap::default(); - let default_provider = DefaultTrieNodeProvider; let provider_factory = create_test_provider_factory(); let mut sparse = ParallelSparseTrie::default().with_updates(true); @@ -5759,7 +5619,7 @@ mod tests { let account = account.into_trie_account(EMPTY_ROOT_HASH); let mut account_rlp = Vec::new(); account.encode(&mut account_rlp); - sparse.update_leaf(key, account_rlp, &default_provider).unwrap(); + sparse.update_leaf(key, account_rlp).unwrap(); } // We need to clone the sparse trie, so that all updated branch nodes are // preserved, and not only those that were changed after the last call to @@ -5808,7 +5668,7 @@ mod tests { // that the sparse trie root still matches the hash builder root for key in &keys_to_delete { state.remove(key).unwrap(); - sparse.remove_leaf(key, &default_provider).unwrap(); + sparse.remove_leaf(key).unwrap(); } // We need to clone the sparse trie, so that all updated branch nodes are @@ -5898,7 +5758,6 @@ mod tests { #[test] fn sparse_trie_two_leaves_at_lower_roots() { - let provider = DefaultTrieNodeProvider; let mut trie = ParallelSparseTrie::default().with_updates(true); let key_50 = Nibbles::unpack(hex!( "0x5000000000000000000000000000000000000000000000000000000000000000" @@ -5912,11 +5771,11 @@ mod tests { account.encode(&mut account_rlp); // Add a leaf and calculate the root. - trie.update_leaf(key_50, account_rlp.clone(), &provider).unwrap(); + trie.update_leaf(key_50, account_rlp.clone()).unwrap(); trie.root(); // Add a second leaf and assert that the root is the expected value. - trie.update_leaf(key_51, account_rlp.clone(), &provider).unwrap(); + trie.update_leaf(key_51, account_rlp.clone()).unwrap(); let expected_root = hex!("0xdaf0ef9f91a2f179bb74501209effdb5301db1697bcab041eca2234b126e25de"); @@ -5956,8 +5815,6 @@ mod tests { Default::default(), [Nibbles::default()], ); - - let provider = DefaultTrieNodeProvider; let masks = match ( branch_node_hash_masks.get(&Nibbles::default()).copied(), branch_node_tree_masks.get(&Nibbles::default()).copied(), @@ -6005,7 +5862,7 @@ mod tests { ); // Insert the leaf for the second key - sparse.update_leaf(key2(), value_encoded(), &provider).unwrap(); + sparse.update_leaf(key2(), value_encoded()).unwrap(); // Check that the branch node was updated and another nibble was set assert_matches!( @@ -6076,8 +5933,6 @@ mod tests { Default::default(), [Nibbles::default()], ); - - let provider = DefaultTrieNodeProvider; let masks = match ( branch_node_hash_masks.get(&Nibbles::default()).copied(), branch_node_tree_masks.get(&Nibbles::default()).copied(), @@ -6126,7 +5981,7 @@ mod tests { ); // Remove the leaf for the first key - sparse.remove_leaf(&key1(), &provider).unwrap(); + sparse.remove_leaf(&key1()).unwrap(); // Check that the branch node was turned into an extension node assert_eq!( @@ -6202,8 +6057,6 @@ mod tests { nodes.sort_unstable_by(|a, b| reth_trie_common::depth_first_cmp(&a.0, &b.0)); let nodes = ProofTrieNodeV2::from_sorted_trie_nodes(nodes); - - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::from_root(nodes[0].node.clone(), nodes[0].masks, false).unwrap(); @@ -6214,7 +6067,7 @@ mod tests { ); // Insert the leaf with a different prefix - sparse.update_leaf(key3(), value_encoded(), &provider).unwrap(); + sparse.update_leaf(key3(), value_encoded()).unwrap(); // Check that the extension node was turned into a branch node assert_eq!( @@ -6276,7 +6129,7 @@ mod tests { // First add leaf 0x1345 - this should create a leaf in upper trie at 0x let (leaf1_path, value1) = ctx.create_test_leaf([0x1, 0x3, 0x4, 0x5], 1); - trie.update_leaf(leaf1_path, value1.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf1_path, value1.clone()).unwrap(); // Verify upper trie has a leaf at the root with key 1345 ctx.assert_upper_subtrie(&trie) @@ -6288,7 +6141,7 @@ mod tests { // Add leaf 0x1234 - this should go first in the upper subtrie let (leaf2_path, value2) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4], 2); - trie.update_leaf(leaf2_path, value2.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf2_path, value2.clone()).unwrap(); // Upper trie should now have a branch at 0x1 ctx.assert_upper_subtrie(&trie) @@ -6298,7 +6151,7 @@ mod tests { // Add leaf 0x1245 - this should cause a branch and create the 0x12 subtrie let (leaf3_path, value3) = ctx.create_test_leaf([0x1, 0x2, 0x4, 0x5], 3); - trie.update_leaf(leaf3_path, value3.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf3_path, value3.clone()).unwrap(); // Verify lower subtrie at 0x12 exists with correct structure ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2])) @@ -6310,7 +6163,7 @@ mod tests { // Add leaf 0x1334 - this should create another lower subtrie let (leaf4_path, value4) = ctx.create_test_leaf([0x1, 0x3, 0x3, 0x4], 4); - trie.update_leaf(leaf4_path, value4.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf4_path, value4.clone()).unwrap(); // Verify lower subtrie at 0x13 exists with correct values ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x3])) @@ -6353,7 +6206,7 @@ mod tests { // First insert a leaf that ends exactly at the boundary (2 nibbles) let (first_leaf_path, first_value) = ctx.create_test_leaf([0x1, 0x2, 0x2, 0x4], 1); - trie.update_leaf(first_leaf_path, first_value.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(first_leaf_path, first_value.clone()).unwrap(); // In an empty trie, the first leaf becomes the root, regardless of path length ctx.assert_upper_subtrie(&trie) @@ -6366,7 +6219,7 @@ mod tests { // Now insert another leaf that shares the same 2-nibble prefix let (second_leaf_path, second_value) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4], 2); - trie.update_leaf(second_leaf_path, second_value.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(second_leaf_path, second_value.clone()).unwrap(); // Now both leaves should be in a lower subtrie at index [0x1, 0x2] ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2])) @@ -6428,7 +6281,7 @@ mod tests { let updated_path = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4])); let (_, updated_value) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4], 100); - trie.update_leaf(updated_path, updated_value.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(updated_path, updated_value.clone()).unwrap(); // Verify the subtrie structure is maintained and value is updated // The branch structure should remain the same and all values should be present @@ -6442,7 +6295,7 @@ mod tests { // Add a new leaf that extends an existing branch let (new_leaf_path, new_leaf_value) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x6], 200); - trie.update_leaf(new_leaf_path, new_leaf_value.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(new_leaf_path, new_leaf_value.clone()).unwrap(); // Verify the branch at [0x1, 0x2, 0x3] now has an additional child ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2])) @@ -6771,8 +6624,8 @@ mod tests { let (leaf1_path, value1) = ctx.create_test_leaf([0xA, 0xB, 0xC, 0xD, 0xE, 0xF], 1); let (leaf2_path, value2) = ctx.create_test_leaf([0xA, 0xB, 0xD, 0xE, 0xF, 0x0], 2); - trie.update_leaf(leaf1_path, value1.clone(), DefaultTrieNodeProvider).unwrap(); - trie.update_leaf(leaf2_path, value2.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf1_path, value1.clone()).unwrap(); + trie.update_leaf(leaf2_path, value2.clone()).unwrap(); // Verify upper trie structure ctx.assert_upper_subtrie(&trie) @@ -6821,9 +6674,9 @@ mod tests { let (leaf2_path, value2) = ctx.create_test_leaf([0x2, 0x3, 0x4, 0x5], 2); let (leaf3_path, value3) = ctx.create_test_leaf([0x2, 0x3, 0x5, 0x6], 3); - trie.update_leaf(leaf1_path, value1, DefaultTrieNodeProvider).unwrap(); - trie.update_leaf(leaf2_path, value2, DefaultTrieNodeProvider).unwrap(); - trie.update_leaf(leaf3_path, value3, DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf1_path, value1).unwrap(); + trie.update_leaf(leaf2_path, value2).unwrap(); + trie.update_leaf(leaf3_path, value3).unwrap(); // Verify initial structure has branch at root ctx.assert_upper_subtrie(&trie).has_branch(&Nibbles::default(), &[0x1, 0x2]); @@ -6836,9 +6689,9 @@ mod tests { // Clear and add new leaves let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); - trie.update_leaf(new_leaf1_path, new_value1.clone(), DefaultTrieNodeProvider).unwrap(); - trie.update_leaf(new_leaf2_path, new_value2.clone(), DefaultTrieNodeProvider).unwrap(); - trie.update_leaf(new_leaf3_path, new_value3.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(new_leaf1_path, new_value1.clone()).unwrap(); + trie.update_leaf(new_leaf2_path, new_value2.clone()).unwrap(); + trie.update_leaf(new_leaf3_path, new_value3.clone()).unwrap(); // Verify new structure has extension ctx.assert_upper_subtrie(&trie) @@ -6882,9 +6735,9 @@ mod tests { let (leaf2_path, value2) = ctx.create_test_leaf([0x2, 0x3, 0x4, 0x5], 2); let (leaf3_path, value3) = ctx.create_test_leaf([0x2, 0x3, 0x5, 0x6], 3); - trie.update_leaf(leaf1_path, value1, DefaultTrieNodeProvider).unwrap(); - trie.update_leaf(leaf2_path, value2.clone(), DefaultTrieNodeProvider).unwrap(); - trie.update_leaf(leaf3_path, value3.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf1_path, value1).unwrap(); + trie.update_leaf(leaf2_path, value2.clone()).unwrap(); + trie.update_leaf(leaf3_path, value3.clone()).unwrap(); // Verify upper trie structure ctx.assert_upper_subtrie(&trie) @@ -6942,7 +6795,7 @@ mod tests { // Step 1: Add first leaf - initially stored as leaf in upper trie let (leaf1_path, value1) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4, 0x5], 1); - trie.update_leaf(leaf1_path, value1.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf1_path, value1.clone()).unwrap(); // Verify leaf node in upper trie (optimized single-leaf case) ctx.assert_upper_subtrie(&trie) @@ -6954,7 +6807,7 @@ mod tests { // Step 2: Add leaf at 0x12346 - creates branch at 0x1234 let (leaf2_path, value2) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4, 0x6], 2); - trie.update_leaf(leaf2_path, value2.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf2_path, value2.clone()).unwrap(); // Verify extension now goes to 0x1234 ctx.assert_upper_subtrie(&trie) @@ -6972,7 +6825,7 @@ mod tests { // Step 3: Add leaf at 0x1235 - creates branch at 0x123 let (leaf3_path, value3) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x5], 3); - trie.update_leaf(leaf3_path, value3.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf3_path, value3.clone()).unwrap(); // Verify extension now goes to 0x123 ctx.assert_upper_subtrie(&trie) @@ -6991,7 +6844,7 @@ mod tests { // Step 4: Add leaf at 0x124 - creates branch at 0x12 (subtrie root) let (leaf4_path, value4) = ctx.create_test_leaf([0x1, 0x2, 0x4], 4); - trie.update_leaf(leaf4_path, value4.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf4_path, value4.clone()).unwrap(); // Verify extension now goes to 0x12 ctx.assert_upper_subtrie(&trie) @@ -7036,8 +6889,8 @@ mod tests { let (leaf1_path, value1) = ctx.create_test_leaf(&path1_nibbles, 1); let (leaf2_path, value2) = ctx.create_test_leaf(&path2_nibbles, 2); - trie.update_leaf(leaf1_path, value1.clone(), DefaultTrieNodeProvider).unwrap(); - trie.update_leaf(leaf2_path, value2.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf1_path, value1.clone()).unwrap(); + trie.update_leaf(leaf2_path, value2.clone()).unwrap(); // The common prefix of 63 F's will create a very long extension let extension_key = vec![0xF; 63]; @@ -7169,7 +7022,7 @@ mod tests { 218, 223, 145, 158, 225, 240, 227, 203, 155, 98, 211, 244, 176, 44, ]; - trie.update_leaf(leaf_full_path, leaf_new_value.clone(), DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf_full_path, leaf_new_value.clone()).unwrap(); // Sanity checks before calculating the root assert_eq!( @@ -7187,12 +7040,11 @@ mod tests { #[test] fn find_leaf_existing_leaf() { // Create a simple trie with one leaf - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::default(); let path = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3])); let value = b"test_value".to_vec(); - sparse.update_leaf(path, value.clone(), &provider).unwrap(); + sparse.update_leaf(path, value.clone()).unwrap(); // Check that the leaf exists let result = sparse.find_leaf(&path, None); @@ -7206,13 +7058,12 @@ mod tests { #[test] fn find_leaf_value_mismatch() { // Create a simple trie with one leaf - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::default(); let path = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3])); let value = b"test_value".to_vec(); let wrong_value = b"wrong_value".to_vec(); - sparse.update_leaf(path, value, &provider).unwrap(); + sparse.update_leaf(path, value).unwrap(); // Check with wrong expected value let result = sparse.find_leaf(&path, Some(&wrong_value)); @@ -7244,10 +7095,9 @@ mod tests { #[test] fn find_leaf_exists_no_value_check() { - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::default(); let path = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4])); - sparse.update_leaf(path, encode_account_value(0), &provider).unwrap(); + sparse.update_leaf(path, encode_account_value(0)).unwrap(); let result = sparse.find_leaf(&path, None); assert_matches!(result, Ok(LeafLookup::Exists)); @@ -7255,11 +7105,10 @@ mod tests { #[test] fn find_leaf_exists_with_value_check_ok() { - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::default(); let path = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4])); let value = encode_account_value(0); - sparse.update_leaf(path, value.clone(), &provider).unwrap(); + sparse.update_leaf(path, value.clone()).unwrap(); let result = sparse.find_leaf(&path, Some(&value)); assert_matches!(result, Ok(LeafLookup::Exists)); @@ -7267,14 +7116,13 @@ mod tests { #[test] fn find_leaf_exclusion_branch_divergence() { - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::default(); let path1 = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4])); // Creates branch at 0x12 let path2 = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x5, 0x6])); // Belongs to same branch let search_path = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x7, 0x8])); // Diverges at nibble 7 - sparse.update_leaf(path1, encode_account_value(0), &provider).unwrap(); - sparse.update_leaf(path2, encode_account_value(1), &provider).unwrap(); + sparse.update_leaf(path1, encode_account_value(0)).unwrap(); + sparse.update_leaf(path2, encode_account_value(1)).unwrap(); let result = sparse.find_leaf(&search_path, None); assert_matches!(result, Ok(LeafLookup::NonExistent)) @@ -7282,14 +7130,13 @@ mod tests { #[test] fn find_leaf_exclusion_extension_divergence() { - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::default(); // This will create an extension node at root with key 0x12 let path1 = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4, 0x5, 0x6])); // This path diverges from the extension key let search_path = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x7, 0x8])); - sparse.update_leaf(path1, encode_account_value(0), &provider).unwrap(); + sparse.update_leaf(path1, encode_account_value(0)).unwrap(); let result = sparse.find_leaf(&search_path, None); assert_matches!(result, Ok(LeafLookup::NonExistent)) @@ -7297,12 +7144,11 @@ mod tests { #[test] fn find_leaf_exclusion_leaf_divergence() { - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::default(); let existing_leaf_path = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4])); let search_path = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4, 0x5, 0x6])); - sparse.update_leaf(existing_leaf_path, encode_account_value(0), &provider).unwrap(); + sparse.update_leaf(existing_leaf_path, encode_account_value(0)).unwrap(); let result = sparse.find_leaf(&search_path, None); assert_matches!(result, Ok(LeafLookup::NonExistent)) @@ -7310,14 +7156,13 @@ mod tests { #[test] fn find_leaf_exclusion_path_ends_at_branch() { - let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::default(); let path1 = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4])); // Creates branch at 0x12 let path2 = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x5, 0x6])); let search_path = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2])); // Path of the branch itself - sparse.update_leaf(path1, encode_account_value(0), &provider).unwrap(); - sparse.update_leaf(path2, encode_account_value(1), &provider).unwrap(); + sparse.update_leaf(path1, encode_account_value(0)).unwrap(); + sparse.update_leaf(path2, encode_account_value(1)).unwrap(); let result = sparse.find_leaf(&search_path, None); assert_matches!(result, Ok(LeafLookup::NonExistent)); @@ -7467,10 +7312,10 @@ mod tests { trie.reveal_nodes(&mut proof_nodes).unwrap(); // Update the leaf in order to reveal it in the trie - trie.update_leaf(leaf_nibbles, leaf_value, NoRevealProvider).unwrap(); + trie.update_leaf(leaf_nibbles, leaf_value).unwrap(); // Now try deleting the leaf - let Err(err) = trie.remove_leaf(&leaf_nibbles, NoRevealProvider) else { + let Err(err) = trie.remove_leaf(&leaf_nibbles) else { panic!("expected blinded node error"); }; assert_matches!(err.kind(), SparseTrieErrorKind::BlindedNode(path) if path == &Nibbles::from_nibbles([0x3, 0x1, 0xc])); @@ -7483,7 +7328,7 @@ mod tests { .unwrap(); // Now remove the leaf again, this should succeed - trie.remove_leaf(&leaf_nibbles, NoRevealProvider).unwrap(); + trie.remove_leaf(&leaf_nibbles).unwrap(); // Compute the root to trigger updates let _ = trie.root(); @@ -7570,8 +7415,6 @@ mod tests { /// missing values stored in `upper_subtrie.inner.values`. #[test] fn test_get_leaf_value_upper_subtrie_via_update_leaf() { - let provider = NoRevealProvider; - // Create an empty trie with an empty root let mut trie = ParallelSparseTrie::default() .with_root(TrieNodeV2::EmptyRoot, None, false) @@ -7583,7 +7426,7 @@ mod tests { // Insert the leaf - since the trie is empty, the leaf node will be created // at the root level (depth 0), which is in the upper subtrie - trie.update_leaf(full_path, value.clone(), provider).unwrap(); + trie.update_leaf(full_path, value.clone()).unwrap(); // Verify the value is stored in upper_subtrie (where update_leaf puts it) assert!( @@ -7601,8 +7444,6 @@ mod tests { /// Test that `get_leaf_value` works for values in both upper and lower subtries. #[test] fn test_get_leaf_value_upper_and_lower_subtries() { - let provider = NoRevealProvider; - // Create an empty trie let mut trie = ParallelSparseTrie::default() .with_root(TrieNodeV2::EmptyRoot, None, false) @@ -7611,12 +7452,12 @@ mod tests { // Insert first leaf - will be at root level (upper subtrie) let path1 = pad_nibbles_right(Nibbles::from_nibbles([0x0, 0xA])); let value1 = encode_account_value(1); - trie.update_leaf(path1, value1.clone(), provider).unwrap(); + trie.update_leaf(path1, value1.clone()).unwrap(); // Insert second leaf with different prefix - creates a branch let path2 = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0xB])); let value2 = encode_account_value(2); - trie.update_leaf(path2, value2.clone(), provider).unwrap(); + trie.update_leaf(path2, value2.clone()).unwrap(); // Both values should be retrievable assert_eq!(trie.get_leaf_value(&path1), Some(&value1)); @@ -7626,8 +7467,6 @@ mod tests { /// Test that `get_leaf_value` works for storage tries which are often very sparse. #[test] fn test_get_leaf_value_sparse_storage_trie() { - let provider = NoRevealProvider; - // Simulate a storage trie with a single slot let mut trie = ParallelSparseTrie::default() .with_root(TrieNodeV2::EmptyRoot, None, false) @@ -7636,7 +7475,7 @@ mod tests { // Single storage slot - leaf will be at root (depth 0) let slot_path = pad_nibbles_right(Nibbles::from_nibbles([0x2, 0x9])); let slot_value = alloy_rlp::encode(U256::from(12345)); - trie.update_leaf(slot_path, slot_value.clone(), provider).unwrap(); + trie.update_leaf(slot_path, slot_value.clone()).unwrap(); // Value should be retrievable assert_eq!(trie.get_leaf_value(&slot_path), Some(&slot_value)); @@ -7647,10 +7486,6 @@ mod tests { // Regression test: when a leaf has an empty suffix key (full path == node path), // the value must be removed when that path becomes a pruned root. // This catches the bug where is_strict_descendant fails to remove p == pruned_root. - - use crate::provider::DefaultTrieNodeProvider; - - let provider = DefaultTrieNodeProvider; let mut parallel = ParallelSparseTrie::default(); // Large value to ensure nodes have hashes (RLP >= 32 bytes) @@ -7671,7 +7506,6 @@ mod tests { .update_leaf( pad_nibbles_right(Nibbles::from_nibbles([i, 0x1, 0x2, 0x3, 0x4, 0x5])), value.clone(), - &provider, ) .unwrap(); } @@ -7707,7 +7541,6 @@ mod tests { #[test] fn test_prune_preserves_root_hash() { - let provider = DefaultTrieNodeProvider; let mut trie = ParallelSparseTrie::default(); let value = large_account_value(); @@ -7717,7 +7550,6 @@ mod tests { trie.update_leaf( pad_nibbles_right(Nibbles::from_nibbles([i, j, 0x3, 0x4, 0x5, 0x6])), value.clone(), - &provider, ) .unwrap(); } @@ -7731,16 +7563,11 @@ mod tests { #[test] fn test_prune_single_leaf_trie() { - let provider = DefaultTrieNodeProvider; let mut trie = ParallelSparseTrie::default(); let value = large_account_value(); - trie.update_leaf( - pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4])), - value, - &provider, - ) - .unwrap(); + trie.update_leaf(pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4])), value) + .unwrap(); let root_before = trie.root(); let nodes_before = trie.size_hint(); @@ -7754,7 +7581,6 @@ mod tests { #[test] fn test_prune_root_hash_preserved() { - let provider = DefaultTrieNodeProvider; let mut trie = ParallelSparseTrie::default(); // Create two 64-nibble paths that differ only in the first nibble @@ -7762,8 +7588,8 @@ mod tests { let key2 = Nibbles::unpack(B256::repeat_byte(0x11)); let large_value = large_account_value(); - trie.update_leaf(key1, large_value.clone(), &provider).unwrap(); - trie.update_leaf(key2, large_value, &provider).unwrap(); + trie.update_leaf(key1, large_value.clone()).unwrap(); + trie.update_leaf(key2, large_value).unwrap(); let root_before = trie.root(); @@ -7774,7 +7600,6 @@ mod tests { #[test] fn test_prune_mixed_embedded_and_hashed() { - let provider = DefaultTrieNodeProvider; let mut trie = ParallelSparseTrie::default(); let large_value = large_account_value(); @@ -7782,12 +7607,8 @@ mod tests { for i in 0..8u8 { let value = if i < 4 { large_value.clone() } else { small_value.clone() }; - trie.update_leaf( - pad_nibbles_right(Nibbles::from_nibbles([i, 0x1, 0x2, 0x3])), - value, - &provider, - ) - .unwrap(); + trie.update_leaf(pad_nibbles_right(Nibbles::from_nibbles([i, 0x1, 0x2, 0x3])), value) + .unwrap(); } let root_before = trie.root(); @@ -7797,8 +7618,6 @@ mod tests { #[test] fn test_prune_all_lower_subtries() { - let provider = DefaultTrieNodeProvider; - let large_value = large_account_value(); let mut keys = Vec::new(); @@ -7813,7 +7632,7 @@ mod tests { let mut trie = ParallelSparseTrie::default(); for key in &keys { - trie.update_leaf(*key, large_value.clone(), &provider).unwrap(); + trie.update_leaf(*key, large_value.clone()).unwrap(); } let root_before = trie.root(); @@ -7830,7 +7649,6 @@ mod tests { #[test] fn test_prune_keeps_only_hot_paths() { - let provider = DefaultTrieNodeProvider; let mut trie = ParallelSparseTrie::default(); let key_keep = pad_nibbles_right(Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4])); @@ -7838,9 +7656,9 @@ mod tests { let key_drop_2 = pad_nibbles_right(Nibbles::from_nibbles([0x9, 0x2, 0x3, 0x4])); let value = large_account_value(); - trie.update_leaf(key_keep, value.clone(), &provider).unwrap(); - trie.update_leaf(key_drop_1, value.clone(), &provider).unwrap(); - trie.update_leaf(key_drop_2, value, &provider).unwrap(); + trie.update_leaf(key_keep, value.clone()).unwrap(); + trie.update_leaf(key_drop_1, value.clone()).unwrap(); + trie.update_leaf(key_drop_2, value).unwrap(); let root_before = trie.root(); @@ -7856,7 +7674,6 @@ mod tests { #[test] fn test_prune_update_after() { // After pruning, we should be able to update leaves without panic. - let provider = DefaultTrieNodeProvider; let mut trie = ParallelSparseTrie::default(); let value = large_account_value(); @@ -7869,7 +7686,6 @@ mod tests { first, second, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, ])), value.clone(), - &provider, ) .unwrap(); } @@ -7886,7 +7702,7 @@ mod tests { // were replaced with Blind(None) let new_path = pad_nibbles_right(Nibbles::from_nibbles([0x5, 0x5, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6])); - trie.update_leaf(new_path, value, &provider).unwrap(); + trie.update_leaf(new_path, value).unwrap(); // The trie should still be functional let _ = trie.root(); @@ -7899,15 +7715,13 @@ mod tests { use crate::LeafUpdate; use alloy_primitives::map::B256Map; use std::cell::RefCell; - - let provider = DefaultTrieNodeProvider; let mut trie = ParallelSparseTrie::default(); // Create a leaf in the trie using a full-length key let b256_key = B256::with_last_byte(42); let key = Nibbles::unpack(b256_key); let value = encode_account_value(1); - trie.update_leaf(key, value, &provider).unwrap(); + trie.update_leaf(key, value).unwrap(); // Create update map with a new value for the same key let new_value = encode_account_value(2); @@ -8049,8 +7863,6 @@ mod tests { use crate::LeafUpdate; use alloy_primitives::map::B256Map; use std::cell::RefCell; - - let provider = DefaultTrieNodeProvider; let mut trie = ParallelSparseTrie::default(); // Create two leaves so removal doesn't result in empty trie issues @@ -8060,8 +7872,8 @@ mod tests { let key1 = Nibbles::unpack(b256_key1); let key2 = Nibbles::unpack(b256_key2); let value = encode_account_value(1); - trie.update_leaf(key1, value.clone(), &provider).unwrap(); - trie.update_leaf(key2, value, &provider).unwrap(); + trie.update_leaf(key1, value.clone()).unwrap(); + trie.update_leaf(key2, value).unwrap(); // Create an update to remove key1 (empty value = removal) let mut updates: B256Map = B256Map::default(); @@ -8275,15 +8087,13 @@ mod tests { use crate::LeafUpdate; use alloy_primitives::map::B256Map; use std::cell::RefCell; - - let provider = DefaultTrieNodeProvider; let mut trie = ParallelSparseTrie::default(); // Create a leaf in the trie using a full-length key let b256_key = B256::with_last_byte(42); let key = Nibbles::unpack(b256_key); let value = encode_account_value(1); - trie.update_leaf(key, value, &provider).unwrap(); + trie.update_leaf(key, value).unwrap(); // Create a Touched update for the existing key let mut updates: B256Map = B256Map::default(); @@ -8695,8 +8505,8 @@ mod tests { // Create an empty trie and update with two leaves let mut trie = ParallelSparseTrie::default(); - trie.update_leaf(leaf1_path, vec![0x1], DefaultTrieNodeProvider).unwrap(); - trie.update_leaf(leaf2_path, vec![0x2], DefaultTrieNodeProvider).unwrap(); + trie.update_leaf(leaf1_path, vec![0x1]).unwrap(); + trie.update_leaf(leaf2_path, vec![0x2]).unwrap(); // Call root() to compute the trie root hash let _root = trie.root(); diff --git a/crates/trie/sparse/src/provider.rs b/crates/trie/sparse/src/provider.rs deleted file mode 100644 index bfd44424cc3..00000000000 --- a/crates/trie/sparse/src/provider.rs +++ /dev/null @@ -1,88 +0,0 @@ -//! Traits and default implementations related to retrieval of blinded trie nodes. - -use alloy_primitives::{Bytes, B256}; -use reth_execution_errors::SparseTrieError; -use reth_trie_common::{Nibbles, TrieMask}; - -/// Factory for instantiating trie node providers. -#[auto_impl::auto_impl(&)] -pub trait TrieNodeProviderFactory { - /// Type capable of fetching blinded account nodes. - type AccountNodeProvider: TrieNodeProvider; - /// Type capable of fetching blinded storage nodes. - type StorageNodeProvider: TrieNodeProvider; - - /// Returns blinded account node provider. - fn account_node_provider(&self) -> Self::AccountNodeProvider; - - /// Returns blinded storage node provider. - fn storage_node_provider(&self, account: B256) -> Self::StorageNodeProvider; -} - -/// Revealed blinded trie node. -#[derive(PartialEq, Eq, Clone, Debug)] -pub struct RevealedNode { - /// Raw trie node. - pub node: Bytes, - /// Branch node tree mask, if any. - pub tree_mask: Option, - /// Branch node hash mask, if any. - pub hash_mask: Option, -} - -/// Trie node provider for retrieving trie nodes. -#[auto_impl::auto_impl(&)] -pub trait TrieNodeProvider { - /// Retrieve trie node by path. - fn trie_node(&self, path: &Nibbles) -> Result, SparseTrieError>; -} - -/// Default trie node provider factory that creates [`DefaultTrieNodeProviderFactory`]. -#[derive(PartialEq, Eq, Clone, Default, Debug)] -pub struct DefaultTrieNodeProviderFactory; - -impl TrieNodeProviderFactory for DefaultTrieNodeProviderFactory { - type AccountNodeProvider = DefaultTrieNodeProvider; - type StorageNodeProvider = DefaultTrieNodeProvider; - - fn account_node_provider(&self) -> Self::AccountNodeProvider { - DefaultTrieNodeProvider - } - - fn storage_node_provider(&self, _account: B256) -> Self::StorageNodeProvider { - DefaultTrieNodeProvider - } -} - -/// Default trie node provider that always returns `Ok(None)`. -#[derive(PartialEq, Eq, Clone, Default, Debug)] -pub struct DefaultTrieNodeProvider; - -impl TrieNodeProvider for DefaultTrieNodeProvider { - fn trie_node(&self, _path: &Nibbles) -> Result, SparseTrieError> { - Ok(None) - } -} - -/// A provider that never reveals nodes from the database. -/// -/// This is used by `update_leaves` to attempt trie operations without -/// performing any database lookups. When the trie encounters a blinded node -/// that would normally trigger a reveal, this provider returns `None`, -/// causing the operation to fail with a `BlindedNode` error. -#[derive(PartialEq, Eq, Clone, Copy, Default, Debug)] -pub struct NoRevealProvider; - -impl TrieNodeProvider for NoRevealProvider { - fn trie_node(&self, _path: &Nibbles) -> Result, SparseTrieError> { - Ok(None) - } -} - -/// Right pad the path with 0s and return as [`B256`]. -#[inline] -pub fn pad_path_to_key(path: &Nibbles) -> B256 { - let mut padded = path.pack(); - padded.resize(32, 0); - B256::from_slice(&padded) -} diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index 8b35d8ab7be..7cccfbbb2e0 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -1,10 +1,8 @@ #[cfg(feature = "trie-debug")] use crate::debug_recorder::TrieDebugRecorder; use crate::{ - lfu::BucketedLfu, - provider::{TrieNodeProvider, TrieNodeProviderFactory}, - traits::SparseTrie as SparseTrieTrait, - ParallelSparseTrie, RevealableSparseTrie, + lfu::BucketedLfu, traits::SparseTrie as SparseTrieTrait, ParallelSparseTrie, + RevealableSparseTrie, }; use alloc::vec::Vec; use alloy_primitives::{map::B256Map, B256}; @@ -13,8 +11,8 @@ use reth_execution_errors::{SparseStateTrieResult, SparseTrieErrorKind}; use reth_primitives_traits::Account; use reth_trie_common::{ updates::{StorageTrieUpdates, TrieUpdates}, - BranchNodeMasks, DecodedMultiProof, MultiProof, Nibbles, ProofTrieNodeV2, TrieAccount, - TrieNodeV2, EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE, + DecodedMultiProof, MultiProof, Nibbles, ProofTrieNodeV2, TrieAccount, TrieNodeV2, + EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE, }; #[cfg(feature = "std")] use tracing::debug; @@ -480,56 +478,30 @@ where } /// Returns mutable reference to the revealed account sparse trie. - /// - /// If the trie is not revealed yet, its root will be revealed using the trie node provider. - fn revealed_trie_mut( - &mut self, - provider_factory: impl TrieNodeProviderFactory, - ) -> SparseStateTrieResult<&mut A> { - match self.state { - RevealableSparseTrie::Blind(_) => { - let (root_node, hash_mask, tree_mask) = provider_factory - .account_node_provider() - .trie_node(&Nibbles::default())? - .map(|node| { - TrieNodeV2::decode(&mut &node.node[..]) - .map(|decoded| (decoded, node.hash_mask, node.tree_mask)) - }) - .transpose()? - .unwrap_or((TrieNodeV2::EmptyRoot, None, None)); - let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); - self.state.reveal_root(root_node, masks, self.retain_updates).map_err(Into::into) - } - RevealableSparseTrie::Revealed(ref mut trie) => Ok(trie), - } + fn revealed_trie_mut(&mut self) -> SparseStateTrieResult<&mut A> { + self.state.as_revealed_mut().ok_or_else(|| SparseTrieErrorKind::Blind.into()) } /// Returns sparse trie root. - /// - /// If the trie has not been revealed, this function reveals the root node and returns its hash. - pub fn root( - &mut self, - provider_factory: impl TrieNodeProviderFactory, - ) -> SparseStateTrieResult { + pub fn root(&mut self) -> SparseStateTrieResult { // record revealed node metrics #[cfg(feature = "metrics")] self.metrics.record(); - Ok(self.revealed_trie_mut(provider_factory)?.root()) + Ok(self.revealed_trie_mut()?.root()) } - /// Returns sparse trie root and trie updates if the trie has been revealed. + /// Returns sparse trie root and trie updates. + /// + /// Returns an error if the account trie is still blind. #[instrument(level = "debug", target = "trie::sparse", skip_all)] - pub fn root_with_updates( - &mut self, - provider_factory: impl TrieNodeProviderFactory, - ) -> SparseStateTrieResult<(B256, TrieUpdates)> { + pub fn root_with_updates(&mut self) -> SparseStateTrieResult<(B256, TrieUpdates)> { // record revealed node metrics #[cfg(feature = "metrics")] self.metrics.record(); let storage_tries = self.storage_trie_updates(); - let revealed = self.revealed_trie_mut(provider_factory)?; + let revealed = self.revealed_trie_mut()?; let (root, updates) = (revealed.root(), revealed.take_updates()); let updates = TrieUpdates { @@ -575,16 +547,16 @@ where } }) } +} +impl SparseStateTrie { /// Update the account leaf node. pub fn update_account_leaf( &mut self, path: Nibbles, value: Vec, - provider_factory: impl TrieNodeProviderFactory, ) -> SparseStateTrieResult<()> { - let provider = provider_factory.account_node_provider(); - self.state.update_leaf(path, value, provider)?; + self.state.update_leaf(path, value)?; Ok(()) } @@ -594,14 +566,12 @@ where address: B256, slot: Nibbles, value: Vec, - provider_factory: impl TrieNodeProviderFactory, ) -> SparseStateTrieResult<()> { - let provider = provider_factory.storage_node_provider(address); self.storage .tries .get_mut(&address) .ok_or(SparseTrieErrorKind::Blind)? - .update_leaf(slot, value, provider)?; + .update_leaf(slot, value)?; Ok(()) } @@ -615,19 +585,18 @@ where &mut self, address: B256, account: Account, - provider_factory: impl TrieNodeProviderFactory, ) -> SparseStateTrieResult { let storage_root = if let Some(storage_trie) = self.storage.tries.get_mut(&address) { trace!(target: "trie::sparse", ?address, "Calculating storage root to update account"); storage_trie.root().ok_or(SparseTrieErrorKind::Blind)? } else if self.is_account_revealed(address) { trace!(target: "trie::sparse", ?address, "Retrieving storage root from account leaf to update account"); - // The account was revealed, either... + // The account path was revealed already, so either... if let Some(value) = self.get_account_value(&address) { - // ..it exists and we should take its current storage root or... + // ...the account leaf exists and we should reuse its current storage root, or... TrieAccount::decode(&mut &value[..])?.storage_root } else { - // ...the account is newly created and the storage trie is empty. + // ...the account is newly created and its storage trie is still empty. EMPTY_ROOT_HASH } } else { @@ -642,7 +611,7 @@ where let nibbles = Nibbles::unpack(address); self.account_rlp_buf.clear(); account.into_trie_account(storage_root).encode(&mut self.account_rlp_buf); - self.update_account_leaf(nibbles, self.account_rlp_buf.clone(), provider_factory)?; + self.update_account_leaf(nibbles, self.account_rlp_buf.clone())?; Ok(true) } @@ -659,13 +628,12 @@ where &mut self, address: B256, account: Option, - provider_factory: impl TrieNodeProviderFactory, ) -> SparseStateTrieResult<()> { let nibbles = Nibbles::unpack(address); let Some(account) = account else { trace!(target: "trie::sparse", ?address, "Removing account"); - return self.remove_account_leaf(&nibbles, provider_factory); + return self.remove_account_leaf(&nibbles); }; let storage_root = if let Some(storage_trie) = self.storage.tries.get_mut(&address) { @@ -681,7 +649,7 @@ where trace!(target: "trie::sparse", ?address, "Updating account"); self.account_rlp_buf.clear(); account.into_trie_account(storage_root).encode(&mut self.account_rlp_buf); - self.update_account_leaf(nibbles, self.account_rlp_buf.clone(), provider_factory)?; + self.update_account_leaf(nibbles, self.account_rlp_buf.clone())?; Ok(()) } @@ -693,11 +661,7 @@ where /// Returns false if the new storage root is empty, and the account info was already empty, /// indicating the account leaf should be removed. #[instrument(level = "debug", target = "trie::sparse", skip_all)] - pub fn update_account_storage_root( - &mut self, - address: B256, - provider_factory: impl TrieNodeProviderFactory, - ) -> SparseStateTrieResult { + pub fn update_account_storage_root(&mut self, address: B256) -> SparseStateTrieResult { if !self.is_account_revealed(address) { return Err(SparseTrieErrorKind::Blind.into()) } @@ -712,8 +676,7 @@ where return Ok(true) }; - // Calculate the new storage root. If the storage trie doesn't exist, the storage root will - // be empty. + // If the storage trie doesn't exist, the new storage root is empty. let storage_root = if let Some(storage_trie) = self.storage.tries.get_mut(&address) { trace!(target: "trie::sparse", ?address, "Calculating storage root to update account"); storage_trie.root().ok_or(SparseTrieErrorKind::Blind)? @@ -724,30 +687,25 @@ where // Update the account with the new storage root. trie_account.storage_root = storage_root; - // If the account is empty, indicate that it should be removed. + // If the account is now empty, indicate that its leaf should be removed. if trie_account == TrieAccount::default() { return Ok(false) } - // Otherwise, update the account leaf. + // Otherwise, rewrite the account leaf with the updated storage root. trace!(target: "trie::sparse", ?address, "Updating account with the new storage root"); let nibbles = Nibbles::unpack(address); self.account_rlp_buf.clear(); trie_account.encode(&mut self.account_rlp_buf); - self.update_account_leaf(nibbles, self.account_rlp_buf.clone(), provider_factory)?; + self.update_account_leaf(nibbles, self.account_rlp_buf.clone())?; Ok(true) } /// Remove the account leaf node. #[instrument(level = "debug", target = "trie::sparse", skip_all)] - pub fn remove_account_leaf( - &mut self, - path: &Nibbles, - provider_factory: impl TrieNodeProviderFactory, - ) -> SparseStateTrieResult<()> { - let provider = provider_factory.account_node_provider(); - self.state.remove_leaf(path, provider)?; + pub fn remove_account_leaf(&mut self, path: &Nibbles) -> SparseStateTrieResult<()> { + self.state.remove_leaf(path)?; Ok(()) } @@ -756,13 +714,11 @@ where &mut self, address: B256, slot: &Nibbles, - provider_factory: impl TrieNodeProviderFactory, ) -> SparseStateTrieResult<()> { let storage_trie = self.storage.tries.get_mut(&address).ok_or(SparseTrieErrorKind::Blind)?; - let provider = provider_factory.storage_node_provider(address); - storage_trie.remove_leaf(slot, provider)?; + storage_trie.remove_leaf(slot)?; Ok(()) } } @@ -1073,7 +1029,7 @@ fn estimate_v2_proof_capacity(nodes: &[ProofTrieNodeV2]) -> usize { #[cfg(test)] mod tests { use super::*; - use crate::{provider::DefaultTrieNodeProviderFactory, LeafLookup, ParallelSparseTrie}; + use crate::{LeafLookup, ParallelSparseTrie}; use alloy_primitives::{ b256, map::{HashMap, HashSet}, @@ -1081,6 +1037,7 @@ mod tests { }; use arbitrary::Arbitrary; use rand::{rngs::StdRng, Rng, SeedableRng}; + use reth_execution_errors::{SparseStateTrieErrorKind, SparseTrieErrorKind}; use reth_primitives_traits::Account; use reth_trie::{updates::StorageTrieUpdates, HashBuilder, MultiProof, EMPTY_ROOT_HASH}; use reth_trie_common::{ @@ -1099,7 +1056,6 @@ mod tests { #[test] fn reveal_account_path_twice() { - let provider_factory = DefaultTrieNodeProviderFactory; let mut sparse = SparseStateTrie::::default(); // Full 64-nibble paths @@ -1148,7 +1104,7 @@ mod tests { // Remove the leaf node and check that the state trie does not contain the leaf node and // value - sparse.remove_account_leaf(&full_path_0, &provider_factory).unwrap(); + sparse.remove_account_leaf(&full_path_0).unwrap(); assert!(matches!( sparse.state_trie_ref().unwrap().find_leaf(&full_path_0, None), Ok(LeafLookup::NonExistent) @@ -1158,7 +1114,6 @@ mod tests { #[test] fn reveal_storage_path_twice() { - let provider_factory = DefaultTrieNodeProviderFactory; let mut sparse = SparseStateTrie::::default(); // Full 64-nibble path @@ -1212,7 +1167,7 @@ mod tests { // Remove the leaf node and check that the storage trie does not contain the leaf node and // value - sparse.remove_storage_leaf(B256::ZERO, &full_path_0, &provider_factory).unwrap(); + sparse.remove_storage_leaf(B256::ZERO, &full_path_0).unwrap(); assert!(matches!( sparse.storage_trie_ref(&B256::ZERO).unwrap().find_leaf(&full_path_0, None), Ok(LeafLookup::NonExistent) @@ -1244,7 +1199,6 @@ mod tests { #[test] fn reveal_v2_proof_nodes() { - let provider_factory = DefaultTrieNodeProviderFactory; let mut sparse = SparseStateTrie::::default(); // Full 64-nibble path @@ -1292,13 +1246,12 @@ mod tests { ); // Remove the leaf node - sparse.remove_account_leaf(&full_path_0, &provider_factory).unwrap(); + sparse.remove_account_leaf(&full_path_0).unwrap(); assert!(sparse.state_trie_ref().unwrap().get_leaf_value(&full_path_0).is_none()); } #[test] fn reveal_storage_v2_proof_nodes() { - let provider_factory = DefaultTrieNodeProviderFactory; let mut sparse = SparseStateTrie::::default(); // Full 64-nibble path @@ -1338,7 +1291,7 @@ mod tests { ); // Remove the leaf node - sparse.remove_storage_leaf(B256::ZERO, &full_path_0, &provider_factory).unwrap(); + sparse.remove_storage_leaf(B256::ZERO, &full_path_0).unwrap(); assert!(sparse .storage_trie_ref(&B256::ZERO) .unwrap() @@ -1346,6 +1299,15 @@ mod tests { .is_none()); } + #[test] + fn root_on_blind_trie_returns_blind_error() { + let mut sparse = SparseStateTrie::::default(); + + let err = sparse.root().unwrap_err(); + + assert!(matches!(err.kind(), SparseStateTrieErrorKind::Sparse(SparseTrieErrorKind::Blind))); + } + #[test] fn take_trie_updates() { reth_tracing::init_test_tracing(); @@ -1400,8 +1362,6 @@ mod tests { let root = hash_builder.root(); let proof_nodes = hash_builder.take_proof_nodes(); - - let provider_factory = DefaultTrieNodeProviderFactory; let mut sparse = SparseStateTrie::::default().with_updates(true); sparse .reveal_decoded_multiproof( @@ -1438,49 +1398,24 @@ mod tests { ) .unwrap(); - assert_eq!(sparse.root(&provider_factory).unwrap(), root); + assert_eq!(sparse.root().unwrap(), root); let address_3 = b256!("0x2000000000000000000000000000000000000000000000000000000000000000"); let address_path_3 = Nibbles::unpack(address_3); let account_3 = Account { nonce: account_1.nonce + 1, ..account_1 }; let trie_account_3 = account_3.into_trie_account(EMPTY_ROOT_HASH); - sparse - .update_account_leaf( - address_path_3, - alloy_rlp::encode(trie_account_3), - &provider_factory, - ) - .unwrap(); + sparse.update_account_leaf(address_path_3, alloy_rlp::encode(trie_account_3)).unwrap(); - sparse - .update_storage_leaf( - address_1, - slot_path_3, - alloy_rlp::encode(value_3), - &provider_factory, - ) - .unwrap(); + sparse.update_storage_leaf(address_1, slot_path_3, alloy_rlp::encode(value_3)).unwrap(); trie_account_1.storage_root = sparse.storage_root(&address_1).unwrap(); - sparse - .update_account_leaf( - address_path_1, - alloy_rlp::encode(trie_account_1), - &provider_factory, - ) - .unwrap(); + sparse.update_account_leaf(address_path_1, alloy_rlp::encode(trie_account_1)).unwrap(); sparse.wipe_storage(address_2).unwrap(); trie_account_2.storage_root = sparse.storage_root(&address_2).unwrap(); - sparse - .update_account_leaf( - address_path_2, - alloy_rlp::encode(trie_account_2), - &provider_factory, - ) - .unwrap(); + sparse.update_account_leaf(address_path_2, alloy_rlp::encode(trie_account_2)).unwrap(); - sparse.root(&provider_factory).unwrap(); + sparse.root().unwrap(); let sparse_updates = sparse.take_trie_updates().unwrap(); // TODO(alexey): assert against real state root calculation updates diff --git a/crates/trie/sparse/src/traits.rs b/crates/trie/sparse/src/traits.rs index 5daf5811da2..06c8aa77acf 100644 --- a/crates/trie/sparse/src/traits.rs +++ b/crates/trie/sparse/src/traits.rs @@ -13,7 +13,6 @@ use reth_trie_common::{BranchNodeMasks, Nibbles, ProofTrieNodeV2, TrieNodeV2}; #[cfg(feature = "trie-debug")] use crate::debug_recorder::TrieDebugRecorder; -use crate::provider::TrieNodeProvider; /// Describes an update to a leaf in the sparse trie. #[derive(Debug, Clone, PartialEq, Eq)] @@ -140,46 +139,6 @@ pub trait SparseTrie: Sized + Debug + Send + Sync { /// each node with [`TrieNodeV2::EmptyRoot`] to avoid cloning. fn reveal_nodes(&mut self, nodes: &mut [ProofTrieNodeV2]) -> SparseTrieResult<()>; - /// Updates the value of a leaf node at the specified path. - /// - /// If the leaf doesn't exist, it will be created. - /// If it does exist, its value will be updated. - /// - /// # Arguments - /// - /// * `full_path` - The full path to the leaf - /// * `value` - The new value for the leaf - /// * `provider` - The trie provider for resolving missing nodes - /// - /// # Returns - /// - /// `Ok(())` if successful, or an error if the update failed. - fn update_leaf( - &mut self, - full_path: Nibbles, - value: Vec, - provider: P, - ) -> SparseTrieResult<()>; - - /// Removes a leaf node at the specified path. - /// - /// This will also handle collapsing the trie structure as needed - /// (e.g., removing branch nodes that become unnecessary). - /// - /// # Arguments - /// - /// * `full_path` - The full path to the leaf to remove - /// * `provider` - The trie node provider for resolving missing nodes - /// - /// # Returns - /// - /// `Ok(())` if successful, or an error if the removal failed. - fn remove_leaf( - &mut self, - full_path: &Nibbles, - provider: P, - ) -> SparseTrieResult<()>; - /// Calculates and returns the root hash of the trie. /// /// This processes any dirty nodes by updating their RLP encodings @@ -459,23 +418,6 @@ mod configurable_sparse_trie { delegate!(self, reveal_nodes, nodes) } - fn update_leaf( - &mut self, - full_path: Nibbles, - value: Vec, - provider: P, - ) -> SparseTrieResult<()> { - delegate!(self, update_leaf, full_path, value, provider) - } - - fn remove_leaf( - &mut self, - full_path: &Nibbles, - provider: P, - ) -> SparseTrieResult<()> { - delegate!(self, remove_leaf, full_path, provider) - } - fn root(&mut self) -> B256 { delegate!(self, root) } diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 4fc1b5b349f..0fa1fd9c66e 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -1,7 +1,4 @@ -use crate::{ - provider::TrieNodeProvider, LeafUpdate, ParallelSparseTrie, SparseTrie as SparseTrieTrait, - SparseTrieUpdates, -}; +use crate::{LeafUpdate, ParallelSparseTrie, SparseTrie as SparseTrieTrait, SparseTrieUpdates}; use alloc::{borrow::Cow, boxed::Box, vec::Vec}; use alloy_primitives::{map::B256Map, B256}; use reth_execution_errors::{SparseTrieErrorKind, SparseTrieResult}; @@ -90,7 +87,7 @@ impl RevealableSparseTrie { /// # Examples /// /// ``` - /// use reth_trie_sparse::{provider::DefaultTrieNodeProvider, RevealableSparseTrie}; + /// use reth_trie_sparse::RevealableSparseTrie; /// /// let trie = ::blind(); /// assert!(trie.is_blind()); @@ -202,39 +199,6 @@ impl RevealableSparseTrie { }; } - /// Updates (or inserts) a leaf at the given key path with the specified RLP-encoded value. - /// - /// # Errors - /// - /// Returns an error if the trie is still blind, or if the update fails. - #[instrument(level = "trace", target = "trie::sparse", skip_all)] - pub fn update_leaf( - &mut self, - path: Nibbles, - value: Vec, - provider: impl TrieNodeProvider, - ) -> SparseTrieResult<()> { - let revealed = self.as_revealed_mut().ok_or(SparseTrieErrorKind::Blind)?; - revealed.update_leaf(path, value, provider)?; - Ok(()) - } - - /// Removes a leaf node at the specified key path. - /// - /// # Errors - /// - /// Returns an error if the trie is still blind, or if the leaf cannot be removed - #[instrument(level = "trace", target = "trie::sparse", skip_all)] - pub fn remove_leaf( - &mut self, - path: &Nibbles, - provider: impl TrieNodeProvider, - ) -> SparseTrieResult<()> { - let revealed = self.as_revealed_mut().ok_or(SparseTrieErrorKind::Blind)?; - revealed.remove_leaf(path, provider)?; - Ok(()) - } - /// Shrinks the capacity of the sparse trie's node storage. /// Works for both revealed and blind tries with allocated storage. pub fn shrink_nodes_to(&mut self, size: usize) { @@ -258,6 +222,32 @@ impl RevealableSparseTrie { } } +impl RevealableSparseTrie { + /// Updates (or inserts) a leaf at the given key path with the specified RLP-encoded value. + /// + /// # Errors + /// + /// Returns an error if the trie is still blind, or if the update fails. + #[instrument(level = "trace", target = "trie::sparse", skip_all)] + pub fn update_leaf(&mut self, path: Nibbles, value: Vec) -> SparseTrieResult<()> { + let revealed = self.as_revealed_mut().ok_or(SparseTrieErrorKind::Blind)?; + revealed.update_leaf(path, value)?; + Ok(()) + } + + /// Removes a leaf node at the specified key path. + /// + /// # Errors + /// + /// Returns an error if the trie is still blind, or if the leaf cannot be removed. + #[instrument(level = "trace", target = "trie::sparse", skip_all)] + pub fn remove_leaf(&mut self, path: &Nibbles) -> SparseTrieResult<()> { + let revealed = self.as_revealed_mut().ok_or(SparseTrieErrorKind::Blind)?; + revealed.remove_leaf(path)?; + Ok(()) + } +} + impl RevealableSparseTrie { /// Applies batch leaf updates to the sparse trie. /// diff --git a/crates/trie/trie/src/proof/mod.rs b/crates/trie/trie/src/proof/mod.rs index 58b9740e103..0c787362f43 100644 --- a/crates/trie/trie/src/proof/mod.rs +++ b/crates/trie/trie/src/proof/mod.rs @@ -23,9 +23,6 @@ use reth_trie_common::{ MultiProof, MultiProofTargets, MultiProofTargetsV2, StorageMultiProof, }; -mod trie_node; -pub use trie_node::*; - /// A struct for generating merkle proofs. /// /// Proof generator adds the target address and slots to the prefix set, enables the proof retainer diff --git a/crates/trie/trie/src/proof/trie_node.rs b/crates/trie/trie/src/proof/trie_node.rs deleted file mode 100644 index 1a7a77f792c..00000000000 --- a/crates/trie/trie/src/proof/trie_node.rs +++ /dev/null @@ -1,151 +0,0 @@ -use super::{Proof, StorageProof}; -use crate::{hashed_cursor::HashedCursorFactory, trie_cursor::TrieCursorFactory}; -use alloy_primitives::{map::HashSet, B256}; -use reth_execution_errors::{SparseTrieError, SparseTrieErrorKind}; -use reth_primitives_traits::FastInstant as Instant; -use reth_trie_common::{MultiProofTargets, Nibbles}; -use reth_trie_sparse::provider::{ - pad_path_to_key, RevealedNode, TrieNodeProvider, TrieNodeProviderFactory, -}; -use tracing::{enabled, trace, Level}; - -/// Factory for instantiating providers capable of retrieving blinded trie nodes via proofs. -#[derive(Debug, Clone)] -pub struct ProofTrieNodeProviderFactory { - /// The cursor factory for traversing trie nodes. - trie_cursor_factory: T, - /// The factory for hashed cursors. - hashed_cursor_factory: H, -} - -impl ProofTrieNodeProviderFactory { - /// Create new proof-based blinded provider factory. - pub const fn new(trie_cursor_factory: T, hashed_cursor_factory: H) -> Self { - Self { trie_cursor_factory, hashed_cursor_factory } - } -} - -impl TrieNodeProviderFactory for ProofTrieNodeProviderFactory -where - T: TrieCursorFactory + Clone, - H: HashedCursorFactory + Clone, -{ - type AccountNodeProvider = ProofBlindedAccountProvider; - type StorageNodeProvider = ProofBlindedStorageProvider; - - fn account_node_provider(&self) -> Self::AccountNodeProvider { - ProofBlindedAccountProvider { - trie_cursor_factory: self.trie_cursor_factory.clone(), - hashed_cursor_factory: self.hashed_cursor_factory.clone(), - } - } - - fn storage_node_provider(&self, account: B256) -> Self::StorageNodeProvider { - ProofBlindedStorageProvider { - trie_cursor_factory: self.trie_cursor_factory.clone(), - hashed_cursor_factory: self.hashed_cursor_factory.clone(), - account, - } - } -} - -/// Blinded provider for retrieving account trie nodes by path. -#[derive(Debug)] -pub struct ProofBlindedAccountProvider { - /// The cursor factory for traversing trie nodes. - trie_cursor_factory: T, - /// The factory for hashed cursors. - hashed_cursor_factory: H, -} - -impl ProofBlindedAccountProvider { - /// Create new proof-based blinded account node provider. - pub const fn new(trie_cursor_factory: T, hashed_cursor_factory: H) -> Self { - Self { trie_cursor_factory, hashed_cursor_factory } - } -} - -impl TrieNodeProvider for ProofBlindedAccountProvider -where - T: TrieCursorFactory, - H: HashedCursorFactory, -{ - fn trie_node(&self, path: &Nibbles) -> Result, SparseTrieError> { - let start = enabled!(target: "trie::proof::blinded", Level::TRACE).then(Instant::now); - - let targets = MultiProofTargets::from_iter([(pad_path_to_key(path), HashSet::default())]); - let mut proof = Proof::new(&self.trie_cursor_factory, &self.hashed_cursor_factory) - .with_branch_node_masks(true) - .multiproof(targets) - .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?; - let node = proof.account_subtree.into_inner().remove(path); - let masks = proof.branch_node_masks.remove(path); - let hash_mask = masks.map(|m| m.hash_mask); - let tree_mask = masks.map(|m| m.tree_mask); - - trace!( - target: "trie::proof::blinded", - elapsed = ?start.unwrap().elapsed(), - ?path, - ?node, - ?tree_mask, - ?hash_mask, - "Blinded node for account trie" - ); - Ok(node.map(|node| RevealedNode { node, tree_mask, hash_mask })) - } -} - -/// Blinded provider for retrieving storage trie nodes by path. -#[derive(Debug)] -pub struct ProofBlindedStorageProvider { - /// The cursor factory for traversing trie nodes. - trie_cursor_factory: T, - /// The factory for hashed cursors. - hashed_cursor_factory: H, - /// Target account. - account: B256, -} - -impl ProofBlindedStorageProvider { - /// Create new proof-based blinded storage node provider. - pub const fn new(trie_cursor_factory: T, hashed_cursor_factory: H, account: B256) -> Self { - Self { trie_cursor_factory, hashed_cursor_factory, account } - } -} - -impl TrieNodeProvider for ProofBlindedStorageProvider -where - T: TrieCursorFactory, - H: HashedCursorFactory, -{ - fn trie_node(&self, path: &Nibbles) -> Result, SparseTrieError> { - let start = enabled!(target: "trie::proof::blinded", Level::TRACE).then(Instant::now); - - let targets = HashSet::from_iter([pad_path_to_key(path)]); - let mut proof = StorageProof::new_hashed( - &self.trie_cursor_factory, - &self.hashed_cursor_factory, - self.account, - ) - .with_branch_node_masks(true) - .storage_multiproof(targets) - .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?; - let node = proof.subtree.into_inner().remove(path); - let masks = proof.branch_node_masks.remove(path); - let hash_mask = masks.map(|m| m.hash_mask); - let tree_mask = masks.map(|m| m.tree_mask); - - trace!( - target: "trie::proof::blinded", - account = ?self.account, - elapsed = ?start.unwrap().elapsed(), - ?path, - ?node, - ?tree_mask, - ?hash_mask, - "Blinded node for storage trie" - ); - Ok(node.map(|node| RevealedNode { node, tree_mask, hash_mask })) - } -}