Compare commits

...

14 Commits

Author SHA1 Message Date
Dan Cline
6f9517636f chore: integrate cli and update_leaves 2026-01-28 22:02:58 -05:00
Dan Cline
3b26415fbe chore: integrate update_leaves 2026-01-28 22:02:58 -05:00
Dan Cline
f3d4fcc1d0 feat(cli): add sparse-trie-as-cache cli flag 2026-01-28 21:20:58 -05:00
Arsenii Kulikov
1ca2013cac wip 2026-01-29 01:55:59 +04:00
Arsenii Kulikov
6d6cbd61c4 wip 2026-01-28 23:32:22 +04:00
Arsenii Kulikov
d9105a29f3 wip 2026-01-28 23:31:43 +04:00
Arsenii Kulikov
ee1aee05e4 wip 2026-01-28 23:31:42 +04:00
Arsenii Kulikov
b3eff8f2fa wip 2026-01-28 23:31:42 +04:00
Arsenii Kulikov
338e6408a8 feat: STaC task scaffolding 2026-01-28 23:31:38 +04:00
yongkangc
1c5727eb46 fix: fmt and docs warning
Amp-Thread-ID: https://ampcode.com/threads/T-019c0549-b6f7-72f8-810a-d2578c54b33d
2026-01-28 16:09:10 +00:00
yongkangc
a1922c0c1c docs: add comments to update_leaves implementation
Amp-Thread-ID: https://ampcode.com/threads/T-019c0549-b6f7-72f8-810a-d2578c54b33d
2026-01-28 15:50:56 +00:00
yongkangc
1d6fec4741 feat(trie): parallelize Touched operations in update_leaves
- Refactor update_leaves to sort and group updates by subtrie
- Split operations into upper (sequential) and lower subtrie ops
- Process Touched operations in parallel using rayon when threshold met
- Add helper functions: apply_single_update, process_touched_ops_parallel
- Add UpdateLeafOutcome enum for clean error handling

Amp-Thread-ID: https://ampcode.com/threads/T-019c0541-585e-744b-aed8-c978033375bc
2026-01-28 15:45:52 +00:00
yongkangc
3bbf31ccda fix(trie): fix update_leaves rollback and add tests
- Fix rollback logic to use upper_subtrie.inner.values where update_leaf inserts
- Extract repeated min_len calculation to request_proof closure
- Add MAX_NIBBLE_PATH_LEN constant instead of magic number 64
- Use remove-then-reinsert pattern to avoid unnecessary clones

Add 5 new tests:
- test_update_leaves_multiple_keys
- test_update_leaves_batch_removal
- test_update_leaves_rollback_on_blinded_node
- test_update_leaves_retry_after_reveal
- test_update_leaves_batch_update_and_insert

Amp-Thread-ID: https://ampcode.com/threads/T-019c04fa-b228-77f9-8370-1d5242106798
2026-01-28 14:47:20 +00:00
yongkangc
d1ad0d1025 feat(trie): add update_leaves method to SparseTrieExt
Implements batch leaf updates for sparse tries with blinded node handling.

- Add LeafUpdate enum (Changed/Touched variants) to traits.rs
- Add update_leaves method to SparseTrieExt trait
- Implement for SerialSparseTrie with snapshot-revert on BlindedNode error
- Implement for ParallelSparseTrie with same semantics
- Add NoRevealProvider helper for short-circuiting pattern
- Add 14 tests covering all scenarios

The method accepts a B256Map of updates and a callback for proof targets.
Successfully applied updates are removed from the map. Updates blocked by
blinded nodes remain in the map and trigger the callback with (path, min_len)
to request the necessary proofs.

Closes RETH-177

Amp-Thread-ID: https://ampcode.com/threads/T-019c0440-00a2-7339-87c5-c55cffe062f0
2026-01-28 14:16:06 +00:00
10 changed files with 1487 additions and 56 deletions

View File

@@ -152,6 +152,8 @@ pub struct TreeConfig {
disable_proof_v2: bool,
/// Whether to disable cache metrics recording (can be expensive with large cached state).
disable_cache_metrics: bool,
/// Whether to enable sparse trie as cache.
enable_sparse_trie_as_cache: bool,
}
impl Default for TreeConfig {
@@ -181,6 +183,7 @@ impl Default for TreeConfig {
account_worker_count: default_account_worker_count(),
disable_proof_v2: false,
disable_cache_metrics: false,
enable_sparse_trie_as_cache: false,
}
}
}
@@ -213,6 +216,7 @@ impl TreeConfig {
account_worker_count: usize,
disable_proof_v2: bool,
disable_cache_metrics: bool,
enable_sparse_trie_as_cache: bool,
) -> Self {
Self {
persistence_threshold,
@@ -239,6 +243,7 @@ impl TreeConfig {
account_worker_count,
disable_proof_v2,
disable_cache_metrics,
enable_sparse_trie_as_cache,
}
}
@@ -540,4 +545,15 @@ impl TreeConfig {
self.disable_cache_metrics = disable_cache_metrics;
self
}
/// Returns whether sparse trie as cache is enabled.
pub const fn enable_sparse_trie_as_cache(&self) -> bool {
self.enable_sparse_trie_as_cache
}
/// Setter for whether to enable sparse trie as cache.
pub const fn with_sparse_trie_as_cache(mut self, enabled: bool) -> Self {
self.enable_sparse_trie_as_cache = enabled;
self
}
}

View File

@@ -7,14 +7,14 @@ use crate::tree::{
prewarm::{PrewarmCacheTask, PrewarmContext, PrewarmMode, PrewarmTaskEvent},
sparse_trie::StateRootComputeOutcome,
},
sparse_trie::SparseTrieTask,
sparse_trie::{SparseTrieCacheTask, SparseTrieTask},
StateProviderBuilder, TreeConfig,
};
use alloy_eip7928::BlockAccessList;
use alloy_eips::eip1898::BlockWithParent;
use alloy_evm::block::StateChangeSource;
use alloy_primitives::B256;
use crossbeam_channel::Sender as CrossbeamSender;
use crossbeam_channel::{Receiver as CrossbeamReceiver, Sender as CrossbeamSender};
use executor::WorkloadExecutor;
use metrics::Counter;
use multiproof::{SparseTrieUpdate, *};
@@ -39,10 +39,7 @@ use reth_trie_parallel::{
proof_task::{ProofTaskCtx, ProofWorkerHandle},
root::ParallelStateRootError,
};
use reth_trie_sparse::{
provider::{TrieNodeProvider, TrieNodeProviderFactory},
ClearedSparseStateTrie, RevealableSparseTrie, SparseStateTrie,
};
use reth_trie_sparse::{ClearedSparseStateTrie, RevealableSparseTrie, SparseStateTrie};
use reth_trie_sparse_parallel::{ParallelSparseTrie, ParallelismThresholds};
use std::{
collections::BTreeMap,
@@ -283,37 +280,45 @@ where
v2_proofs_enabled,
);
let multi_proof_task = MultiProofTask::new(
proof_handle.clone(),
to_sparse_trie,
config.multiproof_chunking_enabled().then_some(config.multiproof_chunk_size()),
to_multi_proof.clone(),
from_multi_proof,
)
.with_v2_proofs_enabled(v2_proofs_enabled);
if !config.enable_sparse_trie_as_cache() {
let multi_proof_task = MultiProofTask::new(
proof_handle.clone(),
to_sparse_trie,
config.multiproof_chunking_enabled().then_some(config.multiproof_chunk_size()),
to_multi_proof.clone(),
from_multi_proof.clone(),
)
.with_v2_proofs_enabled(v2_proofs_enabled);
// spawn multi-proof task
let parent_span = span.clone();
let saved_cache = prewarm_handle.saved_cache.clone();
self.executor.spawn_blocking(move || {
let _enter = parent_span.entered();
// Build a state provider for the multiproof task
let provider = provider_builder.build().expect("failed to build provider");
let provider = if let Some(saved_cache) = saved_cache {
let (cache, metrics, _disable_metrics) = saved_cache.split();
Box::new(CachedStateProvider::new(provider, cache, metrics))
as Box<dyn StateProvider>
} else {
Box::new(provider)
};
multi_proof_task.run(provider);
});
// spawn multi-proof task
let parent_span = span.clone();
let saved_cache = prewarm_handle.saved_cache.clone();
self.executor.spawn_blocking(move || {
let _enter = parent_span.entered();
// Build a state provider for the multiproof task
let provider = provider_builder.build().expect("failed to build provider");
let provider = if let Some(saved_cache) = saved_cache {
let (cache, metrics, _disable_metrics) = saved_cache.split();
Box::new(CachedStateProvider::new(provider, cache, metrics))
as Box<dyn StateProvider>
} else {
Box::new(provider)
};
multi_proof_task.run(provider);
});
}
// wire the sparse trie to the state root response receiver
let (state_root_tx, state_root_rx) = channel();
// Spawn the sparse trie task using any stored trie and parallel trie configuration.
self.spawn_sparse_trie_task(sparse_trie_rx, proof_handle, state_root_tx);
self.spawn_sparse_trie_task(
sparse_trie_rx,
proof_handle,
state_root_tx,
from_multi_proof,
config,
);
PayloadHandle {
to_multi_proof: Some(to_multi_proof),
@@ -493,19 +498,18 @@ where
/// Spawns the [`SparseTrieTask`] for this payload processor.
#[instrument(level = "debug", target = "engine::tree::payload_processor", skip_all)]
fn spawn_sparse_trie_task<BPF>(
fn spawn_sparse_trie_task(
&self,
sparse_trie_rx: mpsc::Receiver<SparseTrieUpdate>,
proof_worker_handle: BPF,
proof_worker_handle: ProofWorkerHandle,
state_root_tx: mpsc::Sender<Result<StateRootComputeOutcome, ParallelStateRootError>>,
) where
BPF: TrieNodeProviderFactory + Clone + Send + Sync + 'static,
BPF::AccountNodeProvider: TrieNodeProvider + Send + Sync,
BPF::StorageNodeProvider: TrieNodeProvider + Send + Sync,
{
from_multi_proof: CrossbeamReceiver<MultiProofMessage>,
config: &TreeConfig,
) {
let cleared_sparse_trie = Arc::clone(&self.sparse_state_trie);
let trie_metrics = self.trie_metrics.clone();
let span = Span::current();
let disable_sparse_trie_as_cache = !config.enable_sparse_trie_as_cache();
self.executor.spawn_blocking(move || {
let _enter = span.entered();
@@ -525,23 +529,36 @@ where
)
});
let task =
SparseTrieTask::<_, ParallelSparseTrie, ParallelSparseTrie>::new_with_cleared_trie(
let (result, trie) = if disable_sparse_trie_as_cache {
SparseTrieTask::new_with_cleared_trie(
sparse_trie_rx,
proof_worker_handle,
trie_metrics,
sparse_state_trie,
);
)
.run()
} else {
SparseTrieCacheTask::new_with_cleared_trie(
from_multi_proof,
proof_worker_handle,
trie_metrics,
sparse_state_trie,
)
.run()
};
let (result, trie) = task.run();
// Send state root computation result
let _ = state_root_tx.send(result);
// Clear the SparseStateTrie, shrink, and replace it back into the mutex _after_ sending
// results to the next step, so that time spent clearing doesn't block the step after
// this one.
// Clear the SparseStateTrie, shrink, and replace it back into the mutex _after_
// sending results to the next step, so that time spent clearing
// doesn't block the step after this one.
let _enter = debug_span!(target: "engine::tree::payload_processor", "clear").entered();
let mut cleared_trie = ClearedSparseStateTrie::from_state_trie(trie);
let mut cleared_trie = if disable_sparse_trie_as_cache {
ClearedSparseStateTrie::from_state_trie(trie)
} else {
ClearedSparseStateTrie::pruned(trie)
};
// Shrink the sparse trie so that we don't have ever increasing memory.
cleared_trie.shrink_to(

View File

@@ -1,15 +1,34 @@
//! Sparse Trie task related functionality.
use crate::tree::payload_processor::multiproof::{MultiProofTaskMetrics, SparseTrieUpdate};
use crate::tree::{
multiproof::{evm_state_to_hashed_post_state, MultiProofMessage, VersionedMultiProofTargets},
payload_processor::multiproof::{MultiProofTaskMetrics, SparseTrieUpdate},
};
use alloy_primitives::B256;
use alloy_rlp::Decodable;
use crossbeam_channel::{Receiver as CrossbeamReceiver, Sender as CrossbeamSender};
use rayon::iter::{ParallelBridge, ParallelIterator};
use reth_trie::{updates::TrieUpdates, Nibbles};
use reth_trie_parallel::{proof_task::ProofResult, root::ParallelStateRootError};
use reth_errors::ProviderError;
use reth_primitives_traits::Account;
use reth_revm::state::EvmState;
use reth_trie::{
proof_v2, updates::TrieUpdates, HashedPostState, Nibbles, TrieAccount, EMPTY_ROOT_HASH,
};
use reth_trie_parallel::{
proof_task::{
AccountMultiproofInput, ProofResult, ProofResultContext, ProofResultMessage,
ProofWorkerHandle,
},
root::ParallelStateRootError,
targets_v2::MultiProofTargetsV2,
};
use reth_trie_sparse::{
errors::{SparseStateTrieResult, SparseTrieErrorKind},
provider::{TrieNodeProvider, TrieNodeProviderFactory},
ClearedSparseStateTrie, SerialSparseTrie, SparseStateTrie, SparseTrie,
ClearedSparseStateTrie, LeafUpdate, SerialSparseTrie, SparseStateTrie, SparseTrie,
SparseTrieExt,
};
use revm_primitives::{hash_map::Entry, B256Map};
use smallvec::SmallVec;
use std::{
sync::mpsc,
@@ -129,6 +148,356 @@ where
}
}
pub(super) struct SparseTrieCacheTask<A = SerialSparseTrie, S = SerialSparseTrie> {
/// Sender for proof results.
proof_result_tx: CrossbeamSender<ProofResultMessage>,
/// Receiver for proof results directly from workers.
proof_result_rx: CrossbeamReceiver<ProofResultMessage>,
/// Receives updates from the state root task.
updates: CrossbeamReceiver<MultiProofMessage>,
/// `SparseStateTrie` used for computing the state root.
trie: SparseStateTrie<A, S>,
/// Handle to the proof worker pools (storage and account).
proof_worker_handle: ProofWorkerHandle,
/// Account trie updates.
account_updates: B256Map<LeafUpdate>,
/// Storage trie updates. hashed address -> slot -> update.
storage_updates: B256Map<B256Map<LeafUpdate>>,
/// Account updates that are blocked by storage root calculation or account reveal.
///
/// Those are being moved into `account_updates` once storage roots
/// are revealed and/or calculated.
///
/// Invariant: for each entry in `pending_account_updates` there is a corresponding
/// [`LeafUpdate::Touched`] entry in `account_updates`.
///
/// Values can be either of:
/// - None: account had a storage update and is awaiting storage root calculation and/or
/// account node reveal to complete.
/// - Some(_): account was changed/destroyed and is awaiting storage root calculation/reveal
/// to complete.
pending_account_updates: B256Map<Option<Option<Account>>>,
/// Metrics for the sparse trie.
metrics: MultiProofTaskMetrics,
}
impl<A, S> SparseTrieCacheTask<A, S>
where
A: SparseTrie + SparseTrieExt + Default,
S: SparseTrie + SparseTrieExt + Default + Clone,
{
/// Creates a new sparse trie, pre-populating with a [`ClearedSparseStateTrie`].
pub(super) fn new_with_cleared_trie(
updates: CrossbeamReceiver<MultiProofMessage>,
proof_worker_handle: ProofWorkerHandle,
metrics: MultiProofTaskMetrics,
sparse_state_trie: ClearedSparseStateTrie<A, S>,
) -> Self {
let (proof_result_tx, proof_result_rx) = crossbeam_channel::unbounded();
Self {
proof_result_tx,
proof_result_rx,
updates,
proof_worker_handle,
trie: sparse_state_trie.into_inner(),
account_updates: Default::default(),
storage_updates: Default::default(),
pending_account_updates: Default::default(),
metrics,
}
}
/// Runs the sparse trie task to completion.
///
/// This waits for new incoming [`SparseTrieUpdate`].
///
/// This concludes once the last trie update has been received.
///
/// # Returns
///
/// - State root computation outcome.
/// - `SparseStateTrie` that needs to be cleared and reused to avoid reallocations.
#[instrument(
level = "debug",
target = "engine::tree::payload_processor::sparse_trie",
skip_all
)]
pub(super) fn run(
mut self,
) -> (Result<StateRootComputeOutcome, ParallelStateRootError>, SparseStateTrie<A, S>) {
// run the main loop to completion
let result = self.run_inner();
(result, self.trie)
}
/// Inner function to run the sparse trie task to completion.
///
/// See [`Self::run`] for more information.
fn run_inner(&mut self) -> Result<StateRootComputeOutcome, ParallelStateRootError> {
let now = Instant::now();
loop {
crossbeam_channel::select_biased! {
recv(self.proof_result_rx) -> message => {
let Ok(result) = message else {
unreachable!("we own the sender half")
};
self.on_proof_result(result)?;
},
recv(self.updates) -> message => {
let update = match message {
Ok(m) => m,
Err(_) => {
break
}
};
match update {
MultiProofMessage::PrefetchProofs(targets) => {
self.on_prewarm_targets(targets);
}
MultiProofMessage::StateUpdate(_, state) => {
self.on_state_update(state);
}
MultiProofMessage::EmptyProof { sequence_number: _, state } => {
self.on_hashed_state_update(state);
}
MultiProofMessage::BlockAccessList(_) => todo!(),
MultiProofMessage::FinishedStateUpdates => {}
}
}
}
self.process_updates()?;
}
loop {
self.process_updates()?;
if !self.account_updates.is_empty() || !self.storage_updates.is_empty() {
let Ok(result) = self.proof_result_rx.recv() else {
unreachable!("we own the receiver half")
};
self.on_proof_result(result)?;
} else {
break
}
}
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 end = Instant::now();
self.metrics.sparse_trie_final_update_duration_histogram.record(end.duration_since(start));
self.metrics.sparse_trie_total_duration_histogram.record(end.duration_since(now));
Ok(StateRootComputeOutcome { state_root, trie_updates })
}
fn on_prewarm_targets(&mut self, targets: VersionedMultiProofTargets) {
let VersionedMultiProofTargets::V2(targets) = targets else {
unreachable!("sparse trie as cache must only be used with V2 multiproof targets");
};
for target in targets.account_targets {
// Only touch accounts that are not yet present in the updates set.
self.account_updates.entry(target.key()).or_insert(LeafUpdate::Touched);
}
for (address, slots) in targets.storage_targets {
for slot in slots {
// Only touch storages that are not yet present in the updates set.
self.storage_updates
.entry(address)
.or_default()
.entry(slot.key())
.or_insert(LeafUpdate::Touched);
}
// Touch corresponding account leaf to make sure its revealed in accounts trie for
// storage root update.
self.account_updates.entry(address).or_insert(LeafUpdate::Touched);
}
}
/// Processes a state update and encodes all state changes as trie updates.
#[instrument(
level = "debug",
target = "engine::tree::payload_processor::sparse_trie",
skip_all,
fields(accounts = update.len())
)]
fn on_state_update(&mut self, update: EvmState) {
let hashed_state_update = evm_state_to_hashed_post_state(update);
self.on_hashed_state_update(hashed_state_update)
}
/// Processes a hashed state update and encodes all state changes as trie updates.
fn on_hashed_state_update(&mut self, hashed_state_update: HashedPostState) {
for (address, storage) in hashed_state_update.storages {
for (slot, value) in storage.storage {
let encoded = if value.is_zero() {
Vec::new()
} else {
alloy_rlp::encode_fixed_size(&value).to_vec()
};
self.storage_updates
.entry(address)
.or_default()
.insert(slot, LeafUpdate::Changed(encoded));
}
// Make sure account is tracked in `account_updates` so that it is revealed in accounts
// trie for storage root update.
self.account_updates.entry(address).or_insert(LeafUpdate::Touched);
// Make sure account is tracked in `pending_account_updates` so that once storage root
// is computed, it will be updated in the accounts trie.
self.pending_account_updates.entry(address).or_insert(None);
}
for (address, account) in hashed_state_update.accounts {
// Track account as touched.
//
// This might overwrite an existing update, which is fine, because storage root from it
// is already tracked in the trie and can be easily fetched again.
self.account_updates.insert(address, LeafUpdate::Touched);
// Track account in `pending_account_updates` so that once storage root is computed,
// it will be updated in the accounts trie.
self.pending_account_updates.insert(address, Some(account));
}
}
fn on_proof_result(
&mut self,
result: ProofResultMessage,
) -> Result<(), ParallelStateRootError> {
let ProofResult::V2(result) = result.result? else {
unreachable!("sparse trie as cache must only be used wit hmultiproof v2");
};
self.trie.reveal_decoded_multiproof_v2(result).map_err(|e| {
ParallelStateRootError::Other(format!("could not reveal multiproof: {e:?}"))
})
}
/// Applies updates to the sparse trie and dispathes requested multiproof targets.
fn process_updates(&mut self) -> Result<(), ProviderError> {
let mut targets = MultiProofTargetsV2::default();
for (addr, updates) in &mut self.storage_updates {
let trie = self.trie.get_or_create_storage_trie_mut(*addr);
trie.update_leaves(updates, |path, min_len| {
let key = B256::from_slice(&path.pack());
targets
.storage_targets
.entry(*addr)
.or_default()
.push(proof_v2::Target::new(key).with_min_len(min_len));
})
.map_err(|e| ProviderError::other(std::io::Error::other(e.to_string())))?;
// If all storage updates were processed, we can now compute the new storage root.
if updates.is_empty() {
let storage_root =
trie.root().expect("updates are drained, trie should be revealed by now");
// If there is a pending account update for this address with known info, we can
// encode it into proper update right away.
if let Entry::Occupied(entry) = self.pending_account_updates.entry(*addr) &&
entry.get().is_some()
{
let account = entry.remove().expect("just checked, should be Some");
let encoded = if account.is_none_or(|account| account.is_empty()) &&
storage_root == EMPTY_ROOT_HASH
{
Vec::new()
} else {
// TODO: optimize allocation
alloy_rlp::encode(
account.unwrap_or_default().into_trie_account(storage_root),
)
};
self.account_updates.insert(*addr, LeafUpdate::Changed(encoded));
}
}
}
// Now handle pending account updates that can be upgraded to a proper update.
self.pending_account_updates.retain(|addr, account| {
// If account has pending storage updates, it is still pending.
if self.storage_updates.get(addr).is_some_and(|updates| !updates.is_empty()) {
return true;
}
// Get the current account state either from the trie or from latest account update.
let trie_account = if let Some(LeafUpdate::Changed(encoded)) = self.account_updates.get(addr) {
Some(encoded).filter(|encoded| !encoded.is_empty())
} else if self.trie.is_account_revealed(*addr) {
self.trie.get_account_value(addr)
} else {
// Needs to be revealed first
return true;
};
let trie_account = trie_account.map(|value| TrieAccount::decode(&mut &value[..]).expect("invalid account RLP"));
let (account, storage_root) = if let Some(account) = account.take() {
// If account is Some(_) here it means it didn't have any storage updates
// and we can fetch the storage root directly from the account trie.
//
// If it did have storage updates, we would've had processed it above when iterating over storage tries.
let storage_root = trie_account.map(|account| account.storage_root).unwrap_or(EMPTY_ROOT_HASH);
(account, storage_root)
} else {
(trie_account.map(Into::into), self.trie.storage_root(addr).expect("account had storage updates that were applies to its trie, storage root must be revealed by now"))
};
let encoded = if account.is_none_or(|account| account.is_empty()) && storage_root == EMPTY_ROOT_HASH {
Vec::new()
} else {
let account = account.unwrap_or_default().into_trie_account(storage_root);
// TODO: optimize allocation
alloy_rlp::encode(account)
};
self.account_updates.insert(*addr, LeafUpdate::Changed(encoded));
false
});
// Process account trie updates and fill the account targets.
self.trie
.update_leaves(&mut self.account_updates, |path, min_len| {
let key = B256::from_slice(&path.pack());
targets.account_targets.push(proof_v2::Target::new(key).with_min_len(min_len));
})
.map_err(|e| ProviderError::other(std::io::Error::other(e.to_string())))?;
if !targets.is_empty() {
self.proof_worker_handle.dispatch_account_multiproof(AccountMultiproofInput::V2 {
targets,
proof_result_sender: ProofResultContext::new(
self.proof_result_tx.clone(),
0,
HashedPostState::default(),
Instant::now(),
),
})?;
}
Ok(())
}
}
/// Outcome of the state root computation, including the state root itself with
/// the trie updates.
#[derive(Debug)]

View File

@@ -324,6 +324,10 @@ pub struct EngineArgs {
/// Disable cache metrics recording, which can take up to 50ms with large cached state.
#[arg(long = "engine.disable-cache-metrics", default_value_t = DefaultEngineValues::get_global().cache_metrics_disabled)]
pub cache_metrics_disabled: bool,
/// Enable experimental sparse trie as cache feature.
#[arg(long = "engine.experimental-sparse-trie-cache", default_value_t = false)]
pub experimental_sparse_trie_cache: bool,
}
#[allow(deprecated)]
@@ -376,6 +380,7 @@ impl Default for EngineArgs {
account_worker_count,
disable_proof_v2,
cache_metrics_disabled,
experimental_sparse_trie_cache: false,
}
}
}
@@ -412,6 +417,7 @@ impl EngineArgs {
config = config.with_disable_proof_v2(self.disable_proof_v2);
config = config.without_cache_metrics(self.cache_metrics_disabled);
config = config.with_sparse_trie_as_cache(self.experimental_sparse_trie_cache);
config
}
@@ -464,6 +470,7 @@ mod tests {
account_worker_count: Some(8),
disable_proof_v2: false,
cache_metrics_disabled: true,
experimental_sparse_trie_cache: true,
};
let parsed_args = CommandParser::<EngineArgs>::parse_from([
@@ -494,6 +501,7 @@ mod tests {
"--engine.account-worker-count",
"8",
"--engine.disable-cache-metrics",
"--engine.experimental-sparse-trie-cache",
])
.args;

View File

@@ -89,6 +89,12 @@ impl From<revm_state::Account> for Account {
}
}
impl From<TrieAccount> for Account {
fn from(value: TrieAccount) -> Self {
Self { balance: value.balance, nonce: value.nonce, bytecode_hash: Some(value.code_hash) }
}
}
impl InMemorySize for Account {
fn size(&self) -> usize {
size_of::<Self>()

View File

@@ -1062,6 +1062,124 @@ impl SparseTrieExt for ParallelSparseTrie {
nodes_converted
}
fn update_leaves(
&mut self,
updates: &mut alloy_primitives::map::B256Map<reth_trie_sparse::LeafUpdate>,
mut proof_required_fn: impl FnMut(Nibbles, u8),
) -> SparseTrieResult<()> {
use alloy_primitives::map::HashSet;
use reth_trie_sparse::LeafUpdate;
if updates.is_empty() {
return Ok(());
}
// Drain updates into a vec, converting B256 keys to Nibbles paths and determining which
// subtrie each belongs to. Sort by (subtrie_type, path) to group operations for efficient
// batch processing.
let mut ops: Vec<(B256, Nibbles, SparseSubtrieType, LeafUpdate)> = updates
.drain()
.map(|(key, update)| {
let full_path = Nibbles::unpack(key);
let subtrie = SparseSubtrieType::from_path(&full_path);
(key, full_path, subtrie, update)
})
.collect();
ops.sort_unstable_by(|(_, path_a, subtrie_a, _), (_, path_b, subtrie_b, _)| {
subtrie_a.cmp(subtrie_b).then(path_a.cmp(path_b))
});
// Split into upper (short paths) and lower (long paths) subtrie operations.
let num_upper = ops
.iter()
.position(|(_, path, _, _)| !SparseSubtrieType::path_len_is_upper(path.len()))
.unwrap_or(ops.len());
let (upper_ops, lower_ops) = ops.split_at(num_upper);
// Track deduplicated proof requests and failed updates to restore later.
let mut requested: HashSet<(Nibbles, u8)> = HashSet::default();
let mut failures: Vec<(B256, LeafUpdate)> = Vec::new();
// Process upper subtrie operations sequentially since they modify shared trie structure.
for (key, full_path, _, update) in upper_ops {
match self.apply_single_update(full_path, update.clone()) {
Ok(()) => {}
Err(UpdateLeafOutcome::Blinded { blinded_path }) => {
let min_len = Self::compute_min_len(&blinded_path);
if requested.insert((*full_path, min_len)) {
proof_required_fn(*full_path, min_len);
}
failures.push((*key, update.clone()));
}
Err(UpdateLeafOutcome::Fatal(e)) => {
for (k, u) in failures {
updates.insert(k, u);
}
return Err(e);
}
}
}
// Process lower subtrie operations.
if lower_ops.is_empty() {
for (k, u) in failures {
updates.insert(k, u);
}
return Ok(());
}
// Separate Touched (read-only) from Changed (mutating) operations. Touched operations
// can run in parallel since find_leaf takes &self, while Changed must be sequential.
let (touched_ops, changed_ops): (Vec<_>, Vec<_>) =
lower_ops.iter().partition(|(_, _, _, update)| matches!(update, LeafUpdate::Touched));
// Process Touched operations with potential parallelism - only checks path accessibility.
let touched_failures = self.process_touched_ops_parallel(&touched_ops);
for (key, full_path, min_len) in touched_failures {
if requested.insert((full_path, min_len)) {
proof_required_fn(full_path, min_len);
}
failures.push((key, LeafUpdate::Touched));
}
// Process Changed operations sequentially since they mutate the trie structure.
for (key, full_path, _, update) in changed_ops {
match self.apply_single_update(full_path, update.clone()) {
Ok(()) => {}
Err(UpdateLeafOutcome::Blinded { blinded_path }) => {
let min_len = Self::compute_min_len(&blinded_path);
if requested.insert((*full_path, min_len)) {
proof_required_fn(*full_path, min_len);
}
failures.push((*key, update.clone()));
}
Err(UpdateLeafOutcome::Fatal(e)) => {
for (k, u) in failures {
updates.insert(k, u);
}
return Err(e);
}
}
}
// Put failed updates back into the map for caller to retry after revealing proofs.
for (k, u) in failures {
updates.insert(k, u);
}
Ok(())
}
}
/// Outcome of attempting to apply a single leaf update
enum UpdateLeafOutcome {
/// Operation hit a blinded node
Blinded { blinded_path: Nibbles },
/// Fatal error that should be propagated
Fatal(reth_execution_errors::SparseTrieError),
}
impl ParallelSparseTrie {
@@ -1094,6 +1212,115 @@ impl ParallelSparseTrie {
num_changed_keys >= self.parallelism_thresholds.min_updated_nodes
}
/// Computes the `min_len` for a proof request from a blinded path.
fn compute_min_len(blinded_path: &Nibbles) -> u8 {
const MAX_NIBBLE_PATH_LEN: u8 = 64;
(blinded_path.len() as u8).min(MAX_NIBBLE_PATH_LEN)
}
/// Applies a single leaf update, returning Ok(()) on success or the outcome on failure.
///
/// Handles three update variants:
/// - `Changed(empty)`: Remove leaf, restore value on blinded node error
/// - `Changed(value)`: Insert/update leaf, clean up on blinded node error
/// - `Touched`: Check path accessibility without mutation
fn apply_single_update(
&mut self,
full_path: &Nibbles,
update: reth_trie_sparse::LeafUpdate,
) -> Result<(), UpdateLeafOutcome> {
use reth_trie_sparse::{provider::NoRevealProvider, LeafUpdate};
match update {
// Removal: empty value triggers leaf deletion. Snapshot the value first so we can
// restore it if removal fails due to a blinded node.
LeafUpdate::Changed(value) if value.is_empty() => {
let old_value = self.get_leaf_value(full_path).cloned();
match self.remove_leaf(full_path, NoRevealProvider) {
Ok(()) => Ok(()),
Err(e) => {
if let SparseTrieErrorKind::BlindedNode { path, .. } = e.kind() {
if let Some(old) = old_value {
// Mirror get_leaf_value logic: check if lower subtrie exists and
// is non-empty before inserting there
if let Some(subtrie) = self.lower_subtrie_for_path_mut(full_path) &&
!subtrie.is_empty()
{
subtrie.inner.values.insert(*full_path, old);
} else {
self.upper_subtrie.inner.values.insert(*full_path, old);
}
}
Err(UpdateLeafOutcome::Blinded { blinded_path: *path })
} else {
Err(UpdateLeafOutcome::Fatal(e))
}
}
}
}
// Insert/update: non-empty value creates or updates the leaf. Track whether the leaf
// existed so we can clean up if update fails due to a blinded node.
LeafUpdate::Changed(value) => {
let existed = self.get_leaf_value(full_path).is_some();
match self.update_leaf(*full_path, value, NoRevealProvider) {
Ok(()) => Ok(()),
Err(e) => {
if let SparseTrieErrorKind::BlindedNode { path, .. } = e.kind() {
if !existed {
self.upper_subtrie.inner.values.remove(full_path);
if let Some(subtrie) = self.lower_subtrie_for_path_mut(full_path) {
subtrie.inner.values.remove(full_path);
}
}
Err(UpdateLeafOutcome::Blinded { blinded_path: *path })
} else {
Err(UpdateLeafOutcome::Fatal(e))
}
}
}
}
// Touched: read-only check for path accessibility, no mutation.
LeafUpdate::Touched => match self.find_leaf(full_path, None) {
Ok(_) | Err(LeafLookupError::ValueMismatch { .. }) => Ok(()),
Err(LeafLookupError::BlindedNode { path, .. }) => {
Err(UpdateLeafOutcome::Blinded { blinded_path: path })
}
},
}
}
/// Process Touched operations with potential parallelism.
/// Returns a list of `(key, full_path, min_len)` for failures that need proof requests.
fn process_touched_ops_parallel(
&self,
touched_ops: &[&(B256, Nibbles, SparseSubtrieType, reth_trie_sparse::LeafUpdate)],
) -> Vec<(B256, Nibbles, u8)> {
#[cfg(feature = "std")]
if touched_ops.len() >= self.parallelism_thresholds.min_updated_nodes {
use rayon::iter::{IntoParallelIterator, ParallelIterator};
return touched_ops
.into_par_iter()
.filter_map(|(key, full_path, _, _)| match self.find_leaf(full_path, None) {
Err(LeafLookupError::BlindedNode { path, .. }) => {
Some((*key, *full_path, Self::compute_min_len(&path)))
}
Ok(_) | Err(LeafLookupError::ValueMismatch { .. }) => None,
})
.collect();
}
touched_ops
.iter()
.filter_map(|(key, full_path, _, _)| match self.find_leaf(full_path, None) {
Err(LeafLookupError::BlindedNode { path, .. }) => {
Some((*key, *full_path, Self::compute_min_len(&path)))
}
Ok(_) | Err(LeafLookupError::ValueMismatch { .. }) => None,
})
.collect()
}
/// Creates a new revealed sparse trie from the given root node.
///
/// This function initializes the internal structures and then reveals the root.
@@ -7680,4 +7907,615 @@ mod tests {
// The trie should still be functional
let _ = trie.root();
}
// ==================== update_leaves tests ====================
#[test]
fn test_update_leaves_successful_update() {
use alloy_primitives::map::B256Map;
use reth_trie_sparse::LeafUpdate;
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();
// Create update map with a new value for the same key
let new_value = encode_account_value(2);
let mut updates: B256Map<LeafUpdate> = B256Map::default();
updates.insert(b256_key, LeafUpdate::Changed(new_value));
let proof_targets = RefCell::new(Vec::new());
trie.update_leaves(&mut updates, |path, min_len| {
proof_targets.borrow_mut().push((path, min_len));
})
.unwrap();
// Update should succeed: map empty, callback not invoked
assert!(updates.is_empty(), "Update map should be empty after successful update");
assert!(
proof_targets.borrow().is_empty(),
"Callback should not be invoked for revealed paths"
);
}
#[test]
fn test_update_leaves_blinded_node() {
use alloy_primitives::map::B256Map;
use reth_trie_sparse::LeafUpdate;
use std::cell::RefCell;
// Create a trie with a blinded node
// Use a small value that fits in RLP encoding
let small_value = alloy_rlp::encode_fixed_size(&U256::from(1)).to_vec();
let leaf = LeafNode::new(
Nibbles::default(), // short key for RLP encoding
small_value,
);
let branch = TrieNode::Branch(BranchNode::new(
vec![
RlpNode::word_rlp(&B256::repeat_byte(1)), // blinded child at 0
RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(), // revealed at 1
],
TrieMask::new(0b11),
));
let mut trie = ParallelSparseTrie::from_root(
branch.clone(),
Some(BranchNodeMasks {
hash_mask: TrieMask::new(0b01),
tree_mask: TrieMask::default(),
}),
false,
)
.unwrap();
// Reveal only the branch and one child, leaving child 0 as a Hash node
trie.reveal_node(
Nibbles::default(),
branch,
Some(BranchNodeMasks {
hash_mask: TrieMask::default(),
tree_mask: TrieMask::new(0b01),
}),
)
.unwrap();
trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap();
// The path 0x0... is blinded (Hash node)
// Create an update targeting the blinded path using a full B256 key
let b256_key = B256::ZERO; // starts with 0x0...
let new_value = encode_account_value(42);
let mut updates: B256Map<LeafUpdate> = B256Map::default();
updates.insert(b256_key, LeafUpdate::Changed(new_value));
let proof_targets = RefCell::new(Vec::new());
trie.update_leaves(&mut updates, |path, min_len| {
proof_targets.borrow_mut().push((path, min_len));
})
.unwrap();
// Update should remain in map (blinded node)
assert!(!updates.is_empty(), "Update should remain in map when hitting blinded node");
// Callback should be invoked
let targets = proof_targets.borrow();
assert!(!targets.is_empty(), "Callback should be invoked for blinded path");
// min_len should equal the blinded node's path length (1 nibble)
assert_eq!(targets[0].1, 1, "min_len should equal blinded node path length");
}
#[test]
fn test_update_leaves_removal() {
use alloy_primitives::map::B256Map;
use reth_trie_sparse::LeafUpdate;
use std::cell::RefCell;
let provider = DefaultTrieNodeProvider;
let mut trie = ParallelSparseTrie::default();
// Create two leaves so removal doesn't result in empty trie issues
// Use full-length keys
let b256_key1 = B256::with_last_byte(1);
let b256_key2 = B256::with_last_byte(2);
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();
// Create an update to remove key1 (empty value = removal)
let mut updates: B256Map<LeafUpdate> = B256Map::default();
updates.insert(b256_key1, LeafUpdate::Changed(vec![])); // empty = removal
let proof_targets = RefCell::new(Vec::new());
trie.update_leaves(&mut updates, |path, min_len| {
proof_targets.borrow_mut().push((path, min_len));
})
.unwrap();
// Removal should succeed: map empty
assert!(updates.is_empty(), "Update map should be empty after successful removal");
}
#[test]
fn test_update_leaves_removal_blinded() {
use alloy_primitives::map::B256Map;
use reth_trie_sparse::LeafUpdate;
use std::cell::RefCell;
// Create a trie with a blinded node
// Use a small value that fits in RLP encoding
let small_value = alloy_rlp::encode_fixed_size(&U256::from(1)).to_vec();
let leaf = LeafNode::new(
Nibbles::default(), // short key for RLP encoding
small_value,
);
let branch = TrieNode::Branch(BranchNode::new(
vec![
RlpNode::word_rlp(&B256::repeat_byte(1)), // blinded child at 0
RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(), // revealed at 1
],
TrieMask::new(0b11),
));
let mut trie = ParallelSparseTrie::from_root(
branch.clone(),
Some(BranchNodeMasks {
hash_mask: TrieMask::new(0b01),
tree_mask: TrieMask::default(),
}),
false,
)
.unwrap();
trie.reveal_node(
Nibbles::default(),
branch,
Some(BranchNodeMasks {
hash_mask: TrieMask::default(),
tree_mask: TrieMask::new(0b01),
}),
)
.unwrap();
trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap();
// Simulate having a known value behind the blinded node
let b256_key = B256::ZERO; // starts with 0x0...
let full_path = Nibbles::unpack(b256_key);
// Insert the value into the trie's values map (simulating we know about it)
let old_value = encode_account_value(99);
trie.upper_subtrie.inner.values.insert(full_path, old_value.clone());
let mut updates: B256Map<LeafUpdate> = B256Map::default();
updates.insert(b256_key, LeafUpdate::Changed(vec![])); // empty = removal
let proof_targets = RefCell::new(Vec::new());
trie.update_leaves(&mut updates, |path, min_len| {
proof_targets.borrow_mut().push((path, min_len));
})
.unwrap();
// Callback should be invoked
assert!(
!proof_targets.borrow().is_empty(),
"Callback should be invoked when removal hits blinded node"
);
// Update should remain in map
assert!(!updates.is_empty(), "Update should remain in map when removal hits blinded node");
// Original value should be preserved (reverted)
assert_eq!(
trie.upper_subtrie.inner.values.get(&full_path),
Some(&old_value),
"Original value should be preserved after failed removal"
);
}
#[test]
fn test_update_leaves_touched() {
use alloy_primitives::map::B256Map;
use reth_trie_sparse::LeafUpdate;
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();
// Create a Touched update for the existing key
let mut updates: B256Map<LeafUpdate> = B256Map::default();
updates.insert(b256_key, LeafUpdate::Touched);
let proof_targets = RefCell::new(Vec::new());
trie.update_leaves(&mut updates, |path, min_len| {
proof_targets.borrow_mut().push((path, min_len));
})
.unwrap();
// Update should be removed (path is accessible)
assert!(updates.is_empty(), "Touched update should be removed for accessible path");
// No callback
assert!(
proof_targets.borrow().is_empty(),
"Callback should not be invoked for accessible path"
);
}
#[test]
fn test_update_leaves_touched_blinded() {
use alloy_primitives::map::B256Map;
use reth_trie_sparse::LeafUpdate;
use std::cell::RefCell;
// Create a trie with a blinded node
// Use a small value that fits in RLP encoding
let small_value = alloy_rlp::encode_fixed_size(&U256::from(1)).to_vec();
let leaf = LeafNode::new(
Nibbles::default(), // short key for RLP encoding
small_value,
);
let branch = TrieNode::Branch(BranchNode::new(
vec![
RlpNode::word_rlp(&B256::repeat_byte(1)), // blinded child at 0
RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(), // revealed at 1
],
TrieMask::new(0b11),
));
let mut trie = ParallelSparseTrie::from_root(
branch.clone(),
Some(BranchNodeMasks {
hash_mask: TrieMask::new(0b01),
tree_mask: TrieMask::default(),
}),
false,
)
.unwrap();
trie.reveal_node(
Nibbles::default(),
branch,
Some(BranchNodeMasks {
hash_mask: TrieMask::default(),
tree_mask: TrieMask::new(0b01),
}),
)
.unwrap();
trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap();
// Create a Touched update targeting the blinded path using full B256 key
let b256_key = B256::ZERO; // starts with 0x0...
let mut updates: B256Map<LeafUpdate> = B256Map::default();
updates.insert(b256_key, LeafUpdate::Touched);
let proof_targets = RefCell::new(Vec::new());
trie.update_leaves(&mut updates, |path, min_len| {
proof_targets.borrow_mut().push((path, min_len));
})
.unwrap();
// Callback should be invoked
assert!(!proof_targets.borrow().is_empty(), "Callback should be invoked for blinded path");
// Update should remain in map
assert!(!updates.is_empty(), "Touched update should remain in map for blinded path");
}
#[test]
fn test_update_leaves_deduplication() {
use alloy_primitives::map::B256Map;
use reth_trie_sparse::LeafUpdate;
use std::cell::RefCell;
// Create a trie with a blinded node
// Use a small value that fits in RLP encoding
let small_value = alloy_rlp::encode_fixed_size(&U256::from(1)).to_vec();
let leaf = LeafNode::new(
Nibbles::default(), // short key for RLP encoding
small_value,
);
let branch = TrieNode::Branch(BranchNode::new(
vec![
RlpNode::word_rlp(&B256::repeat_byte(1)), // blinded child at 0
RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(), // revealed at 1
],
TrieMask::new(0b11),
));
let mut trie = ParallelSparseTrie::from_root(
branch.clone(),
Some(BranchNodeMasks {
hash_mask: TrieMask::new(0b01),
tree_mask: TrieMask::default(),
}),
false,
)
.unwrap();
trie.reveal_node(
Nibbles::default(),
branch,
Some(BranchNodeMasks {
hash_mask: TrieMask::default(),
tree_mask: TrieMask::new(0b01),
}),
)
.unwrap();
trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap();
// Create multiple updates that would all hit the same blinded node at path 0x0
// Use full B256 keys that all start with 0x0
let b256_key1 = B256::ZERO;
let b256_key2 = B256::with_last_byte(1); // still starts with 0x0
let b256_key3 = B256::with_last_byte(2); // still starts with 0x0
let mut updates: B256Map<LeafUpdate> = B256Map::default();
let value = encode_account_value(42);
updates.insert(b256_key1, LeafUpdate::Changed(value.clone()));
updates.insert(b256_key2, LeafUpdate::Changed(value.clone()));
updates.insert(b256_key3, LeafUpdate::Changed(value));
let proof_targets = RefCell::new(Vec::new());
trie.update_leaves(&mut updates, |path, min_len| {
proof_targets.borrow_mut().push((path, min_len));
})
.unwrap();
// The callback should be invoked 3 times - once for each unique full_path
// The deduplication is by (full_path, min_len), not by blinded node
let targets = proof_targets.borrow();
assert_eq!(targets.len(), 3, "Callback should be invoked for each unique key");
// All should have the same min_len (1) since they all hit blinded node at path 0x0
for (_, min_len) in targets.iter() {
assert_eq!(*min_len, 1, "All should have min_len 1 from blinded node at 0x0");
}
}
#[test]
fn test_update_leaves_multiple_keys() {
use alloy_primitives::map::B256Map;
use reth_trie_sparse::LeafUpdate;
use std::cell::RefCell;
let provider = DefaultTrieNodeProvider;
let mut trie = ParallelSparseTrie::default();
// Create multiple leaves with different prefixes
let keys: Vec<B256> =
vec![B256::repeat_byte(0xAB), B256::repeat_byte(0xCD), B256::repeat_byte(0xEF)];
for (i, &b256_key) in keys.iter().enumerate() {
let key = Nibbles::unpack(b256_key);
let value = encode_account_value(i as u64);
trie.update_leaf(key, value, &provider).unwrap();
}
// Update all values using update_leaves
let mut updates: B256Map<LeafUpdate> = B256Map::default();
for (i, &b256_key) in keys.iter().enumerate() {
updates.insert(b256_key, LeafUpdate::Changed(encode_account_value((i + 100) as u64)));
}
let proof_targets = RefCell::new(Vec::new());
trie.update_leaves(&mut updates, |path, min_len| {
proof_targets.borrow_mut().push((path, min_len));
})
.unwrap();
assert!(updates.is_empty(), "All updates should succeed");
assert!(proof_targets.borrow().is_empty(), "No proofs should be needed");
// Verify all values were updated
for (i, &b256_key) in keys.iter().enumerate() {
let key = Nibbles::unpack(b256_key);
let expected = encode_account_value((i + 100) as u64);
assert_eq!(trie.get_leaf_value(&key), Some(&expected), "Value should be updated");
}
}
#[test]
fn test_update_leaves_batch_removal() {
use alloy_primitives::map::B256Map;
use reth_trie_sparse::LeafUpdate;
use std::cell::RefCell;
let provider = DefaultTrieNodeProvider;
let mut trie = ParallelSparseTrie::default();
// Create multiple leaves
let keys: Vec<B256> =
vec![B256::repeat_byte(0xAB), B256::repeat_byte(0xCD), B256::repeat_byte(0xEF)];
for (i, &b256_key) in keys.iter().enumerate() {
let key = Nibbles::unpack(b256_key);
let value = encode_account_value(i as u64);
trie.update_leaf(key, value, &provider).unwrap();
}
// Remove first key using update_leaves (empty value = removal)
let mut updates: B256Map<LeafUpdate> = B256Map::default();
updates.insert(keys[0], LeafUpdate::Changed(vec![]));
let proof_targets = RefCell::new(Vec::new());
trie.update_leaves(&mut updates, |path, min_len| {
proof_targets.borrow_mut().push((path, min_len));
})
.unwrap();
let key1 = Nibbles::unpack(keys[0]);
let key2 = Nibbles::unpack(keys[1]);
let key3 = Nibbles::unpack(keys[2]);
assert!(updates.is_empty(), "Removal should succeed");
assert!(trie.get_leaf_value(&key1).is_none(), "Value should be removed");
assert!(trie.get_leaf_value(&key2).is_some(), "Other values should remain");
assert!(trie.get_leaf_value(&key3).is_some(), "Other values should remain");
}
#[test]
fn test_update_leaves_rollback_on_blinded_node() {
use alloy_primitives::map::B256Map;
use reth_trie_sparse::LeafUpdate;
use std::cell::RefCell;
// Use the same blinded node setup as test_update_leaves_blinded_node
let small_value = alloy_rlp::encode_fixed_size(&U256::from(1)).to_vec();
let leaf = LeafNode::new(Nibbles::default(), small_value);
let branch = TrieNode::Branch(BranchNode::new(
vec![
RlpNode::word_rlp(&B256::repeat_byte(1)), // blinded child at 0
RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(), // revealed at 1
],
TrieMask::new(0b11),
));
let mut trie = ParallelSparseTrie::from_root(
branch.clone(),
Some(BranchNodeMasks {
hash_mask: TrieMask::new(0b01),
tree_mask: TrieMask::default(),
}),
false,
)
.unwrap();
trie.reveal_node(
Nibbles::default(),
branch,
Some(BranchNodeMasks {
hash_mask: TrieMask::default(),
tree_mask: TrieMask::new(0b01),
}),
)
.unwrap();
trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap();
// Attempt update targeting blinded path
let b256_key = B256::ZERO;
let full_path = Nibbles::unpack(b256_key);
// Check no value exists before
assert!(trie.get_leaf_value(&full_path).is_none(), "No value should exist before update");
let mut updates: B256Map<LeafUpdate> = B256Map::default();
updates.insert(b256_key, LeafUpdate::Changed(encode_account_value(42)));
let proof_targets = RefCell::new(Vec::new());
trie.update_leaves(&mut updates, |path, min_len| {
proof_targets.borrow_mut().push((path, min_len));
})
.unwrap();
// After failed update, trie should be unchanged
assert!(
trie.get_leaf_value(&full_path).is_none(),
"No value should exist after failed update (rollback)"
);
assert!(!updates.is_empty(), "Update should remain in map");
assert!(!proof_targets.borrow().is_empty(), "Proof should be requested");
}
#[test]
fn test_update_leaves_retry_after_reveal() {
use alloy_primitives::map::B256Map;
use reth_trie_sparse::LeafUpdate;
use std::cell::RefCell;
let provider = DefaultTrieNodeProvider;
let mut trie = ParallelSparseTrie::default();
// Create two leaves: one at 0x1... (revealed) and prepare to add at 0x0... (initially
// empty)
let revealed_key = B256::with_last_byte(0x10);
let revealed_path = Nibbles::unpack(revealed_key);
trie.update_leaf(revealed_path, encode_account_value(1), &provider).unwrap();
// Compute root to establish state
let _ = trie.root();
// Now try to add a leaf at a new location
let new_key = B256::with_last_byte(0x20);
let new_value = encode_account_value(42);
let mut updates: B256Map<LeafUpdate> = B256Map::default();
updates.insert(new_key, LeafUpdate::Changed(new_value.clone()));
let proof_targets = RefCell::new(Vec::new());
trie.update_leaves(&mut updates, |path, min_len| {
proof_targets.borrow_mut().push((path, min_len));
})
.unwrap();
// In this case, the trie is fully revealed so update should succeed
assert!(updates.is_empty(), "Update should succeed on revealed trie");
// Verify the value was added
let new_path = Nibbles::unpack(new_key);
assert_eq!(trie.get_leaf_value(&new_path), Some(&new_value));
}
#[test]
fn test_update_leaves_batch_update_and_insert() {
use alloy_primitives::map::B256Map;
use reth_trie_sparse::LeafUpdate;
use std::cell::RefCell;
let provider = DefaultTrieNodeProvider;
let mut trie = ParallelSparseTrie::default();
// Create an existing leaf
let existing_key = B256::repeat_byte(0xAA);
let existing_path = Nibbles::unpack(existing_key);
trie.update_leaf(existing_path, encode_account_value(1), &provider).unwrap();
// Now use update_leaves to both update existing and insert new
let new_key = B256::repeat_byte(0xBB);
let mut updates: B256Map<LeafUpdate> = B256Map::default();
updates.insert(existing_key, LeafUpdate::Changed(encode_account_value(100))); // update
updates.insert(new_key, LeafUpdate::Changed(encode_account_value(200))); // insert
let proof_targets = RefCell::new(Vec::new());
trie.update_leaves(&mut updates, |path, min_len| {
proof_targets.borrow_mut().push((path, min_len));
})
.unwrap();
// Both should succeed
assert!(updates.is_empty(), "All updates should succeed");
assert!(proof_targets.borrow().is_empty(), "No proofs should be needed");
// Verify values
let new_path = Nibbles::unpack(new_key);
assert_eq!(
trie.get_leaf_value(&existing_path),
Some(&encode_account_value(100)),
"Existing value should be updated"
);
assert_eq!(
trie.get_leaf_value(&new_path),
Some(&encode_account_value(200)),
"New value should be inserted"
);
}
}

View File

@@ -64,6 +64,21 @@ impl TrieNodeProvider for DefaultTrieNodeProvider {
}
}
/// 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<Option<RevealedNode>, SparseTrieError> {
Ok(None)
}
}
/// Right pad the path with 0s and return as [`B256`].
#[inline]
pub fn pad_path_to_key(path: &Nibbles) -> B256 {

View File

@@ -1,7 +1,8 @@
use crate::{
provider::{TrieNodeProvider, TrieNodeProviderFactory},
traits::{SparseTrie as SparseTrieTrait, SparseTrieExt},
RevealableSparseTrie, SerialSparseTrie,
RevealableSparseTrie, SerialSparseTrie, DEFAULT_MAX_PRESERVED_STORAGE_TRIES,
DEFAULT_SPARSE_TRIE_PRUNE_DEPTH,
};
use alloc::{collections::VecDeque, vec::Vec};
use alloy_primitives::{
@@ -76,6 +77,18 @@ where
}
}
impl<A, S> ClearedSparseStateTrie<A, S>
where
A: SparseTrieExt + Default,
S: SparseTrieExt + Default + Clone,
{
/// Creates a [`ClearedSparseStateTrie`] by pruning the given [`SparseStateTrie`].
pub fn pruned(mut trie: SparseStateTrie<A, S>) -> Self {
trie.prune(DEFAULT_SPARSE_TRIE_PRUNE_DEPTH, DEFAULT_MAX_PRESERVED_STORAGE_TRIES);
Self(trie)
}
}
#[derive(Debug)]
/// Sparse state trie representing lazy-loaded Ethereum state trie.
pub struct SparseStateTrie<
@@ -224,6 +237,14 @@ where
self.storage.tries.insert(address, storage_trie);
}
/// Returns mutable reference to storage sparse trie, creating a blind one if it doesn't exist.
pub fn get_or_create_storage_trie_mut(
&mut self,
address: B256,
) -> &mut RevealableSparseTrie<S> {
self.storage.get_or_create_trie_mut(address)
}
/// Reveal unknown trie paths from multiproof.
/// NOTE: This method does not extensively validate the proof.
pub fn reveal_multiproof(&mut self, multiproof: MultiProof) -> SparseStateTrieResult<()> {
@@ -714,8 +735,8 @@ where
}
/// Returns storage sparse trie root if the trie has been revealed.
pub fn storage_root(&mut self, account: B256) -> Option<B256> {
self.storage.tries.get_mut(&account).and_then(|trie| trie.root())
pub fn storage_root(&mut self, account: &B256) -> Option<B256> {
self.storage.tries.get_mut(account).and_then(|trie| trie.root())
}
/// Returns mutable reference to the revealed account sparse trie.
@@ -831,6 +852,28 @@ where
Ok(())
}
/// Applies batch leaf updates to the account trie, collecting proof targets for blinded paths.
///
/// See [`SparseTrieExt::update_leaves`] for detailed documentation.
///
/// If the account trie is blind, all keys in `updates` are treated as needing proofs
/// (with `min_len = 0`), and the updates remain in the map for retry after proofs are revealed.
///
/// # Errors
///
/// Returns an error if a non-blinded-node error occurs during update.
pub fn update_leaves(
&mut self,
updates: &mut B256Map<crate::LeafUpdate>,
proof_required_fn: impl FnMut(Nibbles, u8),
) -> SparseStateTrieResult<()>
where
A: SparseTrieExt,
{
self.state.update_leaves(updates, proof_required_fn)?;
Ok(())
}
/// Update the leaf node of a revealed storage trie at the provided address.
pub fn update_storage_leaf(
&mut self,
@@ -1170,6 +1213,13 @@ impl<S: SparseTrieTrait + Clone> StorageTries<S> {
(trie, revealed_paths)
}
// Returns mutable reference to storage sparse trie, creating a blind one if it doesn't exist.
fn get_or_create_trie_mut(&mut self, address: B256) -> &mut RevealableSparseTrie<S> {
self.tries.entry(address).or_insert_with(|| {
self.cleared_tries.pop().unwrap_or_else(|| self.default_trie.clone())
})
}
/// Takes the storage trie for the account from the internal `HashMap`, creating it if it
/// doesn't already exist.
#[cfg(feature = "std")]
@@ -1772,7 +1822,7 @@ mod tests {
&provider_factory,
)
.unwrap();
trie_account_1.storage_root = sparse.storage_root(address_1).unwrap();
trie_account_1.storage_root = sparse.storage_root(&address_1).unwrap();
sparse
.update_account_leaf(
address_path_1,
@@ -1782,7 +1832,7 @@ mod tests {
.unwrap();
sparse.wipe_storage(address_2).unwrap();
trie_account_2.storage_root = sparse.storage_root(address_2).unwrap();
trie_account_2.storage_root = sparse.storage_root(&address_2).unwrap();
sparse
.update_account_leaf(
address_path_2,

View File

@@ -4,7 +4,7 @@ use core::fmt::Debug;
use alloc::{borrow::Cow, vec, vec::Vec};
use alloy_primitives::{
map::{HashMap, HashSet},
map::{B256Map, HashMap, HashSet},
B256,
};
use alloy_trie::BranchNodeCompact;
@@ -13,6 +13,40 @@ use reth_trie_common::{BranchNodeMasks, Nibbles, ProofTrieNode, TrieNode};
use crate::provider::TrieNodeProvider;
/// Describes an update to a leaf in the sparse trie.
///
/// Used with [`SparseTrieExt::update_leaves`] for batch leaf operations that
/// gracefully handle blinded nodes by collecting proof targets instead of failing.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LeafUpdate {
/// The leaf value has been changed to the given RLP-encoded value.
///
/// - Non-empty `Vec`: Updates or inserts the leaf with this value
/// - Empty `Vec`: Removes the leaf from the trie
///
/// # Example
///
/// ```ignore
/// // Update a leaf
/// updates.insert(key, LeafUpdate::Changed(rlp_encoded_value));
///
/// // Remove a leaf
/// updates.insert(key, LeafUpdate::Changed(vec![]));
/// ```
Changed(Vec<u8>),
/// The leaf value is likely changed, but the new value is not yet known.
///
/// Used in prewarming contexts where transactions may execute out-of-order.
/// This variant triggers proof fetching for the path without modifying the
/// trie structure, enabling optimistic revelation of trie nodes that will
/// likely be needed.
///
/// When processed:
/// - If the path is fully revealed: the update is considered complete
/// - If blinded nodes block the path: proof targets are collected
Touched,
}
/// Trait defining common operations for revealed sparse trie implementations.
///
/// This trait abstracts over different sparse trie implementations (serial vs parallel)
@@ -260,6 +294,58 @@ pub trait SparseTrieExt: SparseTrie {
///
/// The number of nodes converted to hash stubs.
fn prune(&mut self, max_depth: usize) -> usize;
/// Applies batch leaf updates to the sparse trie.
///
/// This method enables efficient batch processing of leaf updates while handling
/// blinded (unrevealed) nodes gracefully. Instead of failing on blinded nodes,
/// it collects proof targets that can be used to fetch the necessary proofs.
///
/// # Arguments
///
/// * `updates` - A mutable map of leaf keys to their updates. Successfully applied updates are
/// removed from this map, while updates blocked by blinded nodes remain.
/// * `proof_required_fn` - A callback invoked when a proof is needed. Receives the full path
/// (as `Nibbles`) and `min_len` (minimum proof depth to skip already-revealed nodes). The
/// caller can construct a `proof_v2::Target` using these values.
///
/// # Update Types
///
/// * [`LeafUpdate::Changed`] - Updates or inserts a leaf with the given value. An empty `Vec`
/// indicates removal of the leaf.
/// * [`LeafUpdate::Touched`] - Marks a leaf as likely changed without providing a value. Used
/// for prewarming/optimistic revelation where the actual value isn't yet known.
///
/// # Behavior
///
/// For each update in the map:
/// 1. Attempts to apply the update using the existing revealed trie structure
/// 2. On success: removes the key from `updates`
/// 3. On blinded node: keeps the key in `updates`, invokes `proof_required_fn` with the path
/// and `min_len = blinded_path.len()` (to skip already-revealed prefixes)
///
/// Proof targets are deduplicated - the callback is invoked at most once per unique
/// `(path, min_len)` pair within a single call.
///
/// # Workflow
///
/// ```text
/// 1. Call update_leaves(updates, callback)
/// 2. Callback collects proof targets for blinded paths
/// 3. Fetch proofs for those targets
/// 4. Reveal proofs via reveal_nodes()
/// 5. Call update_leaves(updates, callback) again with remaining updates
/// 6. Repeat until updates map is empty
/// ```
///
/// # Returns
///
/// `Ok(())` on success, or an error if a non-blinded-node error occurs.
fn update_leaves(
&mut self,
updates: &mut B256Map<LeafUpdate>,
proof_required_fn: impl FnMut(Nibbles, u8),
) -> SparseTrieResult<()>;
}
/// Tracks modifications to the sparse trie structure.

View File

@@ -285,6 +285,32 @@ impl<T: SparseTrieTrait> RevealableSparseTrie<T> {
_ => {}
}
}
/// Applies batch leaf updates, collecting proof targets for blinded paths.
///
/// See [`SparseTrieExt::update_leaves`] for detailed documentation.
///
/// If the trie is blind, all keys in `updates` are treated as needing proofs
/// (with `min_len = 0`), and the updates remain in the map for retry after
/// proofs are revealed.
///
/// # Errors
///
/// Returns an error if a non-blinded-node error occurs during update.
pub fn update_leaves(
&mut self,
updates: &mut alloy_primitives::map::B256Map<crate::LeafUpdate>,
mut proof_required_fn: impl FnMut(Nibbles, u8),
) -> SparseTrieResult<()>
where
T: crate::SparseTrieExt,
{
if let Some(revealed) = self.as_revealed_mut() {
revealed.update_leaves(updates, proof_required_fn)
} else {
Ok(())
}
}
}
/// The representation of revealed sparse trie.