diff --git a/crates/engine/tree/src/tree/payload_processor/multiproof.rs b/crates/engine/tree/src/tree/payload_processor/multiproof.rs index 8dd5fc58e7..f2efcff135 100644 --- a/crates/engine/tree/src/tree/payload_processor/multiproof.rs +++ b/crates/engine/tree/src/tree/payload_processor/multiproof.rs @@ -12,8 +12,8 @@ use reth_primitives_traits::FastInstant as Instant; use reth_provider::AccountReader; use reth_revm::state::EvmState; use reth_trie::{ - added_removed_keys::MultiAddedRemovedKeys, proof_v2, HashedPostState, HashedStorage, - MultiProofTargets, + added_removed_keys::{default_added_removed_keys, MultiAddedRemovedKeys}, + proof_v2, HashedPostState, HashedStorage, MultiProofTargets, }; #[cfg(test)] use reth_trie_parallel::stats::ParallelTrieTracker; @@ -919,7 +919,7 @@ impl MultiProofTask { .storages .get(account) .cloned() - .unwrap_or_default(), + .unwrap_or_else(default_added_removed_keys), ); } } diff --git a/crates/trie/common/src/added_removed_keys.rs b/crates/trie/common/src/added_removed_keys.rs index 34a4561dad..00aeef83d1 100644 --- a/crates/trie/common/src/added_removed_keys.rs +++ b/crates/trie/common/src/added_removed_keys.rs @@ -15,7 +15,7 @@ pub struct MultiAddedRemovedKeys { /// Returns [`AddedRemovedKeys`] with default parameters. This is necessary while we are not yet /// tracking added keys. -fn default_added_removed_keys() -> AddedRemovedKeys { +pub fn default_added_removed_keys() -> AddedRemovedKeys { AddedRemovedKeys::default().with_assume_added(true) } diff --git a/crates/trie/common/src/lib.rs b/crates/trie/common/src/lib.rs index 2a3d540016..17531853ad 100644 --- a/crates/trie/common/src/lib.rs +++ b/crates/trie/common/src/lib.rs @@ -37,7 +37,7 @@ mod key; pub use key::{KeccakKeyHasher, KeyHasher}; mod nibbles; -pub use nibbles::{Nibbles, StoredNibbles, StoredNibblesSubKey}; +pub use nibbles::{depth_first_cmp, Nibbles, StoredNibbles, StoredNibblesSubKey}; mod storage; pub use storage::StorageTrieEntry; @@ -48,6 +48,9 @@ pub use subnode::StoredSubNode; mod trie; pub use trie::{BranchNodeMasks, BranchNodeMasksMap, ProofTrieNode}; +mod trie_node_v2; +pub use trie_node_v2::*; + /// The implementation of a container for storing intermediate changes to a trie. /// The container indicates when the trie has been modified. pub mod prefix_set; diff --git a/crates/trie/common/src/nibbles.rs b/crates/trie/common/src/nibbles.rs index 82d710395f..de1e9a321a 100644 --- a/crates/trie/common/src/nibbles.rs +++ b/crates/trie/common/src/nibbles.rs @@ -1,7 +1,28 @@ use alloc::vec::Vec; +use core::cmp::Ordering; use derive_more::Deref; pub use nybbles::Nibbles; +/// Compares two [`Nibbles`] in depth-first order. +/// +/// In depth-first ordering: +/// - Descendants come before their ancestors (children before parents) +/// - Siblings are ordered lexicographically +pub fn depth_first_cmp(a: &Nibbles, b: &Nibbles) -> Ordering { + if a.len() == b.len() { + return a.cmp(b) + } + + let common_prefix_len = a.common_prefix_length(b); + if a.len() == common_prefix_len { + return Ordering::Greater + } else if b.len() == common_prefix_len { + return Ordering::Less + } + + a.get_unchecked(common_prefix_len).cmp(&b.get_unchecked(common_prefix_len)) +} + /// The representation of nibbles of the merkle trie stored in the database. #[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, derive_more::Index)] #[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))] diff --git a/crates/trie/common/src/proofs.rs b/crates/trie/common/src/proofs.rs index c782a075d9..361d134da5 100644 --- a/crates/trie/common/src/proofs.rs +++ b/crates/trie/common/src/proofs.rs @@ -1,6 +1,6 @@ //! Merkle trie proofs. -use crate::{BranchNodeMasksMap, Nibbles, ProofTrieNode, TrieAccount}; +use crate::{BranchNodeMasksMap, Nibbles, ProofTrieNodeV2, TrieAccount}; use alloc::{borrow::Cow, vec::Vec}; use alloy_consensus::constants::KECCAK_EMPTY; use alloy_primitives::{ @@ -448,9 +448,9 @@ impl TryFrom for DecodedMultiProof { #[derive(Clone, Debug, PartialEq, Eq, Default)] pub struct DecodedMultiProofV2 { /// Account trie proof nodes - pub account_proofs: Vec, + pub account_proofs: Vec, /// Storage trie proof nodes indexed by account - pub storage_proofs: B256Map>, + pub storage_proofs: B256Map>, } impl DecodedMultiProofV2 { @@ -477,6 +477,34 @@ impl DecodedMultiProofV2 { } } +impl From for DecodedMultiProofV2 { + fn from(proof: DecodedMultiProof) -> Self { + let account_proofs = + decoded_proof_nodes_to_v2(proof.account_subtree, &proof.branch_node_masks); + let storage_proofs = proof + .storages + .into_iter() + .map(|(address, storage)| { + (address, decoded_proof_nodes_to_v2(storage.subtree, &storage.branch_node_masks)) + }) + .collect(); + Self { account_proofs, storage_proofs } + } +} + +/// Converts a [`DecodedProofNodes`] (path → [`TrieNode`] map) into a `Vec`, +/// merging extension nodes into their child branch nodes. +fn decoded_proof_nodes_to_v2( + nodes: DecodedProofNodes, + masks: &BranchNodeMasksMap, +) -> Vec { + let mut sorted: Vec<_> = nodes.into_inner().into_iter().collect(); + sorted.sort_unstable_by(|a, b| crate::depth_first_cmp(&a.0, &b.0)); + ProofTrieNodeV2::from_sorted_trie_nodes( + sorted.into_iter().map(|(path, node)| (path, node, masks.get(&path).copied())), + ) +} + /// The merkle multiproof of storage trie. #[derive(Clone, Debug, PartialEq, Eq)] pub struct StorageMultiProof { diff --git a/crates/trie/common/src/trie_node_v2.rs b/crates/trie/common/src/trie_node_v2.rs new file mode 100644 index 0000000000..a434b59d29 --- /dev/null +++ b/crates/trie/common/src/trie_node_v2.rs @@ -0,0 +1,219 @@ +//! Version 2 types related to representing nodes in an MPT. + +use crate::BranchNodeMasks; +use alloc::vec::Vec; +use alloy_primitives::hex; +use alloy_rlp::{bytes, Decodable, Encodable, EMPTY_STRING_CODE}; +use alloy_trie::{ + nodes::{BranchNodeRef, ExtensionNode, ExtensionNodeRef, LeafNode, RlpNode, TrieNode}, + Nibbles, TrieMask, +}; +use core::fmt; + +/// Carries all information needed by a sparse trie to reveal a particular node. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProofTrieNodeV2 { + /// Path of the node. + pub path: Nibbles, + /// The node itself. + pub node: TrieNodeV2, + /// Tree and hash masks for the node, if known. + /// Both masks are always set together (from database branch nodes). + pub masks: Option, +} + +impl ProofTrieNodeV2 { + /// Converts an iterator of `(path, TrieNode, masks)` tuples into `Vec`, + /// merging extension nodes into their child branch nodes. + /// + /// The input **must** be sorted in depth-first order (children before parents) for extension + /// merging to work correctly. + pub fn from_sorted_trie_nodes( + iter: impl IntoIterator)>, + ) -> Vec { + let iter = iter.into_iter(); + let mut result = Vec::with_capacity(iter.size_hint().0); + + for (path, node, masks) in iter { + match node { + TrieNode::EmptyRoot => { + result.push(Self { path, node: TrieNodeV2::EmptyRoot, masks }); + } + TrieNode::Leaf(leaf) => { + result.push(Self { path, node: TrieNodeV2::Leaf(leaf), masks }); + } + TrieNode::Branch(branch) => { + result.push(Self { + path, + node: TrieNodeV2::Branch(BranchNodeV2 { + key: Nibbles::new(), + branch_rlp_node: None, + stack: branch.stack, + state_mask: branch.state_mask, + }), + masks, + }); + } + TrieNode::Extension(ext) => { + // In depth-first order, the child branch comes BEFORE the parent + // extension. The child branch should be the last item we added to + // result, at path extension.path + extension.key. + let expected_branch_path = path.join(&ext.key); + + // Check if the last item in result is the child branch + if let Some(last) = result.last_mut() && + last.path == expected_branch_path && + let TrieNodeV2::Branch(branch_v2) = &mut last.node + { + debug_assert!( + branch_v2.key.is_empty(), + "Branch at {:?} already has extension key {:?}", + last.path, + branch_v2.key + ); + branch_v2.key = ext.key; + branch_v2.branch_rlp_node = Some(ext.child); + last.path = path; + } + + // If we reach here, the extension's child is not a branch in the + // result. This happens when the child branch is hashed (not revealed + // in the proof). In V2 format, extension nodes are always combined + // with their child branch, so we skip extension nodes whose child + // isn't revealed. + } + } + } + + result + } +} + +/// Enum representing an MPT trie node. +/// +/// This is a V2 representiation, differing from [`TrieNode`] in that branch and extension nodes are +/// compressed into a single node. +#[derive(PartialEq, Eq, Clone, Debug)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +pub enum TrieNodeV2 { + /// Variant representing empty root node. + EmptyRoot, + /// Variant representing a [`BranchNodeV2`]. + Branch(BranchNodeV2), + /// Variant representing a [`LeafNode`]. + Leaf(LeafNode), + /// Variant representing an [`ExtensionNode`]. + /// + /// This will only be used for extension nodes for which child is not inlined. This variant + /// will never be produced by proof workers that will always reveal a full path to a requested + /// leaf. + Extension(ExtensionNode), +} + +impl Encodable for TrieNodeV2 { + fn encode(&self, out: &mut dyn bytes::BufMut) { + match self { + Self::EmptyRoot => { + out.put_u8(EMPTY_STRING_CODE); + } + Self::Leaf(leaf) => { + leaf.as_ref().encode(out); + } + Self::Branch(branch) => branch.encode(out), + Self::Extension(ext) => { + ext.encode(out); + } + } + } +} + +impl Decodable for TrieNodeV2 { + fn decode(buf: &mut &[u8]) -> Result { + match TrieNode::decode(buf)? { + TrieNode::EmptyRoot => Ok(Self::EmptyRoot), + TrieNode::Leaf(leaf) => Ok(Self::Leaf(leaf)), + TrieNode::Branch(branch) => Ok(Self::Branch(BranchNodeV2::new( + Default::default(), + branch.stack, + branch.state_mask, + None, + ))), + TrieNode::Extension(ext) => { + if ext.child.is_hash() { + Ok(Self::Extension(ext)) + } else { + let Self::Branch(mut branch) = Self::decode(&mut ext.child.as_ref())? else { + return Err(alloy_rlp::Error::Custom( + "extension node child is not a branch", + )); + }; + + branch.key = ext.key; + + Ok(Self::Branch(branch)) + } + } + } + } +} + +/// A branch node in an Ethereum Merkle Patricia Trie. +/// +/// Branch node is a 17-element array consisting of 16 slots that correspond to each hexadecimal +/// character and an additional slot for a value. We do exclude the node value since all paths have +/// a fixed size. +/// +/// This node also encompasses the possible parent extension node of a branch via the `key` field. +#[derive(PartialEq, Eq, Clone, Default)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +pub struct BranchNodeV2 { + /// The key for the branch's parent extension. if key is empty then the branch does not have a + /// parent extension. + pub key: Nibbles, + /// The collection of RLP encoded children. + pub stack: Vec, + /// The bitmask indicating the presence of children at the respective nibble positions + pub state_mask: TrieMask, + /// [`RlpNode`] encoding of the branch node. Always provided when `key` is not empty (i.e this + /// is an extension node). + pub branch_rlp_node: Option, +} + +impl fmt::Debug for BranchNodeV2 { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BranchNode") + .field("key", &self.key) + .field("stack", &self.stack.iter().map(hex::encode).collect::>()) + .field("state_mask", &self.state_mask) + .field("branch_rlp_node", &self.branch_rlp_node) + .finish() + } +} + +impl BranchNodeV2 { + /// Creates a new branch node with the given short key, stack, and state mask. + pub const fn new( + key: Nibbles, + stack: Vec, + state_mask: TrieMask, + branch_rlp_node: Option, + ) -> Self { + Self { key, stack, state_mask, branch_rlp_node } + } +} + +impl Encodable for BranchNodeV2 { + fn encode(&self, out: &mut dyn bytes::BufMut) { + if self.key.is_empty() { + BranchNodeRef::new(&self.stack, self.state_mask).encode(out); + return; + } + + let branch_rlp_node = self + .branch_rlp_node + .as_ref() + .expect("branch_rlp_node must always be present for extension nodes"); + + ExtensionNodeRef::new(&self.key, branch_rlp_node.as_slice()).encode(out); + } +} diff --git a/crates/trie/parallel/src/proof.rs b/crates/trie/parallel/src/proof.rs index c5b972303f..dbbf4f7356 100644 --- a/crates/trie/parallel/src/proof.rs +++ b/crates/trie/parallel/src/proof.rs @@ -1,20 +1,14 @@ use crate::{ metrics::ParallelTrieMetrics, - proof_task::{ - AccountMultiproofInput, ProofResult, ProofResultContext, ProofWorkerHandle, - StorageProofInput, StorageProofResultMessage, - }, + proof_task::{AccountMultiproofInput, ProofResult, ProofResultContext, ProofWorkerHandle}, root::ParallelStateRootError, StorageRootTargets, }; -use alloy_primitives::{map::B256Set, B256}; -use crossbeam_channel::{unbounded as crossbeam_unbounded, Receiver as CrossbeamReceiver}; -use reth_execution_errors::StorageRootError; +use crossbeam_channel::unbounded as crossbeam_unbounded; use reth_primitives_traits::FastInstant as Instant; -use reth_storage_errors::db::DatabaseError; use reth_trie::{ - prefix_set::{PrefixSet, PrefixSetMut, TriePrefixSets, TriePrefixSetsMut}, - DecodedMultiProof, DecodedStorageMultiProof, HashedPostState, MultiProofTargets, Nibbles, + prefix_set::{PrefixSetMut, TriePrefixSets, TriePrefixSetsMut}, + DecodedMultiProof, HashedPostState, MultiProofTargets, Nibbles, }; use reth_trie_common::added_removed_keys::MultiAddedRemovedKeys; use std::sync::Arc; @@ -78,78 +72,6 @@ impl ParallelProof { self.multi_added_removed_keys = multi_added_removed_keys; self } - /// Queues a storage proof task and returns a receiver for the result. - fn send_storage_proof( - &self, - hashed_address: B256, - prefix_set: PrefixSet, - target_slots: B256Set, - ) -> Result, ParallelStateRootError> { - let (result_tx, result_rx) = crossbeam_channel::unbounded(); - - let input = if self.v2_proofs_enabled { - StorageProofInput::new( - hashed_address, - target_slots.into_iter().map(Into::into).collect(), - ) - } else { - StorageProofInput::legacy( - hashed_address, - prefix_set, - target_slots, - self.collect_branch_node_masks, - self.multi_added_removed_keys.clone(), - ) - }; - - self.proof_worker_handle - .dispatch_storage_proof(input, result_tx) - .map_err(|e| ParallelStateRootError::Other(e.to_string()))?; - - Ok(result_rx) - } - - /// Generate a storage multiproof according to the specified targets and hashed address. - pub fn storage_proof( - self, - hashed_address: B256, - target_slots: B256Set, - ) -> Result { - let total_targets = target_slots.len(); - let prefix_set = if self.v2_proofs_enabled { - PrefixSet::default() - } else { - PrefixSetMut::from(target_slots.iter().map(Nibbles::unpack)).freeze() - }; - - trace!( - target: "trie::parallel_proof", - total_targets, - ?hashed_address, - "Starting storage proof generation" - ); - - let receiver = self.send_storage_proof(hashed_address, prefix_set, target_slots)?; - let proof_msg = receiver.recv().map_err(|_| { - ParallelStateRootError::StorageRoot(StorageRootError::Database(DatabaseError::Other( - format!("channel closed for {hashed_address}"), - ))) - })?; - - // Extract storage proof directly from the result - let proof_result = proof_msg.result?; - let storage_proof = Into::>::into(proof_result) - .expect("Partial proofs are not yet supported"); - - trace!( - target: "trie::parallel_proof", - total_targets, - ?hashed_address, - "Storage proof generation completed" - ); - - Ok(storage_proof) - } /// Extends prefix sets with the given multiproof targets and returns the frozen result. /// @@ -251,7 +173,7 @@ mod tests { use alloy_primitives::{ keccak256, map::{B256Set, DefaultHashBuilder, HashMap}, - Address, U256, + Address, B256, U256, }; use rand::Rng; use reth_primitives_traits::{Account, StorageEntry}; diff --git a/crates/trie/parallel/src/proof_task.rs b/crates/trie/parallel/src/proof_task.rs index 0135ddea33..3143c6b445 100644 --- a/crates/trie/parallel/src/proof_task.rs +++ b/crates/trie/parallel/src/proof_task.rs @@ -59,7 +59,7 @@ use reth_trie::{ trie_cursor::{InstrumentedTrieCursor, TrieCursorFactory, TrieCursorMetricsCache}, walker::TrieWalker, DecodedMultiProof, DecodedMultiProofV2, DecodedStorageMultiProof, HashBuilder, HashedPostState, - MultiProofTargets, Nibbles, ProofTrieNode, TRIE_ACCOUNT_RLP_MAX_SIZE, + MultiProofTargets, Nibbles, ProofTrieNodeV2, TRIE_ACCOUNT_RLP_MAX_SIZE, }; use reth_trie_common::{ added_removed_keys::MultiAddedRemovedKeys, @@ -737,7 +737,7 @@ pub(crate) enum StorageProofResult { }, V2 { /// The calculated V2 proof nodes - proof: Vec, + proof: Vec, /// The storage root calculated by the V2 proof root: Option, }, @@ -753,23 +753,6 @@ impl StorageProofResult { } } -impl From for Option { - /// Returns None if the V2 proof result doesn't have a calculated root hash. - fn from(proof_result: StorageProofResult) -> Self { - match proof_result { - StorageProofResult::Legacy { proof } => Some(proof), - StorageProofResult::V2 { proof, root } => root.map(|root| { - let branch_node_masks = proof - .iter() - .filter_map(|node| node.masks.map(|masks| (node.path, masks))) - .collect(); - let subtree = proof.into_iter().map(|node| (node.path, node.node)).collect(); - DecodedStorageMultiProof { root, subtree, branch_node_masks } - }), - } - } -} - /// Message containing a completed storage proof result with metadata. #[derive(Debug)] pub struct StorageProofResultMessage { @@ -1649,17 +1632,10 @@ where proof_msg.hashed_address, hashed_address, "storage worker must return same address" ); - let proof_result = proof_msg.result?; - let Some(root) = proof_result.root() else { - trace!( - target: "trie::proof_task", - ?proof_result, - "Received proof_result without root", - ); - panic!("Partial proofs are not yet supported"); + let StorageProofResult::Legacy { proof } = proof_msg.result? else { + unreachable!("v2 result in legacy worker") }; - let proof = Into::>::into(proof_result) - .expect("Partial proofs are not yet supported (into)"); + let root = proof.root; collected_decoded_storages.insert(hashed_address, proof); root } @@ -1735,9 +1711,9 @@ where let wait_start = Instant::now(); if let Ok(proof_msg) = receiver.recv() { *storage_wait_time += wait_start.elapsed(); - let proof_result = proof_msg.result?; - let proof = Into::>::into(proof_result) - .expect("Partial proofs are not yet supported"); + let StorageProofResult::Legacy { proof } = proof_msg.result? else { + unreachable!("v2 result in legacy worker") + }; collected_decoded_storages.insert(hashed_address, proof); } } diff --git a/crates/trie/parallel/src/value_encoder.rs b/crates/trie/parallel/src/value_encoder.rs index f28918c19f..5e63ffddf8 100644 --- a/crates/trie/parallel/src/value_encoder.rs +++ b/crates/trie/parallel/src/value_encoder.rs @@ -10,7 +10,7 @@ use reth_trie::{ hashed_cursor::HashedStorageCursor, proof_v2::{DeferredValueEncoder, LeafValueEncoder, StorageProofCalculator}, trie_cursor::TrieStorageCursor, - ProofTrieNode, + ProofTrieNodeV2, }; use std::{ rc::Rc, @@ -59,7 +59,7 @@ pub(crate) enum AsyncAccountDeferredValueEncoder { proof_result_rx: Option, DatabaseError>>, /// Shared storage proof results. - storage_proof_results: Rc>>>, + storage_proof_results: Rc>>>, /// Shared stats for tracking wait time and counts. stats: Rc>, /// Shared storage proof calculator for synchronous fallback when dispatched proof has no @@ -226,7 +226,7 @@ pub(crate) struct AsyncAccountValueEncoder { cached_storage_roots: Arc>, /// Tracks storage proof results received from the storage workers. [`Rc`] + [`RefCell`] is /// required because [`DeferredValueEncoder`] cannot have a lifetime. - storage_proof_results: Rc>>>, + storage_proof_results: Rc>>>, /// Shared storage proof calculator for synchronous computation. Reuses cursors and internal /// buffers across multiple storage root calculations. storage_calculator: Rc>>, @@ -267,7 +267,7 @@ impl AsyncAccountValueEncoder { /// been dropped. pub(crate) fn finalize( self, - ) -> Result<(B256Map>, ValueEncoderStats), StateProofError> { + ) -> Result<(B256Map>, ValueEncoderStats), StateProofError> { let mut storage_proof_results = Rc::into_inner(self.storage_proof_results) .expect("no deferred encoders are still allocated") .into_inner(); diff --git a/crates/trie/sparse/src/parallel.rs b/crates/trie/sparse/src/parallel.rs index 516841bbc3..83151e713b 100644 --- a/crates/trie/sparse/src/parallel.rs +++ b/crates/trie/sparse/src/parallel.rs @@ -18,7 +18,7 @@ use reth_primitives_traits::FastInstant as Instant; use reth_trie_common::{ prefix_set::{PrefixSet, PrefixSetMut}, BranchNodeMasks, BranchNodeMasksMap, BranchNodeRef, ExtensionNodeRef, LeafNodeRef, Nibbles, - ProofTrieNode, RlpNode, TrieNode, + ProofTrieNodeV2, RlpNode, TrieNode, TrieNodeV2, }; use smallvec::SmallVec; use tracing::{debug, instrument, trace}; @@ -158,7 +158,7 @@ impl Default for ParallelSparseTrie { impl SparseTrie for ParallelSparseTrie { fn set_root( &mut self, - root: TrieNode, + root: TrieNodeV2, masks: Option, retain_updates: bool, ) -> SparseTrieResult<()> { @@ -177,7 +177,7 @@ impl SparseTrie for ParallelSparseTrie { self.updates = retain_updates.then(Default::default); } - fn reveal_nodes(&mut self, nodes: &mut [ProofTrieNode]) -> SparseTrieResult<()> { + fn reveal_nodes(&mut self, nodes: &mut [ProofTrieNodeV2]) -> SparseTrieResult<()> { if nodes.is_empty() { return Ok(()) } @@ -185,7 +185,7 @@ impl SparseTrie for ParallelSparseTrie { // Sort nodes first by their subtrie, and secondarily by their path. This allows for // grouping nodes by their subtrie using `chunk_by`. nodes.sort_unstable_by( - |ProofTrieNode { path: path_a, .. }, ProofTrieNode { path: path_b, .. }| { + |ProofTrieNodeV2 { path: path_a, .. }, ProofTrieNodeV2 { path: path_b, .. }| { let subtrie_type_a = SparseSubtrieType::from_path(path_a); let subtrie_type_b = SparseSubtrieType::from_path(path_b); subtrie_type_a.cmp(&subtrie_type_b).then_with(|| path_a.cmp(path_b)) @@ -194,9 +194,19 @@ impl SparseTrie for ParallelSparseTrie { // Update the top-level branch node masks. This is simple and can't be done in parallel. self.branch_node_masks.reserve(nodes.len()); - for ProofTrieNode { path, masks, .. } in nodes.iter() { + for ProofTrieNodeV2 { path, masks, node } in nodes.iter() { if let Some(branch_masks) = masks { - self.branch_node_masks.insert(*path, *branch_masks); + // Use proper path for branch nodes by combining path and extension key. + let path = if let TrieNodeV2::Branch(branch) = node && + !branch.key.is_empty() + { + let mut path = *path; + path.extend(&branch.key); + path + } else { + *path + }; + self.branch_node_masks.insert(path, *branch_masks); } } @@ -368,7 +378,7 @@ impl SparseTrie for ParallelSparseTrie { &mut self, full_path: Nibbles, value: Vec, - provider: P, + _provider: P, ) -> SparseTrieResult<()> { trace!( target: "trie::parallel_sparse", @@ -396,8 +406,6 @@ impl SparseTrie for ParallelSparseTrie { return Ok(()); } - let retain_updates = self.updates_enabled(); - // 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()); @@ -414,8 +422,6 @@ impl SparseTrie for ParallelSparseTrie { let mut next = Some(Nibbles::default()); // Track the original node that was modified (path, original_node) for rollback let mut modified_original: Option<(Nibbles, SparseNode)> = None; - // Track inserted branch masks for rollback - let mut inserted_masks: Vec = Vec::new(); // Traverse the upper subtrie to find the node to update or the subtrie to update. // @@ -433,8 +439,7 @@ impl SparseTrie for ParallelSparseTrie { // 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, retain_updates); + 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); @@ -447,90 +452,8 @@ impl SparseTrie for ParallelSparseTrie { // Clear modified_original since we haven't actually modified anything yet modified_original = None; } - LeafUpdateStep::Complete { inserted_nodes, reveal_path } => { + LeafUpdateStep::Complete { inserted_nodes } => { new_nodes.extend(inserted_nodes); - - if let Some(reveal_path) = reveal_path { - let subtrie = self.subtrie_for_path_mut(&reveal_path); - let reveal_masks = if subtrie - .nodes - .get(&reveal_path) - .expect("node must exist") - .is_hash() - { - debug!( - target: "trie::parallel_sparse", - child_path = ?reveal_path, - leaf_full_path = ?full_path, - "Extension node child not revealed in update_leaf, falling back to db", - ); - let revealed_node = match provider.trie_node(&reveal_path) { - Ok(node) => node, - Err(e) => { - self.rollback_insert( - &full_path, - &new_nodes, - &inserted_masks, - modified_original.take(), - ); - return Err(e); - } - }; - if let Some(RevealedNode { node, tree_mask, hash_mask }) = revealed_node - { - let decoded = match TrieNode::decode(&mut &node[..]) { - Ok(d) => d, - Err(e) => { - self.rollback_insert( - &full_path, - &new_nodes, - &inserted_masks, - modified_original.take(), - ); - return Err(e.into()); - } - }; - trace!( - target: "trie::parallel_sparse", - ?reveal_path, - ?decoded, - ?tree_mask, - ?hash_mask, - "Revealing child (from upper)", - ); - let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); - if let Err(e) = subtrie.reveal_node(reveal_path, &decoded, masks) { - self.rollback_insert( - &full_path, - &new_nodes, - &inserted_masks, - modified_original.take(), - ); - return Err(e); - } - masks - } else { - self.rollback_insert( - &full_path, - &new_nodes, - &inserted_masks, - modified_original.take(), - ); - return Err(SparseTrieErrorKind::NodeNotFoundInProvider { - path: reveal_path, - } - .into()) - } - } else { - None - }; - - if let Some(_masks) = reveal_masks { - self.branch_node_masks.insert(reveal_path, _masks); - inserted_masks.push(reveal_path); - } - } - next = None; } LeafUpdateStep::NodeNotFound => { @@ -598,22 +521,12 @@ impl SparseTrie for ParallelSparseTrie { // 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. - match subtrie.update_leaf(full_path, value, provider, retain_updates) { - Ok(Some((revealed_path, revealed_masks))) => { - self.branch_node_masks.insert(revealed_path, revealed_masks); - } - Ok(None) => {} - Err(e) => { - // 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); - } - // Clean up any branch masks that were inserted during upper subtrie traversal - for mask_path in &inserted_masks { - self.branch_node_masks.remove(mask_path); - } - return Err(e); + 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); } } @@ -837,7 +750,6 @@ impl SparseTrie for ParallelSparseTrie { provider, full_path, &remaining_child_path, - true, // recurse_into_extension )?; let (new_branch_node, remove_child) = Self::branch_changes_on_leaf_removal( @@ -1093,7 +1005,6 @@ impl SparseTrie for ParallelSparseTrie { match Self::find_next_to_leaf(&curr_path, curr_node, full_path) { FindNextToLeafOutcome::NotFound => return Ok(LeafLookup::NonExistent), FindNextToLeafOutcome::BlindedNode(hash) => { - // We hit a blinded node - cannot determine if leaf exists return Err(LeafLookupError::BlindedNode { path: curr_path, hash }); } FindNextToLeafOutcome::Found => { @@ -1185,6 +1096,7 @@ impl SparseTrie for ParallelSparseTrie { } // Get children to visit from current node (immutable access) + let mut is_extension = false; let children: SmallVec<[Nibbles; 16]> = { let Some(subtrie) = self.subtrie_for_path(&path) else { continue }; let Some(node) = subtrie.nodes.get(&path) else { continue }; @@ -1196,6 +1108,7 @@ impl SparseTrie for ParallelSparseTrie { SparseNode::Extension { key, .. } => { let mut child = path; child.extend(key); + is_extension = true; SmallVec::from_slice(&[child]) } SparseNode::Branch { state_mask, .. } => { @@ -1216,18 +1129,26 @@ impl SparseTrie for ParallelSparseTrie { // Process children - either continue traversal or prune for child in children { if depth == max_depth { + let path_to_prune = if is_extension { + // If this is a child of extension node, we want to prune the extension node + // itself to preserve invariant of both extension and branch nodes being + // revealed. + path + } else { + child + }; // Check if child has a computed hash and replace inline let hash = self - .subtrie_for_path(&child) - .and_then(|s| s.nodes.get(&child)) + .subtrie_for_path(&path_to_prune) + .and_then(|s| s.nodes.get(&path_to_prune)) .filter(|n| !n.is_hash()) .and_then(|n| n.hash()); if let Some(hash) = hash { // Use untracked access to avoid marking subtrie as modified during pruning - if let Some(subtrie) = self.subtrie_for_path_mut_untracked(&child) { - subtrie.nodes.insert(child, SparseNode::Hash(hash)); - effective_pruned_roots.push((child, hash)); + if let Some(subtrie) = self.subtrie_for_path_mut_untracked(&path_to_prune) { + subtrie.nodes.insert(path_to_prune, SparseNode::Hash(hash)); + effective_pruned_roots.push((path_to_prune, hash)); } } } else { @@ -1457,37 +1378,6 @@ impl ParallelSparseTrie { (target_key, min_len) } - /// Rolls back a partial update by removing the value, removing any inserted nodes, - /// removing any inserted branch masks, and restoring any modified original node. - /// This ensures `update_leaf` is atomic - either it succeeds completely or leaves the trie - /// unchanged. - fn rollback_insert( - &mut self, - full_path: &Nibbles, - inserted_nodes: &[Nibbles], - inserted_masks: &[Nibbles], - modified_original: Option<(Nibbles, SparseNode)>, - ) { - self.upper_subtrie.inner.values.remove(full_path); - for node_path in inserted_nodes { - // Try upper subtrie first - nodes may be there even if path length suggests lower - if self.upper_subtrie.nodes.remove(node_path).is_none() { - // Not in upper, try lower subtrie - if let Some(subtrie) = self.lower_subtrie_for_path_mut(node_path) { - subtrie.nodes.remove(node_path); - } - } - } - // Remove any branch masks that were inserted - for mask_path in inserted_masks { - self.branch_node_masks.remove(mask_path); - } - // Restore the original node that was modified - if let Some((path, original_node)) = modified_original { - self.upper_subtrie.nodes.insert(path, original_node); - } - } - /// Creates a new revealed sparse trie from the given root node. /// /// This function initializes the internal structures and then reveals the root. @@ -1503,7 +1393,7 @@ impl ParallelSparseTrie { /// /// Self if successful, or an error if revealing fails. pub fn from_root( - root: TrieNode, + root: TrieNodeV2, masks: Option, retain_updates: bool, ) -> SparseTrieResult { @@ -1828,20 +1718,16 @@ impl ParallelSparseTrie { if let TrieNode::Extension(ext) = decoded { let mut grandchild_path = *path; grandchild_path.extend(&ext.key); + return self.pre_validate_reveal_chain(&grandchild_path, provider); } + Ok(()) } // Provider cannot reveal this node - operation would fail None => Err(SparseTrieErrorKind::BlindedNode { path: *path, hash: *hash }.into()), }, - // Already-revealed extension: recursively validate its child - Some(SparseNode::Extension { key, .. }) => { - let mut child_path = *path; - child_path.extend(key); - self.pre_validate_reveal_chain(&child_path, provider) - } - // Leaf, Branch, Empty, or missing: no further validation needed + // Leaf, Extension, Branch, Empty, or missing: no further validation needed _ => Ok(()), } } @@ -1860,7 +1746,6 @@ impl ParallelSparseTrie { provider: P, full_path: &Nibbles, // only needed for logs remaining_child_path: &Nibbles, - recurse_into_extension: bool, ) -> SparseTrieResult { let remaining_child_subtrie = self.subtrie_for_path_mut(remaining_child_path); @@ -1879,7 +1764,7 @@ impl ParallelSparseTrie { if let Some(RevealedNode { node, tree_mask, hash_mask }) = provider.trie_node(remaining_child_path)? { - let decoded = TrieNode::decode(&mut &node[..])?; + let decoded = TrieNodeV2::decode(&mut &node[..])?; trace!( target: "trie::parallel_sparse", ?remaining_child_path, @@ -1910,32 +1795,6 @@ impl ParallelSparseTrie { self.branch_node_masks.insert(*remaining_child_path, masks); } - // If `recurse_into_extension` is true, and the remaining child is an extension node, then - // its child will be ensured to be revealed as well. This is required for generation of - // trie updates; without revealing the grandchild branch it's not always possible to know - // if the tree mask bit should be set for the child extension on its parent branch. - if let SparseNode::Extension { key, .. } = &remaining_child_node && - recurse_into_extension - { - let mut remaining_grandchild_path = *remaining_child_path; - remaining_grandchild_path.extend(key); - - trace!( - target: "trie::parallel_sparse", - remaining_grandchild_path = ?remaining_grandchild_path, - child_path = ?remaining_child_path, - leaf_full_path = ?full_path, - "Revealing child of extension node, which is the last remaining child of the branch" - ); - - self.reveal_remaining_child_on_leaf_removal( - provider, - full_path, - &remaining_grandchild_path, - false, // recurse_into_extension - )?; - } - Ok(remaining_child_node) } @@ -2154,7 +2013,7 @@ impl ParallelSparseTrie { fn reveal_upper_node( &mut self, path: Nibbles, - node: &TrieNode, + node: &TrieNodeV2, masks: Option, ) -> SparseTrieResult<()> { // Only reveal nodes that can be reached given the current state of the upper trie. If they @@ -2165,7 +2024,19 @@ impl ParallelSparseTrie { // Exit early if the node was already revealed before. if !self.upper_subtrie.reveal_node(path, node, masks)? { - return Ok(()) + if let TrieNodeV2::Branch(branch) = node { + if branch.key.is_empty() { + return Ok(()); + } + + // We might still potentially need to reveal a child branch node in the lower + // subtrie, even if the upper subtrie already knew about the extension node. + if SparseSubtrieType::path_len_is_upper(path.len() + branch.key.len()) { + return Ok(()) + } + } else { + return Ok(()); + } } // The previous upper_trie.reveal_node call will not have revealed any child nodes via @@ -2173,30 +2044,44 @@ impl ParallelSparseTrie { // here by manually checking the specific cases where this could happen, and calling // reveal_node_or_hash for each. match node { - TrieNode::Branch(branch) => { - // If a branch is at the cutoff level of the trie then it will be in the upper trie, - // but all of its children will be in a lower trie. Check if a child node would be - // in the lower subtrie, and reveal accordingly. - if !SparseSubtrieType::path_len_is_upper(path.len() + 1) { - let mut stack_ptr = branch.as_ref().first_child_index(); - for idx in branch.state_mask.iter() { - let mut child_path = path; + TrieNodeV2::Branch(branch) => { + let mut branch_path = path; + branch_path.extend(&branch.key); + + // If only the parent extension belongs to the upper trie, we need to reveal the + // actual branch node in the corresponding lower subtrie. + if !SparseSubtrieType::path_len_is_upper(branch_path.len()) { + self.lower_subtrie_for_path_mut(&branch_path) + .expect("branch_path must have a lower subtrie") + .reveal_branch( + branch_path, + branch.state_mask, + &branch.stack, + masks, + branch.branch_rlp_node.clone(), + )? + } else if !SparseSubtrieType::path_len_is_upper(branch_path.len() + 1) { + // If a branch is at the cutoff level of the trie then it will be in the upper + // trie, but all of its children will be in a lower trie. + // Check if a child node would be in the lower subtrie, and + // reveal accordingly. + for (stack_ptr, idx) in branch.state_mask.iter().enumerate() { + let mut child_path = branch_path; child_path.push_unchecked(idx); self.lower_subtrie_for_path_mut(&child_path) .expect("child_path must have a lower subtrie") .reveal_node_or_hash(child_path, &branch.stack[stack_ptr])?; - stack_ptr += 1; } } } - TrieNode::Extension(ext) => { + TrieNodeV2::Extension(ext) => { let mut child_path = path; child_path.extend(&ext.key); if let Some(subtrie) = self.lower_subtrie_for_path_mut(&child_path) { subtrie.reveal_node_or_hash(child_path, &ext.child)?; } } - TrieNode::EmptyRoot | TrieNode::Leaf(_) => (), + TrieNodeV2::EmptyRoot | TrieNodeV2::Leaf(_) => (), } Ok(()) @@ -2300,11 +2185,11 @@ impl ParallelSparseTrie { fn is_boundary_leaf_reachable( upper_nodes: &HashMap, path: &Nibbles, - node: &TrieNode, + node: &TrieNodeV2, ) -> bool { debug_assert_eq!(path.len(), UPPER_TRIE_MAX_DEPTH); - if !matches!(node, TrieNode::Leaf(_)) { + if !matches!(node, TrieNodeV2::Leaf(_)) { return true } @@ -2526,24 +2411,17 @@ impl SparseSubtrie { /// /// 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, - provider: impl TrieNodeProvider, - retain_updates: bool, - ) -> SparseTrieResult> { + pub fn update_leaf(&mut self, full_path: Nibbles, value: Vec) -> SparseTrieResult<()> { debug_assert!(full_path.starts_with(&self.path)); // Check if value already exists - if so, just update it (no structural changes needed) if let Entry::Occupied(mut e) = self.inner.values.entry(full_path) { e.insert(value); - return Ok(None) + return Ok(()) } // Here we are starting at the root of the subtrie, and traversing from there. let mut current = Some(self.path); - let mut revealed = None; // Track inserted nodes and modified original for rollback on error let mut inserted_nodes: Vec = Vec::new(); @@ -2557,12 +2435,11 @@ impl SparseSubtrie { modified_original = Some((current_path, node.clone())); } - let step_result = self.update_next_node(current_path, &full_path, retain_updates); + let step_result = self.update_next_node(current_path, &full_path); - // Handle errors from update_next_node - rollback and propagate - if let Err(e) = step_result { + if let Err(err) = step_result { self.rollback_leaf_insert(&full_path, &inserted_nodes, modified_original.take()); - return Err(e); + return Err(err); } match step_result? { @@ -2571,77 +2448,8 @@ impl SparseSubtrie { // Clear modified_original since we haven't actually modified anything yet modified_original = None; } - LeafUpdateStep::Complete { inserted_nodes: new_inserted, reveal_path } => { + LeafUpdateStep::Complete { inserted_nodes: new_inserted } => { inserted_nodes.extend(new_inserted); - - if let Some(reveal_path) = reveal_path && - self.nodes.get(&reveal_path).expect("node must exist").is_hash() - { - debug!( - target: "trie::parallel_sparse", - child_path = ?reveal_path, - leaf_full_path = ?full_path, - "Extension node child not revealed in update_leaf, falling back to db", - ); - let revealed_node = match provider.trie_node(&reveal_path) { - Ok(node) => node, - Err(e) => { - self.rollback_leaf_insert( - &full_path, - &inserted_nodes, - modified_original.take(), - ); - return Err(e); - } - }; - if let Some(RevealedNode { node, tree_mask, hash_mask }) = revealed_node { - let decoded = match TrieNode::decode(&mut &node[..]) { - Ok(d) => d, - Err(e) => { - self.rollback_leaf_insert( - &full_path, - &inserted_nodes, - modified_original.take(), - ); - return Err(e.into()); - } - }; - trace!( - target: "trie::parallel_sparse", - ?reveal_path, - ?decoded, - ?tree_mask, - ?hash_mask, - "Revealing child (from lower)", - ); - let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); - if let Err(e) = self.reveal_node(reveal_path, &decoded, masks) { - self.rollback_leaf_insert( - &full_path, - &inserted_nodes, - modified_original.take(), - ); - return Err(e); - } - - debug_assert_eq!( - revealed, None, - "Only a single blinded node should be revealed during update_leaf" - ); - revealed = masks.map(|masks| (reveal_path, masks)); - } else { - self.rollback_leaf_insert( - &full_path, - &inserted_nodes, - modified_original.take(), - ); - return Err(SparseTrieErrorKind::NodeNotFoundInProvider { - path: reveal_path, - } - .into()) - } - } - current = None; } LeafUpdateStep::NodeNotFound => { @@ -2653,7 +2461,7 @@ impl SparseSubtrie { // Only insert the value after all structural changes succeed self.inner.values.insert(full_path, value); - Ok(revealed) + Ok(()) } /// Rollback structural changes made during a failed leaf insert. @@ -2690,7 +2498,6 @@ impl SparseSubtrie { &mut self, mut current: Nibbles, path: &Nibbles, - retain_updates: bool, ) -> SparseTrieResult { debug_assert!(path.starts_with(&self.path)); debug_assert!(current.starts_with(&self.path)); @@ -2698,13 +2505,14 @@ impl SparseSubtrie { let Some(node) = self.nodes.get_mut(¤t) else { return Ok(LeafUpdateStep::NodeNotFound); }; + match node { SparseNode::Empty => { // We need to insert the node with a different path and key depending on the path of // the subtrie. let path = path.slice(self.path.len()..); *node = SparseNode::new_leaf(path); - Ok(LeafUpdateStep::complete_with_insertions(vec![current], None)) + Ok(LeafUpdateStep::complete_with_insertions(vec![current])) } SparseNode::Hash(hash) => { Err(SparseTrieErrorKind::BlindedNode { path: current, hash: *hash }.into()) @@ -2742,10 +2550,11 @@ impl SparseSubtrie { self.nodes .insert(existing_leaf_path, SparseNode::new_leaf(current.slice(common + 1..))); - Ok(LeafUpdateStep::complete_with_insertions( - vec![branch_path, new_leaf_path, existing_leaf_path], - None, - )) + Ok(LeafUpdateStep::complete_with_insertions(vec![ + branch_path, + new_leaf_path, + existing_leaf_path, + ])) } SparseNode::Extension { key, .. } => { current.extend(key); @@ -2755,11 +2564,6 @@ impl SparseSubtrie { let common = current.common_prefix_length(path); *key = current.slice(current.len() - key.len()..common); - // If branch node updates retention is enabled, we need to query the - // extension node child to later set the hash mask for a parent branch node - // correctly. - let reveal_path = retain_updates.then_some(current); - // create state mask for new branch node // NOTE: this might overwrite the current extension node self.nodes.reserve(3); @@ -2786,7 +2590,7 @@ impl SparseSubtrie { inserted_nodes.push(ext_path); } - return Ok(LeafUpdateStep::complete_with_insertions(inserted_nodes, reveal_path)) + return Ok(LeafUpdateStep::complete_with_insertions(inserted_nodes)) } Ok(LeafUpdateStep::continue_with(current)) @@ -2798,7 +2602,7 @@ impl SparseSubtrie { state_mask.set_bit(nibble); let new_leaf = SparseNode::new_leaf(path.slice(current.len()..)); self.nodes.insert(current, new_leaf); - return Ok(LeafUpdateStep::complete_with_insertions(vec![current], None)) + return Ok(LeafUpdateStep::complete_with_insertions(vec![current])) } // If the nibble is set, we can continue traversing the branch. @@ -2807,18 +2611,95 @@ impl SparseSubtrie { } } + /// Reveals a branch node at the given path. + fn reveal_branch( + &mut self, + path: Nibbles, + state_mask: TrieMask, + children: &[RlpNode], + masks: Option, + rlp_node: Option, + ) -> SparseTrieResult<()> { + match self.nodes.entry(path) { + Entry::Occupied(mut entry) => { + match entry.get() { + SparseNode::Hash(hash) => { + // Replace a hash node with a fully revealed branch node. + entry.insert(SparseNode::Branch { + state_mask, + // Memoize the hash of a previously blinded node in a new branch + // node. + hash: Some(*hash), + store_in_db_trie: Some(masks.is_some_and(|m| { + !m.hash_mask.is_empty() || !m.tree_mask.is_empty() + })), + }); + } + // Branch already revealed, do nothing + _ => return Ok(()), + } + } + Entry::Vacant(entry) => { + entry.insert(SparseNode::Branch { + state_mask, + hash: rlp_node.as_ref().and_then(|n| n.as_hash()), + store_in_db_trie: Some( + masks.is_some_and(|m| !m.hash_mask.is_empty() || !m.tree_mask.is_empty()), + ), + }); + } + } + + // For a branch node, iterate over all children. This must happen second so leaf + // children can check connectivity with parent branch. + for (stack_ptr, idx) in state_mask.iter().enumerate() { + let mut child_path = path; + child_path.push_unchecked(idx); + if Self::is_child_same_level(&path, &child_path) { + // Reveal each child node or hash it has, but only if the child is on + // the same level as the parent. + self.reveal_node_or_hash(child_path, &children[stack_ptr])?; + } + } + + Ok(()) + } + /// Internal implementation of the method of the same name on `ParallelSparseTrie`. fn reveal_node( &mut self, path: Nibbles, - node: &TrieNode, + node: &TrieNodeV2, masks: Option, ) -> SparseTrieResult { debug_assert!(path.starts_with(&self.path)); + let mut skip_extension_reveal = false; + // If the node is already revealed and it's not a hash node, do nothing. if self.nodes.get(&path).is_some_and(|node| !node.is_hash()) { - return Ok(false) + // Make sure that we reveal branch nodes properly even when extension was already + // revealed. + if let TrieNodeV2::Branch(branch) = node && + !branch.key.is_empty() + { + let mut branch_path = path; + branch_path.extend(&branch.key); + + // If branch does not belong to this subtrie, we can exit. + if !Self::is_child_same_level(&path, &branch_path) { + return Ok(false); + } + + // If branch is already revealed, we can exit. + if self.nodes.get(&branch_path).is_some_and(|node| !node.is_hash()) { + return Ok(false); + } + + skip_extension_reveal = true; + } else { + return Ok(false) + } } trace!( @@ -2830,51 +2711,77 @@ impl SparseSubtrie { ); match node { - TrieNode::EmptyRoot => { + TrieNodeV2::EmptyRoot => { // For an empty root, ensure that we are at the root path, and at the upper subtrie. debug_assert!(path.is_empty()); debug_assert!(self.path.is_empty()); self.nodes.insert(path, SparseNode::Empty); } - TrieNode::Branch(branch) => { - // Update the branch node entry in the nodes map, handling cases where a blinded - // node is now replaced with a revealed node. - match self.nodes.entry(path) { - Entry::Occupied(mut entry) => match entry.get() { - // Replace a hash node with a fully revealed branch node. - SparseNode::Hash(hash) => { - entry.insert(SparseNode::Branch { - state_mask: branch.state_mask, - // Memoize the hash of a previously blinded node in a new branch - // node. - hash: Some(*hash), + TrieNodeV2::Branch(branch) => { + if branch.key.is_empty() { + self.reveal_branch( + path, + branch.state_mask, + &branch.stack, + masks, + branch.branch_rlp_node.clone(), + )?; + return Ok(true); + } + + if !skip_extension_reveal { + match self.nodes.entry(path) { + Entry::Occupied(mut entry) => match entry.get() { + SparseNode::Hash(hash) => { + // Replace a hash node with a revealed extension node. + entry.insert(SparseNode::Extension { + key: branch.key, + // Memoize the hash of a previously blinded node in a new + // extension node. + hash: Some(*hash), + // Inherit `store_in_db_trie` from the child branch + // node masks so that the memoized hash can be used + // without needing to fetch the child branch. + store_in_db_trie: Some(masks.is_some_and(|m| { + !m.hash_mask.is_empty() || !m.tree_mask.is_empty() + })), + }); + } + _ => unreachable!("checked that node is either a hash or non-existent"), + }, + Entry::Vacant(entry) => { + entry.insert(SparseNode::Extension { + key: branch.key, + hash: None, + // Inherit `store_in_db_trie` from the child branch + // node masks so that the memoized hash can be used + // without needing to fetch the child branch. store_in_db_trie: Some(masks.is_some_and(|m| { !m.hash_mask.is_empty() || !m.tree_mask.is_empty() })), }); } - _ => unreachable!("checked that node is either a hash or non-existent"), - }, - Entry::Vacant(entry) => { - entry.insert(SparseNode::new_branch(branch.state_mask)); } } - // For a branch node, iterate over all children. This must happen second so leaf - // children can check connectivity with parent branch. - let mut stack_ptr = branch.as_ref().first_child_index(); - for idx in branch.state_mask.iter() { - let mut child_path = path; - child_path.push_unchecked(idx); - if Self::is_child_same_level(&path, &child_path) { - // Reveal each child node or hash it has, but only if the child is on - // the same level as the parent. - self.reveal_node_or_hash(child_path, &branch.stack[stack_ptr])?; - } - stack_ptr += 1; + let mut branch_path = path; + branch_path.extend(&branch.key); + + // Exit early if the actual branch node does not belong to this subtrie. + if !Self::is_child_same_level(&path, &branch_path) { + return Ok(true); } + + // Reveal the actual branch node. + self.reveal_branch( + branch_path, + branch.state_mask, + &branch.stack, + masks, + branch.branch_rlp_node.clone(), + )?; } - TrieNode::Extension(ext) => match self.nodes.entry(path) { + TrieNodeV2::Extension(ext) => match self.nodes.entry(path) { Entry::Occupied(mut entry) => match entry.get() { // Replace a hash node with a revealed extension node. SparseNode::Hash(hash) => { @@ -2902,7 +2809,7 @@ impl SparseSubtrie { } } }, - TrieNode::Leaf(leaf) => { + TrieNodeV2::Leaf(leaf) => { // Skip the reachability check when path.len() == UPPER_TRIE_MAX_DEPTH because // at that boundary the leaf is in the lower subtrie but its parent branch is in // the upper subtrie. The subtrie cannot check connectivity across the upper/lower @@ -2998,7 +2905,7 @@ impl SparseSubtrie { return Ok(()) } - self.reveal_node(path, &TrieNode::decode(&mut &child[..])?, None)?; + self.reveal_node(path, &TrieNodeV2::decode(&mut &child[..])?, None)?; Ok(()) } @@ -3336,7 +3243,6 @@ impl SparseSubtrieInner { if should_set_tree_mask_bit { tree_mask.set_bit(last_child_nibble); } - // Set the hash mask. If a child node is a revealed branch node OR // is a blinded node that has its hash mask bit set according to the // database, set the hash mask bit and save the hash. @@ -3467,8 +3373,6 @@ pub enum LeafUpdateStep { Complete { /// The node paths that were inserted during this step inserted_nodes: Vec, - /// Path to a node which may need to be revealed - reveal_path: Option, }, /// The node was not found #[default] @@ -3482,11 +3386,8 @@ impl LeafUpdateStep { } /// Creates a step indicating completion with inserted nodes - pub const fn complete_with_insertions( - inserted_nodes: Vec, - reveal_path: Option, - ) -> Self { - Self::Complete { inserted_nodes, reveal_path } + pub const fn complete_with_insertions(inserted_nodes: Vec) -> Self { + Self::Complete { inserted_nodes } } } @@ -3718,8 +3619,8 @@ mod tests { prefix_set::PrefixSetMut, proof::{ProofNodes, ProofRetainer}, updates::TrieUpdates, - BranchNode, BranchNodeMasks, BranchNodeMasksMap, ExtensionNode, HashBuilder, LeafNode, - ProofTrieNode, RlpNode, TrieMask, TrieNode, EMPTY_ROOT_HASH, + BranchNodeMasks, BranchNodeMasksMap, BranchNodeV2, ExtensionNode, HashBuilder, LeafNode, + ProofTrieNodeV2, RlpNode, TrieMask, TrieNodeV2, EMPTY_ROOT_HASH, }; use reth_trie_db::DatabaseTrieCursorFactory; use std::collections::{BTreeMap, BTreeSet}; @@ -3941,19 +3842,6 @@ mod tests { self } - fn has_hash(self, path: &Nibbles, expected_hash: &B256) -> Self { - match self.subtrie.nodes.get(path) { - Some(SparseNode::Hash(hash)) => { - assert_eq!( - *hash, *expected_hash, - "Expected hash at {path:?} to be {expected_hash:?}, found {hash:?}", - ); - } - node => panic!("Expected hash node at {path:?}, found {node:?}"), - } - self - } - fn has_value(self, path: &Nibbles, expected_value: &[u8]) -> Self { let actual = self.subtrie.inner.values.get(path); assert_eq!( @@ -3971,12 +3859,15 @@ mod tests { } } - fn create_leaf_node(key: impl AsRef<[u8]>, value_nonce: u64) -> TrieNode { - TrieNode::Leaf(LeafNode::new(Nibbles::from_nibbles(key), encode_account_value(value_nonce))) + fn create_leaf_node(key: impl AsRef<[u8]>, value_nonce: u64) -> TrieNodeV2 { + TrieNodeV2::Leaf(LeafNode::new( + Nibbles::from_nibbles(key), + encode_account_value(value_nonce), + )) } - fn create_extension_node(key: impl AsRef<[u8]>, child_hash: B256) -> TrieNode { - TrieNode::Extension(ExtensionNode::new( + fn create_extension_node(key: impl AsRef<[u8]>, child_hash: B256) -> TrieNodeV2 { + TrieNodeV2::Extension(ExtensionNode::new( Nibbles::from_nibbles(key), RlpNode::word_rlp(&child_hash), )) @@ -3985,7 +3876,7 @@ mod tests { fn create_branch_node_with_children( children_indices: &[u8], child_hashes: impl IntoIterator, - ) -> TrieNode { + ) -> TrieNodeV2 { let mut stack = Vec::new(); let mut state_mask = TrieMask::default(); @@ -3994,7 +3885,7 @@ mod tests { stack.push(hash); } - TrieNode::Branch(BranchNode::new(stack, state_mask)) + TrieNodeV2::Branch(BranchNodeV2::new(Nibbles::default(), stack, state_mask, None)) } /// Calculate the state root by feeding the provided state to the hash builder and retaining the @@ -4113,7 +4004,7 @@ mod tests { let proof_nodes = proof_nodes .into_nodes_sorted() .into_iter() - .map(|(path, node)| (path, TrieNode::decode(&mut node.as_ref()).unwrap())); + .map(|(path, node)| (path, TrieNodeV2::decode(&mut node.as_ref()).unwrap())); let all_sparse_nodes = parallel_sparse_trie_nodes(sparse_trie); @@ -4124,20 +4015,20 @@ mod tests { let equals = match (&proof_node, &sparse_node) { // Both nodes are empty - (TrieNode::EmptyRoot, SparseNode::Empty) => true, + (TrieNodeV2::EmptyRoot, SparseNode::Empty) => true, // Both nodes are branches and have the same state mask ( - TrieNode::Branch(BranchNode { state_mask: proof_state_mask, .. }), + TrieNodeV2::Branch(BranchNodeV2 { state_mask: proof_state_mask, .. }), SparseNode::Branch { state_mask: sparse_state_mask, .. }, ) => proof_state_mask == sparse_state_mask, // Both nodes are extensions and have the same key ( - TrieNode::Extension(ExtensionNode { key: proof_key, .. }), + TrieNodeV2::Extension(ExtensionNode { key: proof_key, .. }), SparseNode::Extension { key: sparse_key, .. }, ) | // Both nodes are leaves and have the same key ( - TrieNode::Leaf(LeafNode { key: proof_key, .. }), + TrieNodeV2::Leaf(LeafNode { key: proof_key, .. }), SparseNode::Leaf { key: sparse_key, .. }, ) => proof_key == sparse_key, // Empty and hash nodes are specific to the sparse trie, skip them @@ -4316,7 +4207,7 @@ mod tests { let node = create_leaf_node([0x2, 0x3], 42); let masks = None; - trie.reveal_nodes(&mut [ProofTrieNode { path, node, masks }]).unwrap(); + trie.reveal_nodes(&mut [ProofTrieNodeV2 { path, node, masks }]).unwrap(); assert_matches!( trie.upper_subtrie.nodes.get(&path), @@ -4340,7 +4231,7 @@ mod tests { let branch_at_1 = create_branch_node_with_children(&[0x2], [RlpNode::word_rlp(&B256::repeat_byte(0xBB))]); let mut trie = ParallelSparseTrie::from_root(root_branch, None, false).unwrap(); - trie.reveal_nodes(&mut [ProofTrieNode { + trie.reveal_nodes(&mut [ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x1]), node: branch_at_1, masks: None, @@ -4352,7 +4243,7 @@ mod tests { let node = create_leaf_node([0x3, 0x4], 42); let masks = None; - trie.reveal_nodes(&mut [ProofTrieNode { path, node, masks }]).unwrap(); + trie.reveal_nodes(&mut [ProofTrieNodeV2 { path, node, masks }]).unwrap(); // Check that the lower subtrie was created let idx = path_subtrie_index_unchecked(&path); @@ -4376,7 +4267,7 @@ mod tests { let node = create_leaf_node([0x4, 0x5], 42); let masks = None; - trie.reveal_nodes(&mut [ProofTrieNode { path, node, masks }]).unwrap(); + trie.reveal_nodes(&mut [ProofTrieNodeV2 { path, node, masks }]).unwrap(); // Check that the lower subtrie's path hasn't changed let idx = path_subtrie_index_unchecked(&path); @@ -4441,7 +4332,7 @@ mod tests { let node = create_extension_node([0x2], child_hash); let masks = None; - trie.reveal_nodes(&mut [ProofTrieNode { path, node, masks }]).unwrap(); + trie.reveal_nodes(&mut [ProofTrieNodeV2 { path, node, masks }]).unwrap(); // Extension node should be in upper trie, hash is memoized from the previous Hash node assert_matches!( @@ -4507,7 +4398,7 @@ mod tests { let node = create_branch_node_with_children(&[0x0, 0x7, 0xf], child_hashes.clone()); let masks = None; - trie.reveal_nodes(&mut [ProofTrieNode { path, node, masks }]).unwrap(); + trie.reveal_nodes(&mut [ProofTrieNodeV2 { path, node, masks }]).unwrap(); // Branch node should be in upper trie, hash is memoized from the previous Hash node assert_matches!( @@ -4570,9 +4461,9 @@ mod tests { // Reveal the existing nodes trie.reveal_nodes(&mut [ - ProofTrieNode { path: branch_path, node: branch_node, masks: None }, - ProofTrieNode { path: leaf_1_path, node: leaf_1, masks: None }, - ProofTrieNode { path: leaf_2_path, node: leaf_2, masks: None }, + ProofTrieNodeV2 { path: branch_path, node: branch_node, masks: None }, + ProofTrieNodeV2 { path: leaf_1_path, node: leaf_1, masks: None }, + ProofTrieNodeV2 { path: leaf_2_path, node: leaf_2, masks: None }, ]) .unwrap(); @@ -5271,185 +5162,6 @@ mod tests { ); } - #[test] - fn test_remove_leaf_remaining_extension_node_child_is_revealed() { - let branch_path = Nibbles::from_nibbles([0x4, 0xf, 0x8, 0x8, 0x0, 0x7]); - let removed_branch_path = Nibbles::from_nibbles([0x4, 0xf, 0x8, 0x8, 0x0, 0x7, 0x2]); - - // Convert the logs into reveal_nodes call on a fresh ParallelSparseTrie - let mut nodes = vec![ - // Branch at 0x4f8807 - ProofTrieNode { - path: branch_path, - node: { - TrieNode::Branch(BranchNode::new( - vec![ - RlpNode::word_rlp(&B256::from(hex!( - "dede882d52f0e0eddfb5b89293a10c87468b4a73acd0d4ae550054a92353f6d5" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "8746f18e465e2eed16117306b6f2eef30bc9d2978aee4a7838255e39c41a3222" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "35a4ea861548af5f0262a9b6d619b4fc88fce6531cbd004eab1530a73f34bbb1" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "47d5c2bf9eea5c1ee027e4740c2b86159074a27d52fd2f6a8a8c86c77e48006f" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "eb76a359b216e1d86b1f2803692a9fe8c3d3f97a9fe6a82b396e30344febc0c1" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "437656f2697f167b23e33cb94acc8550128cfd647fc1579d61e982cb7616b8bc" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "45a1ac2faf15ea8a4da6f921475974e0379f39c3d08166242255a567fa88ce6c" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "7dbb299d714d3dfa593f53bc1b8c66d5c401c30a0b5587b01254a56330361395" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "ae407eb14a74ed951c9949c1867fb9ee9ba5d5b7e03769eaf3f29c687d080429" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "768d0fe1003f0e85d3bc76e4a1fa0827f63b10ca9bca52d56c2b1cceb8eb8b08" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "e5127935143493d5094f4da6e4f7f5a0f62d524fbb61e7bb9fb63d8a166db0f3" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "7f3698297308664fbc1b9e2c41d097fbd57d8f364c394f6ad7c71b10291fbf42" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "4a2bc7e19cec63cb5ef5754add0208959b50bcc79f13a22a370f77b277dbe6db" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "40764b8c48de59258e62a3371909a107e76e1b5e847cfa94dbc857e9fd205103" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "2985dca29a7616920d95c43ab62eb013a40e6a0c88c284471e4c3bd22f3b9b25" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "1b6511f7a385e79477239f7dd4a49f52082ecac05aa5bd0de18b1d55fe69d10c" - ))), - ], - TrieMask::new(0b1111111111111111), - )) - }, - masks: Some(BranchNodeMasks { - hash_mask: TrieMask::new(0b1111111111111111), - tree_mask: TrieMask::new(0b0011110100100101), - }), - }, - // Branch at 0x4f88072 - ProofTrieNode { - path: removed_branch_path, - node: { - let stack = vec![ - RlpNode::word_rlp(&B256::from(hex!( - "15fd4993a41feff1af3b629b32572ab05acddd97c681d82ec2eb89c8a8e3ab9e" - ))), - RlpNode::word_rlp(&B256::from(hex!( - "a272b0b94ced4e6ec7adb41719850cf4a167ad8711d0dda6a810d129258a0d94" - ))), - ]; - let branch_node = BranchNode::new(stack, TrieMask::new(0b0001000000000100)); - TrieNode::Branch(branch_node) - }, - masks: Some(BranchNodeMasks { - hash_mask: TrieMask::new(0b0000000000000000), - tree_mask: TrieMask::new(0b0000000000000100), - }), - }, - // Extension at 0x4f880722 - ProofTrieNode { - path: Nibbles::from_nibbles([0x4, 0xf, 0x8, 0x8, 0x0, 0x7, 0x2, 0x2]), - node: { - let extension_node = ExtensionNode::new( - Nibbles::from_nibbles([0x6]), - RlpNode::word_rlp(&B256::from(hex!( - "56fab2b106a97eae9c7197f86d03bca292da6e0ac725b783082f7d950cc4e0fc" - ))), - ); - TrieNode::Extension(extension_node) - }, - masks: None, - }, - // Leaf at 0x4f88072c - ProofTrieNode { - path: Nibbles::from_nibbles([0x4, 0xf, 0x8, 0x8, 0x0, 0x7, 0x2, 0xc]), - node: { - let leaf_node = LeafNode::new( - Nibbles::from_nibbles([ - 0x0, 0x7, 0x7, 0xf, 0x8, 0x6, 0x6, 0x1, 0x3, 0x0, 0x8, 0x8, 0xd, 0xf, - 0xc, 0xa, 0xe, 0x6, 0x4, 0x8, 0xa, 0xb, 0xe, 0x8, 0x3, 0x1, 0xf, 0xa, - 0xd, 0xc, 0xa, 0x5, 0x5, 0xa, 0xd, 0x4, 0x3, 0xa, 0xb, 0x1, 0x6, 0x5, - 0xd, 0x1, 0x6, 0x8, 0x0, 0xd, 0xd, 0x5, 0x6, 0x7, 0xb, 0x5, 0xd, 0x6, - ]), - hex::decode("8468d3971d").unwrap(), - ); - TrieNode::Leaf(leaf_node) - }, - masks: None, - }, - ]; - - // Create a fresh ParallelSparseTrie - let mut trie = ParallelSparseTrie::from_root( - TrieNode::Extension(ExtensionNode::new( - Nibbles::from_nibbles([0x4, 0xf, 0x8, 0x8, 0x0, 0x7]), - RlpNode::word_rlp(&B256::from(hex!( - "56fab2b106a97eae9c7197f86d03bca292da6e0ac725b783082f7d950cc4e0fc" - ))), - )), - None, - true, - ) - .unwrap(); - - // Call reveal_nodes - trie.reveal_nodes(&mut nodes).unwrap(); - - // Remove the leaf at "0x4f88072c077f86613088dfcae648abe831fadca55ad43ab165d1680dd567b5d6" - let leaf_key = Nibbles::from_nibbles([ - 0x4, 0xf, 0x8, 0x8, 0x0, 0x7, 0x2, 0xc, 0x0, 0x7, 0x7, 0xf, 0x8, 0x6, 0x6, 0x1, 0x3, - 0x0, 0x8, 0x8, 0xd, 0xf, 0xc, 0xa, 0xe, 0x6, 0x4, 0x8, 0xa, 0xb, 0xe, 0x8, 0x3, 0x1, - 0xf, 0xa, 0xd, 0xc, 0xa, 0x5, 0x5, 0xa, 0xd, 0x4, 0x3, 0xa, 0xb, 0x1, 0x6, 0x5, 0xd, - 0x1, 0x6, 0x8, 0x0, 0xd, 0xd, 0x5, 0x6, 0x7, 0xb, 0x5, 0xd, 0x6, - ]); - - let mut provider = MockTrieNodeProvider::new(); - let revealed_branch = create_branch_node_with_children(&[], []); - let mut encoded = Vec::new(); - revealed_branch.encode(&mut encoded); - provider.add_revealed_node( - Nibbles::from_nibbles([0x4, 0xf, 0x8, 0x8, 0x0, 0x7, 0x2, 0x2, 0x6]), - RevealedNode { - node: encoded.into(), - tree_mask: None, - // Give it a fake hashmask so that it appears like it will be stored in the db - hash_mask: Some(TrieMask::new(0b1111)), - }, - ); - - trie.remove_leaf(&leaf_key, provider).unwrap(); - - // Calculate root so that updates are calculated. - trie.root(); - - // Take updates and assert they are correct - let updates = trie.take_updates(); - assert_eq!( - updates.removed_nodes.into_iter().collect::>(), - vec![removed_branch_path] - ); - assert_eq!(updates.updated_nodes.len(), 1); - let updated_node = updates.updated_nodes.get(&branch_path).unwrap(); - - // Second bit must be set, indicating that the extension's child is in the db - assert_eq!(updated_node.tree_mask, TrieMask::new(0b011110100100101),) - } - #[test] fn test_parallel_sparse_trie_root() { // Step 1: Create the trie structure @@ -5495,9 +5207,9 @@ mod tests { // Step 2: Reveal nodes in the trie let mut trie = ParallelSparseTrie::from_root(extension, None, true).unwrap(); trie.reveal_nodes(&mut [ - ProofTrieNode { path: branch_path, node: branch, masks: None }, - ProofTrieNode { path: leaf_1_path, node: leaf_1, masks: None }, - ProofTrieNode { path: leaf_2_path, node: leaf_2, masks: None }, + ProofTrieNodeV2 { path: branch_path, node: branch, masks: None }, + ProofTrieNodeV2 { path: leaf_1_path, node: leaf_1, masks: None }, + ProofTrieNodeV2 { path: leaf_2_path, node: leaf_2, masks: None }, ]) .unwrap(); @@ -6009,12 +5721,14 @@ mod tests { Nibbles::default(), alloy_rlp::encode_fixed_size(&U256::from(1)).to_vec(), ); - let branch = TrieNode::Branch(BranchNode::new( + let branch = TrieNodeV2::Branch(BranchNodeV2::new( + Nibbles::default(), vec![ RlpNode::word_rlp(&B256::repeat_byte(1)), RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(), ], TrieMask::new(0b11), + None, )); let provider = DefaultTrieNodeProvider; @@ -6035,7 +5749,7 @@ mod tests { // └── 1 -> Leaf (Path = 1) sparse .reveal_nodes(&mut [ - ProofTrieNode { + ProofTrieNodeV2 { path: Nibbles::default(), node: branch, masks: Some(BranchNodeMasks { @@ -6043,9 +5757,9 @@ mod tests { tree_mask: TrieMask::new(0b01), }), }, - ProofTrieNode { + ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x1]), - node: TrieNode::Leaf(leaf), + node: TrieNodeV2::Leaf(leaf), masks: None, }, ]) @@ -6064,12 +5778,14 @@ mod tests { Nibbles::default(), alloy_rlp::encode_fixed_size(&U256::from(1)).to_vec(), ); - let branch = TrieNode::Branch(BranchNode::new( + let branch = TrieNodeV2::Branch(BranchNodeV2::new( + Nibbles::default(), vec![ RlpNode::word_rlp(&B256::repeat_byte(1)), RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(), ], TrieMask::new(0b11), + None, )); let provider = DefaultTrieNodeProvider; @@ -6090,7 +5806,7 @@ mod tests { // └── 1 -> Leaf (Path = 1) sparse .reveal_nodes(&mut [ - ProofTrieNode { + ProofTrieNodeV2 { path: Nibbles::default(), node: branch, masks: Some(BranchNodeMasks { @@ -6098,9 +5814,9 @@ mod tests { tree_mask: TrieMask::new(0b01), }), }, - ProofTrieNode { + ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x1]), - node: TrieNode::Leaf(leaf), + node: TrieNodeV2::Leaf(leaf), masks: None, }, ]) @@ -6339,7 +6055,7 @@ mod tests { (None, None) => None, }; let mut sparse = ParallelSparseTrie::from_root( - TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(), + TrieNodeV2::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(), masks, false, ) @@ -6353,14 +6069,14 @@ mod tests { Default::default(), [key1()], ); - let mut revealed_nodes: Vec = hash_builder_proof_nodes + let mut revealed_nodes: Vec = hash_builder_proof_nodes .nodes_sorted() .into_iter() .map(|(path, node)| { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); - ProofTrieNode { path, node: TrieNode::decode(&mut &node[..]).unwrap(), masks } + ProofTrieNodeV2 { path, node: TrieNodeV2::decode(&mut &node[..]).unwrap(), masks } }) .collect(); sparse.reveal_nodes(&mut revealed_nodes).unwrap(); @@ -6368,7 +6084,11 @@ mod tests { // Check that the branch node exists with only two nibbles set assert_eq!( sparse.upper_subtrie.nodes.get(&Nibbles::default()), - Some(&SparseNode::new_branch(0b101.into())) + Some(&SparseNode::Branch { + state_mask: 0b101.into(), + hash: None, + store_in_db_trie: Some(false) + }) ); // Insert the leaf for the second key @@ -6377,7 +6097,11 @@ mod tests { // Check that the branch node was updated and another nibble was set assert_eq!( sparse.upper_subtrie.nodes.get(&Nibbles::default()), - Some(&SparseNode::new_branch(0b111.into())) + Some(&SparseNode::Branch { + state_mask: 0b111.into(), + hash: None, + store_in_db_trie: Some(false) + }) ); // Generate the proof for the third key and reveal it in the sparse trie @@ -6388,14 +6112,14 @@ mod tests { Default::default(), [key3()], ); - let mut revealed_nodes: Vec = hash_builder_proof_nodes + let mut revealed_nodes: Vec = hash_builder_proof_nodes .nodes_sorted() .into_iter() .map(|(path, node)| { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); - ProofTrieNode { path, node: TrieNode::decode(&mut &node[..]).unwrap(), masks } + ProofTrieNodeV2 { path, node: TrieNodeV2::decode(&mut &node[..]).unwrap(), masks } }) .collect(); sparse.reveal_nodes(&mut revealed_nodes).unwrap(); @@ -6403,7 +6127,11 @@ mod tests { // Check that nothing changed in the branch node assert_eq!( sparse.upper_subtrie.nodes.get(&Nibbles::default()), - Some(&SparseNode::new_branch(0b111.into())) + Some(&SparseNode::Branch { + state_mask: 0b111.into(), + hash: None, + store_in_db_trie: Some(false) + }) ); // Generate the nodes for the full trie with all three key using the hash builder, and @@ -6459,7 +6187,7 @@ mod tests { (None, None) => None, }; let mut sparse = ParallelSparseTrie::from_root( - TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(), + TrieNodeV2::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(), masks, false, ) @@ -6474,14 +6202,14 @@ mod tests { Default::default(), [key1(), Nibbles::from_nibbles_unchecked([0x01])], ); - let mut revealed_nodes: Vec = hash_builder_proof_nodes + let mut revealed_nodes: Vec = hash_builder_proof_nodes .nodes_sorted() .into_iter() .map(|(path, node)| { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); - ProofTrieNode { path, node: TrieNode::decode(&mut &node[..]).unwrap(), masks } + ProofTrieNodeV2 { path, node: TrieNodeV2::decode(&mut &node[..]).unwrap(), masks } }) .collect(); sparse.reveal_nodes(&mut revealed_nodes).unwrap(); @@ -6489,7 +6217,11 @@ mod tests { // Check that the branch node exists assert_eq!( sparse.upper_subtrie.nodes.get(&Nibbles::default()), - Some(&SparseNode::new_branch(0b11.into())) + Some(&SparseNode::Branch { + state_mask: 0b11.into(), + hash: None, + store_in_db_trie: Some(true) + }) ); // Remove the leaf for the first key @@ -6509,14 +6241,14 @@ mod tests { Default::default(), [key2()], ); - let mut revealed_nodes: Vec = hash_builder_proof_nodes + let mut revealed_nodes: Vec = hash_builder_proof_nodes .nodes_sorted() .into_iter() .map(|(path, node)| { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); - ProofTrieNode { path, node: TrieNode::decode(&mut &node[..]).unwrap(), masks } + ProofTrieNodeV2 { path, node: TrieNodeV2::decode(&mut &node[..]).unwrap(), masks } }) .collect(); sparse.reveal_nodes(&mut revealed_nodes).unwrap(); @@ -6572,7 +6304,7 @@ mod tests { (None, None) => None, }; let mut sparse = ParallelSparseTrie::from_root( - TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(), + TrieNodeV2::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(), masks, false, ) @@ -6601,14 +6333,14 @@ mod tests { Default::default(), [key1()], ); - let mut revealed_nodes: Vec = hash_builder_proof_nodes + let mut revealed_nodes: Vec = hash_builder_proof_nodes .nodes_sorted() .into_iter() .map(|(path, node)| { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); - ProofTrieNode { path, node: TrieNode::decode(&mut &node[..]).unwrap(), masks } + ProofTrieNodeV2 { path, node: TrieNodeV2::decode(&mut &node[..]).unwrap(), masks } }) .collect(); sparse.reveal_nodes(&mut revealed_nodes).unwrap(); @@ -6623,7 +6355,7 @@ mod tests { #[test] fn test_update_leaf_cross_level() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // Test adding leaves that demonstrate the cross-level behavior // Based on the example: leaves 0x1234, 0x1245, 0x1334, 0x1345 @@ -6703,7 +6435,7 @@ mod tests { #[test] fn test_update_leaf_split_at_level_boundary() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // This test demonstrates what happens when we insert leaves that cause // splitting exactly at the upper/lower trie boundary (2 nibbles). @@ -6750,7 +6482,7 @@ mod tests { #[test] fn test_update_subtrie_with_multiple_leaves() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // First, add multiple leaves that will create a subtrie structure // All leaves share the prefix [0x1, 0x2] to ensure they create a subtrie @@ -6818,7 +6550,7 @@ mod tests { #[test] fn test_update_subtrie_extension_node_subtrie() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // All leaves share the prefix [0x1, 0x2] to ensure they create a subtrie // @@ -6847,7 +6579,7 @@ mod tests { #[test] fn update_subtrie_extension_node_cross_level() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // First, add multiple leaves that will create a subtrie structure // All leaves share the prefix [0x1, 0x2] to ensure they create a branch node and subtrie @@ -6879,7 +6611,7 @@ mod tests { #[test] fn test_update_single_nibble_paths() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // Test edge case: single nibble paths that create branches in upper trie // @@ -6923,7 +6655,7 @@ mod tests { #[test] fn test_update_deep_extension_chain() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // Test edge case: deep extension chains that span multiple levels // @@ -6967,7 +6699,7 @@ mod tests { #[test] fn test_update_branch_with_all_nibbles() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // Test edge case: branch node with all 16 possible nibble children // @@ -7016,7 +6748,7 @@ mod tests { #[test] fn test_update_creates_multiple_subtries() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // Test edge case: updates that create multiple subtries at once // @@ -7069,7 +6801,7 @@ mod tests { #[test] fn test_update_extension_to_branch_transformation() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // Test edge case: extension node transforms to branch when split // @@ -7117,75 +6849,10 @@ mod tests { .has_value(&leaf3_path, &value3); } - #[test] - fn test_update_upper_extension_reveal_lower_hash_node() { - let ctx = ParallelSparseTrieTestContext; - - // Test edge case: extension pointing to hash node that gets updated to branch - // and reveals the hash node from lower trie - // - // Setup: - // Upper trie: - // 0x: Extension { key: 0xAB } - // └── Subtrie (0xAB): pointer - // Lower trie (0xAB): - // 0xAB: Hash - // - // After update: - // Upper trie: - // 0x: Extension { key: 0xA } - // └── 0xA: Branch { state_mask: 0b100000000001 } - // ├── 0xA0: Leaf { value: ... } - // └── 0xAB: pointer - // Lower trie (0xAB): - // 0xAB: Branch { state_mask: 0b11 } - // ├── 0xAB1: Hash - // └── 0xAB2: Hash - - // Create a mock provider that will provide the hash node - let mut provider = MockTrieNodeProvider::new(); - - // Create revealed branch which will get revealed and add it to the mock provider - let child_hashes = [ - RlpNode::word_rlp(&B256::repeat_byte(0x11)), - RlpNode::word_rlp(&B256::repeat_byte(0x22)), - ]; - let revealed_branch = create_branch_node_with_children(&[0x1, 0x2], child_hashes); - let mut encoded = Vec::new(); - revealed_branch.encode(&mut encoded); - provider.add_revealed_node( - Nibbles::from_nibbles([0xA, 0xB]), - RevealedNode { node: encoded.into(), tree_mask: None, hash_mask: None }, - ); - - let mut trie = new_test_trie( - [ - (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0xA, 0xB]))), - (Nibbles::from_nibbles([0xA, 0xB]), SparseNode::Hash(B256::repeat_byte(0x42))), - ] - .into_iter(), - ); - - // Now add a leaf that will force the hash node to become a branch - let (leaf_path, value) = ctx.create_test_leaf([0xA, 0x0], 1); - trie.update_leaf(leaf_path, value, provider).unwrap(); - - // Verify the structure: extension should now terminate in a branch on the upper trie - ctx.assert_upper_subtrie(&trie) - .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0xA])) - .has_branch(&Nibbles::from_nibbles([0xA]), &[0x0, 0xB]); - - // Verify the lower trie now has a branch structure - ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0xA, 0xB])) - .has_branch(&Nibbles::from_nibbles([0xA, 0xB]), &[0x1, 0x2]) - .has_hash(&Nibbles::from_nibbles([0xA, 0xB, 0x1]), &B256::repeat_byte(0x11)) - .has_hash(&Nibbles::from_nibbles([0xA, 0xB, 0x2]), &B256::repeat_byte(0x22)); - } - #[test] fn test_update_long_shared_prefix_at_boundary() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // Test edge case: leaves with long shared prefix that ends exactly at 2-nibble boundary // @@ -7228,7 +6895,7 @@ mod tests { #[test] fn test_update_branch_to_extension_collapse() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // Test creating a trie with leaves that share a long common prefix // @@ -7273,7 +6940,7 @@ mod tests { let (new_leaf3_path, new_value3) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x6], 12); // Clear and add new leaves - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + 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(); @@ -7299,7 +6966,7 @@ mod tests { #[test] fn test_update_shared_prefix_patterns() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // Test edge case: different patterns of shared prefixes // @@ -7342,7 +7009,7 @@ mod tests { #[test] fn test_progressive_branch_creation() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // Test starting with a single leaf and progressively adding leaves // that create branch nodes at shorter and shorter paths @@ -7450,7 +7117,7 @@ mod tests { #[test] fn test_update_max_depth_paths() { let ctx = ParallelSparseTrieTestContext; - let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNodeV2::EmptyRoot, None, true).unwrap(); // Test edge case: very long paths (64 nibbles - max for addresses/storage) // @@ -7513,9 +7180,11 @@ mod tests { .map(|hex_str| RlpNode::from_raw_rlp(&hex_str[..]).unwrap()) .collect(); - let root_branch_node = BranchNode::new( + let root_branch_node = BranchNodeV2::new( + Default::default(), root_branch_rlp_stack, TrieMask::new(0b1111111111111111), // state_mask: all 16 children present + None, ); let root_branch_masks = Some(BranchNodeMasks { @@ -7524,7 +7193,7 @@ mod tests { }); let mut trie = ParallelSparseTrie::from_root( - TrieNode::Branch(root_branch_node), + TrieNodeV2::Branch(root_branch_node), root_branch_masks, true, ) @@ -7555,9 +7224,11 @@ mod tests { .map(|hex_str| RlpNode::from_raw_rlp(&hex_str[..]).unwrap()) .collect(); - let branch_0x3_node = BranchNode::new( + let branch_0x3_node = BranchNodeV2::new( + Default::default(), branch_0x3_rlp_stack, TrieMask::new(0b1111111111111111), // state_mask: all 16 children present + None, ); let branch_0x3_masks = Some(BranchNodeMasks { @@ -7576,12 +7247,16 @@ mod tests { let leaf_masks = None; trie.reveal_nodes(&mut [ - ProofTrieNode { + ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x3]), - node: TrieNode::Branch(branch_0x3_node), + node: TrieNodeV2::Branch(branch_0x3_node), masks: branch_0x3_masks, }, - ProofTrieNode { path: leaf_path, node: TrieNode::Leaf(leaf_node), masks: leaf_masks }, + ProofTrieNodeV2 { + path: leaf_path, + node: TrieNodeV2::Leaf(leaf_node), + masks: leaf_masks, + }, ]) .unwrap(); @@ -7886,7 +7561,7 @@ mod tests { // Reveal the trie structure using ProofTrieNode let mut proof_nodes = vec![ - ProofTrieNode { + ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x3]), node: branch_0x3_node, masks: Some(BranchNodeMasks { @@ -7894,7 +7569,7 @@ mod tests { hash_mask: TrieMask::new(65535), }), }, - ProofTrieNode { + ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x3, 0x1]), node: branch_0x31_node, masks: Some(BranchNodeMasks { @@ -7907,7 +7582,7 @@ mod tests { // Create a sparse trie and reveal nodes let mut trie = ParallelSparseTrie::default() .with_root( - TrieNode::Extension(ExtensionNode { + TrieNodeV2::Extension(ExtensionNode { key: Nibbles::from_nibbles([0x3]), child: RlpNode::word_rlp(&B256::ZERO), }), @@ -7963,7 +7638,7 @@ mod tests { let branch_at_1 = create_branch_node_with_children(&[0x2], [RlpNode::word_rlp(&B256::repeat_byte(0xBB))]); let mut trie = ParallelSparseTrie::from_root(root_branch, None, false).unwrap(); - trie.reveal_nodes(&mut [ProofTrieNode { + trie.reveal_nodes(&mut [ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x1]), node: branch_at_1, masks: None, @@ -7976,7 +7651,7 @@ mod tests { let leaf_node = create_leaf_node(leaf_key.to_vec(), 42); // Reveal the leaf node - trie.reveal_nodes(&mut [ProofTrieNode { path: leaf_path, node: leaf_node, masks: None }]) + trie.reveal_nodes(&mut [ProofTrieNodeV2 { path: leaf_path, node: leaf_node, masks: None }]) .unwrap(); // The full path is leaf_path + leaf_key @@ -8013,7 +7688,7 @@ mod tests { // Create an empty trie with an empty root let mut trie = ParallelSparseTrie::default() - .with_root(TrieNode::EmptyRoot, None, false) + .with_root(TrieNodeV2::EmptyRoot, None, false) .expect("root revealed"); // Create a full 64-nibble path (like a real account hash) @@ -8044,7 +7719,7 @@ mod tests { // Create an empty trie let mut trie = ParallelSparseTrie::default() - .with_root(TrieNode::EmptyRoot, None, false) + .with_root(TrieNodeV2::EmptyRoot, None, false) .expect("root revealed"); // Insert first leaf - will be at root level (upper subtrie) @@ -8069,7 +7744,7 @@ mod tests { // Simulate a storage trie with a single slot let mut trie = ParallelSparseTrie::default() - .with_root(TrieNode::EmptyRoot, None, false) + .with_root(TrieNodeV2::EmptyRoot, None, false) .expect("root revealed"); // Single storage slot - leaf will be at root (depth 0) @@ -8510,12 +8185,14 @@ mod tests { Nibbles::default(), // short key for RLP encoding small_value, ); - let branch = TrieNode::Branch(BranchNode::new( + let branch = TrieNodeV2::Branch(BranchNodeV2::new( + Nibbles::default(), 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), + None, )); let mut trie = ParallelSparseTrie::from_root( @@ -8538,7 +8215,7 @@ mod tests { }), ) .unwrap(); - trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap(); + trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNodeV2::Leaf(leaf), None).unwrap(); // The path 0x0... is blinded (Hash node) // Create an update targeting the blinded path using a full B256 key @@ -8619,12 +8296,14 @@ mod tests { Nibbles::default(), // short key for RLP encoding small_value, ); - let branch = TrieNode::Branch(BranchNode::new( + let branch = TrieNodeV2::Branch(BranchNodeV2::new( + Nibbles::default(), 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), + None, )); let mut trie = ParallelSparseTrie::from_root( @@ -8646,7 +8325,7 @@ mod tests { }), ) .unwrap(); - trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap(); + trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNodeV2::Leaf(leaf), None).unwrap(); // Simulate having a known value behind the blinded node let b256_key = B256::ZERO; // starts with 0x0... @@ -8701,12 +8380,14 @@ mod tests { // - Child at nibble 1: a revealed Leaf 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( + let branch = TrieNodeV2::Branch(BranchNodeV2::new( + Nibbles::default(), vec![ RlpNode::word_rlp(&B256::repeat_byte(1)), // blinded child at nibble 0 RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(), /* leaf at nibble 1 */ ], TrieMask::new(0b11), + None, )); let mut trie = ParallelSparseTrie::from_root( @@ -8729,7 +8410,7 @@ mod tests { }), ) .unwrap(); - trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap(); + trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNodeV2::Leaf(leaf), None).unwrap(); // Insert the leaf's value into the values map for the revealed leaf // Use B256 key that starts with nibble 1 (0x10 has first nibble = 1) @@ -8898,12 +8579,14 @@ mod tests { Nibbles::default(), // short key for RLP encoding small_value, ); - let branch = TrieNode::Branch(BranchNode::new( + let branch = TrieNodeV2::Branch(BranchNodeV2::new( + Nibbles::default(), 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), + None, )); let mut trie = ParallelSparseTrie::from_root( @@ -8925,7 +8608,7 @@ mod tests { }), ) .unwrap(); - trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap(); + trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNodeV2::Leaf(leaf), None).unwrap(); // Create a Touched update targeting the blinded path using full B256 key let b256_key = B256::ZERO; // starts with 0x0... @@ -8967,12 +8650,14 @@ mod tests { Nibbles::default(), // short key for RLP encoding small_value, ); - let branch = TrieNode::Branch(BranchNode::new( + let branch = TrieNodeV2::Branch(BranchNodeV2::new( + Nibbles::default(), 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), + None, )); let mut trie = ParallelSparseTrie::from_root( @@ -8994,7 +8679,7 @@ mod tests { }), ) .unwrap(); - trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap(); + trie.reveal_node(Nibbles::from_nibbles([0x1]), TrieNodeV2::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 @@ -9026,109 +8711,6 @@ mod tests { } } - #[test] - fn test_update_leaves_node_not_found_in_provider_atomicity() { - use crate::LeafUpdate; - use alloy_primitives::map::B256Map; - use std::cell::RefCell; - - // Create a trie with retain_updates enabled (this triggers the code path that - // can return NodeNotFoundInProvider when an extension node's child needs revealing). - // - // Structure: Extension at root -> Hash node (blinded child) - // When we try to insert a new leaf that would split the extension, with retain_updates - // enabled, it tries to reveal the hash child via the provider. With NoRevealProvider, - // this returns NodeNotFoundInProvider. - - let child_hash = B256::repeat_byte(0xAB); - let extension = TrieNode::Extension(ExtensionNode::new( - Nibbles::from_nibbles([0x1, 0x2, 0x3]), - RlpNode::word_rlp(&child_hash), - )); - - // Create trie with retain_updates = true - let mut trie = - ParallelSparseTrie::from_root(extension, None, true).expect("from_root failed"); - - // Record state before update_leaves - let prefix_set_len_before = trie.prefix_set.len(); - let node_count_before = trie.upper_subtrie.nodes.len() + - trie.lower_subtries - .iter() - .filter_map(|s| s.as_revealed_ref()) - .map(|s| s.nodes.len()) - .sum::(); - let value_count_before = trie.upper_subtrie.inner.values.len() + - trie.lower_subtries - .iter() - .filter_map(|s| s.as_revealed_ref()) - .map(|s| s.inner.values.len()) - .sum::(); - - // Create an update that would cause an extension split. - // The key starts with 0x1 but diverges from 0x123... at the second nibble. - let b256_key = { - let mut k = B256::ZERO; - k.0[0] = 0x14; // nibbles: 1, 4 - matches first nibble, diverges at second - k - }; - - let new_value = encode_account_value(42); - let mut updates: B256Map = 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)); - }) - .expect("update_leaves should succeed"); - - // Assert: update remains in map (NodeNotFoundInProvider is retriable) - assert!( - !updates.is_empty(), - "Update should remain in map when NodeNotFoundInProvider occurs" - ); - assert!( - updates.contains_key(&b256_key), - "The specific key should be re-inserted for retry" - ); - - // Assert: callback was invoked - let targets = proof_targets.borrow(); - assert!(!targets.is_empty(), "Callback should be invoked for NodeNotFoundInProvider"); - - // Assert: prefix_set unchanged (atomic - no partial state) - assert_eq!( - trie.prefix_set.len(), - prefix_set_len_before, - "prefix_set should be unchanged after atomic failure" - ); - - // Assert: node count unchanged (no structural changes persisted) - let node_count_after = trie.upper_subtrie.nodes.len() + - trie.lower_subtries - .iter() - .filter_map(|s| s.as_revealed_ref()) - .map(|s| s.nodes.len()) - .sum::(); - assert_eq!( - node_count_before, node_count_after, - "Node count should be unchanged after atomic failure" - ); - - // Assert: value count unchanged (no values left dangling) - let value_count_after = trie.upper_subtrie.inner.values.len() + - trie.lower_subtries - .iter() - .filter_map(|s| s.as_revealed_ref()) - .map(|s| s.inner.values.len()) - .sum::(); - assert_eq!( - value_count_before, value_count_after, - "Value count should be unchanged after atomic failure (no dangling values)" - ); - } - #[test] fn test_nibbles_to_padded_b256() { // Empty nibbles should produce all zeros @@ -9179,12 +8761,12 @@ mod tests { let branch_at_5 = create_branch_node_with_children(&[0x6], [RlpNode::word_rlp(&B256::repeat_byte(0xDD))]); trie.reveal_nodes(&mut [ - ProofTrieNode { + ProofTrieNodeV2 { path: Nibbles::from_nibbles_unchecked([0x1]), node: branch_at_1, masks: None, }, - ProofTrieNode { + ProofTrieNodeV2 { path: Nibbles::from_nibbles_unchecked([0x5]), node: branch_at_5, masks: None, @@ -9193,17 +8775,17 @@ mod tests { .unwrap(); let mut nodes = vec![ - ProofTrieNode { + ProofTrieNodeV2 { path: Nibbles::from_nibbles_unchecked([0x1, 0x2]), - node: TrieNode::Leaf(LeafNode { + node: TrieNodeV2::Leaf(LeafNode { key: Nibbles::from_nibbles_unchecked([0x3, 0x4]), value: vec![1, 2, 3], }), masks: None, }, - ProofTrieNode { + ProofTrieNodeV2 { path: Nibbles::from_nibbles_unchecked([0x5, 0x6]), - node: TrieNode::Leaf(LeafNode { + node: TrieNodeV2::Leaf(LeafNode { key: Nibbles::from_nibbles_unchecked([0x7, 0x8]), value: vec![4, 5, 6], }), diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index a43bc971f3..321f297d29 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -3,23 +3,20 @@ use crate::{ traits::SparseTrie as SparseTrieTrait, ParallelSparseTrie, RevealableSparseTrie, }; -use alloc::{collections::VecDeque, vec::Vec}; +use alloc::vec::Vec; use alloy_primitives::{ map::{B256Map, B256Set, HashSet}, - Bytes, B256, + B256, }; use alloy_rlp::{Decodable, Encodable}; -use alloy_trie::proof::DecodedProofNodes; use reth_execution_errors::{SparseStateTrieErrorKind, SparseStateTrieResult, SparseTrieErrorKind}; use reth_primitives_traits::Account; #[cfg(feature = "std")] use reth_primitives_traits::FastInstant as Instant; use reth_trie_common::{ - proof::ProofNodes, updates::{StorageTrieUpdates, TrieUpdates}, - BranchNodeMasks, BranchNodeMasksMap, DecodedMultiProof, DecodedStorageMultiProof, MultiProof, - Nibbles, ProofTrieNode, RlpNode, StorageMultiProof, TrieAccount, TrieNode, EMPTY_ROOT_HASH, - TRIE_ACCOUNT_RLP_MAX_SIZE, + BranchNodeMasks, DecodedMultiProof, MultiProof, Nibbles, ProofTrieNodeV2, TrieAccount, + TrieNodeV2, EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE, }; #[cfg(feature = "std")] use tracing::debug; @@ -32,7 +29,7 @@ use tracing::{instrument, trace}; #[derive(Debug, Default)] pub struct DeferredDrops { /// Each nodes reveal operation creates a new buffer, uses it, and pushes it here. - pub proof_nodes_bufs: Vec>, + pub proof_nodes_bufs: Vec>, } #[derive(Debug)] @@ -269,83 +266,7 @@ where &mut self, multiproof: DecodedMultiProof, ) -> SparseStateTrieResult<()> { - let DecodedMultiProof { account_subtree, storages, branch_node_masks } = multiproof; - - // first reveal the account proof nodes - self.reveal_decoded_account_multiproof(account_subtree, branch_node_masks)?; - - #[cfg(not(feature = "std"))] - // If nostd then serially reveal storage proof nodes for each storage trie - { - for (account, storage_subtree) in storages { - self.reveal_decoded_storage_multiproof(account, storage_subtree)?; - // Mark this storage trie as hot (accessed this tick) - self.storage.modifications.mark_accessed(account); - } - - Ok(()) - } - - #[cfg(feature = "std")] - // If std then reveal storage proofs in parallel - { - use rayon::iter::ParallelIterator; - use reth_primitives_traits::ParallelBridgeBuffered; - - let retain_updates = self.retain_updates; - - // Process all storage trie revealings in parallel, having first removed the - // `reveal_nodes` tracking and `RevealableSparseTrie`s for each account from their - // HashMaps. These will be returned after processing. - let results: Vec<_> = storages - .into_iter() - .map(|(account, storage_subtree)| { - let revealed_nodes = self.storage.take_or_create_revealed_paths(&account); - let trie = self.storage.take_or_create_trie(&account); - (account, storage_subtree, revealed_nodes, trie) - }) - .par_bridge_buffered() - .map(|(account, storage_subtree, mut revealed_nodes, mut trie)| { - let mut bufs = Vec::new(); - let result = Self::reveal_decoded_storage_multiproof_inner( - account, - storage_subtree, - &mut revealed_nodes, - &mut trie, - &mut bufs, - retain_updates, - ); - - (account, revealed_nodes, trie, result, bufs) - }) - .collect(); - - // Return `revealed_nodes` and `RevealableSparseTrie` for each account, incrementing - // metrics and returning the last error seen if any. - let mut any_err = Ok(()); - for (account, revealed_nodes, trie, result, bufs) in results { - self.storage.revealed_paths.insert(account, revealed_nodes); - self.storage.tries.insert(account, trie); - // Mark this storage trie as hot (accessed this tick) - self.storage.modifications.mark_accessed(account); - if let Ok(_metric_values) = result { - #[cfg(feature = "metrics")] - { - self.metrics - .increment_total_storage_nodes(_metric_values.total_nodes as u64); - self.metrics - .increment_skipped_storage_nodes(_metric_values.skipped_nodes as u64); - } - } else { - any_err = result.map(|_| ()); - } - - // Keep buffers for deferred dropping - self.deferred_drops.proof_nodes_bufs.extend(bufs); - } - - any_err - } + self.reveal_decoded_multiproof_v2(multiproof.into()) } /// Reveals a V2 decoded multiproof. @@ -442,66 +363,13 @@ where } } - /// Reveals an account multiproof. - pub fn reveal_account_multiproof( - &mut self, - account_subtree: ProofNodes, - branch_node_masks: BranchNodeMasksMap, - ) -> SparseStateTrieResult<()> { - // decode the multiproof first - let decoded_multiproof = account_subtree.try_into()?; - self.reveal_decoded_account_multiproof(decoded_multiproof, branch_node_masks) - } - - /// Reveals a decoded account multiproof. - pub fn reveal_decoded_account_multiproof( - &mut self, - account_subtree: DecodedProofNodes, - branch_node_masks: BranchNodeMasksMap, - ) -> SparseStateTrieResult<()> { - let FilterMappedProofNodes { - root_node, - mut nodes, - new_nodes, - metric_values: _metric_values, - } = filter_map_revealed_nodes( - account_subtree, - &mut self.revealed_account_paths, - &branch_node_masks, - )?; - #[cfg(feature = "metrics")] - { - self.metrics.increment_total_account_nodes(_metric_values.total_nodes as u64); - self.metrics.increment_skipped_account_nodes(_metric_values.skipped_nodes as u64); - } - - if let Some(root_node) = root_node { - // Reveal root node if it wasn't already. - trace!(target: "trie::sparse", ?root_node, "Revealing root account node"); - let trie = - self.state.reveal_root(root_node.node, root_node.masks, self.retain_updates)?; - - // Reserve the capacity for new nodes ahead of time, if the trie implementation - // supports doing so. - trie.reserve_nodes(new_nodes); - - trace!(target: "trie::sparse", total_nodes = ?nodes.len(), "Revealing account nodes"); - trie.reveal_nodes(&mut nodes)?; - } - - // Keep buffer for deferred dropping - self.deferred_drops.proof_nodes_bufs.push(nodes); - - Ok(()) - } - /// Reveals account proof nodes from a V2 proof. /// /// V2 proofs already include the masks in the `ProofTrieNode` structure, /// so no separate masks map is needed. pub fn reveal_account_v2_proof_nodes( &mut self, - mut nodes: Vec, + mut nodes: Vec, ) -> SparseStateTrieResult<()> { if self.skip_proof_node_filtering { let capacity = estimate_v2_proof_capacity(&nodes); @@ -562,7 +430,7 @@ where pub fn reveal_storage_v2_proof_nodes( &mut self, account: B256, - nodes: Vec, + nodes: Vec, ) -> SparseStateTrieResult<()> { let (trie, revealed_paths) = self.storage.get_trie_and_revealed_paths_mut(account); let _metric_values = Self::reveal_storage_v2_proof_nodes_inner( @@ -588,10 +456,10 @@ where /// designed to handle a variety of associated public functions. fn reveal_storage_v2_proof_nodes_inner( account: B256, - mut nodes: Vec, + mut nodes: Vec, revealed_nodes: &mut HashSet, trie: &mut RevealableSparseTrie, - bufs: &mut Vec>, + bufs: &mut Vec>, retain_updates: bool, skip_filtering: bool, ) -> SparseStateTrieResult { @@ -636,180 +504,6 @@ where Ok(metric_values) } - /// Reveals a storage multiproof for the given address. - pub fn reveal_storage_multiproof( - &mut self, - account: B256, - storage_subtree: StorageMultiProof, - ) -> SparseStateTrieResult<()> { - // decode the multiproof first - let decoded_multiproof = storage_subtree.try_into()?; - self.reveal_decoded_storage_multiproof(account, decoded_multiproof) - } - - /// Reveals a decoded storage multiproof for the given address. - pub fn reveal_decoded_storage_multiproof( - &mut self, - account: B256, - storage_subtree: DecodedStorageMultiProof, - ) -> SparseStateTrieResult<()> { - let (trie, revealed_paths) = self.storage.get_trie_and_revealed_paths_mut(account); - let _metric_values = Self::reveal_decoded_storage_multiproof_inner( - account, - storage_subtree, - revealed_paths, - trie, - &mut self.deferred_drops.proof_nodes_bufs, - self.retain_updates, - )?; - - #[cfg(feature = "metrics")] - { - self.metrics.increment_total_storage_nodes(_metric_values.total_nodes as u64); - self.metrics.increment_skipped_storage_nodes(_metric_values.skipped_nodes as u64); - } - - Ok(()) - } - - /// Reveals a decoded storage multiproof for the given address. This is internal static function - /// is designed to handle a variety of associated public functions. - fn reveal_decoded_storage_multiproof_inner( - account: B256, - storage_subtree: DecodedStorageMultiProof, - revealed_nodes: &mut HashSet, - trie: &mut RevealableSparseTrie, - bufs: &mut Vec>, - retain_updates: bool, - ) -> SparseStateTrieResult { - let FilterMappedProofNodes { root_node, mut nodes, new_nodes, metric_values } = - filter_map_revealed_nodes( - storage_subtree.subtree, - revealed_nodes, - &storage_subtree.branch_node_masks, - )?; - - if let Some(root_node) = root_node { - // Reveal root node if it wasn't already. - trace!(target: "trie::sparse", ?account, ?root_node, "Revealing root storage node"); - let trie = trie.reveal_root(root_node.node, root_node.masks, retain_updates)?; - - // Reserve the capacity for new nodes ahead of time, if the trie implementation - // supports doing so. - trie.reserve_nodes(new_nodes); - - trace!(target: "trie::sparse", ?account, total_nodes = ?nodes.len(), "Revealing storage nodes"); - trie.reveal_nodes(&mut nodes)?; - } - - // Keep buffer for deferred dropping - bufs.push(nodes); - - Ok(metric_values) - } - - /// Reveal state witness with the given state root. - /// The state witness is expected to be a map of `keccak(rlp(node)): rlp(node).` - /// NOTE: This method does not extensively validate the witness. - pub fn reveal_witness( - &mut self, - state_root: B256, - witness: &B256Map, - ) -> SparseStateTrieResult<()> { - // Create a `(hash, path, maybe_account)` queue for traversing witness trie nodes - // starting from the root node. - let mut queue = VecDeque::from([(state_root, Nibbles::default(), None)]); - - while let Some((hash, path, maybe_account)) = queue.pop_front() { - // Retrieve the trie node and decode it. - let Some(trie_node_bytes) = witness.get(&hash) else { continue }; - let trie_node = TrieNode::decode(&mut &trie_node_bytes[..])?; - - // Push children nodes into the queue. - match &trie_node { - TrieNode::Branch(branch) => { - for (idx, maybe_child) in branch.as_ref().children() { - if let Some(child_hash) = maybe_child.and_then(RlpNode::as_hash) { - let mut child_path = path; - child_path.push_unchecked(idx); - queue.push_back((child_hash, child_path, maybe_account)); - } - } - } - TrieNode::Extension(ext) => { - if let Some(child_hash) = ext.child.as_hash() { - let mut child_path = path; - child_path.extend(&ext.key); - queue.push_back((child_hash, child_path, maybe_account)); - } - } - TrieNode::Leaf(leaf) => { - let mut full_path = path; - full_path.extend(&leaf.key); - if maybe_account.is_none() { - let hashed_address = B256::from_slice(&full_path.pack()); - let account = TrieAccount::decode(&mut &leaf.value[..])?; - if account.storage_root != EMPTY_ROOT_HASH { - queue.push_back(( - account.storage_root, - Nibbles::default(), - Some(hashed_address), - )); - } - } - } - TrieNode::EmptyRoot => {} // nothing to do here - }; - - // Reveal the node itself. - if let Some(account) = maybe_account { - // Check that the path was not already revealed. - if self - .storage - .revealed_paths - .get(&account) - .is_none_or(|paths| !paths.contains(&path)) - { - let retain_updates = self.retain_updates; - let (storage_trie_entry, revealed_storage_paths) = - self.storage.get_trie_and_revealed_paths_mut(account); - - if path.is_empty() { - // Handle special storage state root node case. - storage_trie_entry.reveal_root(trie_node, None, retain_updates)?; - } else { - // Reveal non-root storage trie node. - storage_trie_entry - .as_revealed_mut() - .ok_or(SparseTrieErrorKind::Blind)? - .reveal_node(path, trie_node, None)?; - } - - // Track the revealed path. - revealed_storage_paths.insert(path); - } - } - // Check that the path was not already revealed. - else if !self.revealed_account_paths.contains(&path) { - if path.is_empty() { - // Handle special state root node case. - self.state.reveal_root(trie_node, None, self.retain_updates)?; - } else { - // Reveal non-root state trie node. - self.state - .as_revealed_mut() - .ok_or(SparseTrieErrorKind::Blind)? - .reveal_node(path, trie_node, None)?; - } - - // Track the revealed path. - self.revealed_account_paths.insert(path); - } - } - - Ok(()) - } - /// Wipe the storage trie at the provided address. pub fn wipe_storage(&mut self, address: B256) -> SparseStateTrieResult<()> { if let Some(trie) = self.storage.tries.get_mut(&address) { @@ -846,11 +540,11 @@ where .account_node_provider() .trie_node(&Nibbles::default())? .map(|node| { - TrieNode::decode(&mut &node.node[..]) + TrieNodeV2::decode(&mut &node.node[..]) .map(|decoded| (decoded, node.hash_mask, node.tree_mask)) }) .transpose()? - .unwrap_or((TrieNode::EmptyRoot, None, None)); + .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) } @@ -1491,97 +1185,13 @@ struct ProofNodesMetricValues { skipped_nodes: usize, } -/// Result of [`filter_map_revealed_nodes`]. -#[derive(Debug, PartialEq, Eq)] -struct FilterMappedProofNodes { - /// Root node which was pulled out of the original node set to be handled specially. - root_node: Option, - /// Filtered, decoded and unsorted proof nodes. Root node is removed. - nodes: Vec, - /// Number of new nodes that will be revealed. This includes all children of branch nodes, even - /// if they are not in the proof. - new_nodes: usize, - /// Values which are being returned so they can be incremented into metrics. - metric_values: ProofNodesMetricValues, -} - -/// Filters the decoded nodes that are already revealed, maps them to `SparseTrieNode`s, -/// separates the root node if present, and returns additional information about the number of -/// total, skipped, and new nodes. -fn filter_map_revealed_nodes( - proof_nodes: DecodedProofNodes, - revealed_nodes: &mut HashSet, - branch_node_masks: &BranchNodeMasksMap, -) -> SparseStateTrieResult { - let mut result = FilterMappedProofNodes { - root_node: None, - nodes: Vec::with_capacity(proof_nodes.len()), - new_nodes: 0, - metric_values: Default::default(), - }; - - let proof_nodes_len = proof_nodes.len(); - for (path, proof_node) in proof_nodes.into_inner() { - result.metric_values.total_nodes += 1; - - let is_root = path.is_empty(); - - // If the node is already revealed, skip it. We don't ever skip the root node, nor do we add - // it to `revealed_nodes`. - if !is_root && !revealed_nodes.insert(path) { - result.metric_values.skipped_nodes += 1; - continue - } - - result.new_nodes += 1; - - // Extract hash/tree masks based on the node type (only branch nodes have masks). At the - // same time increase the new_nodes counter if the node is a type which has children. - let masks = match &proof_node { - TrieNode::Branch(branch) => { - // If it's a branch node, increase the number of new nodes by the number of children - // according to the state mask. - result.new_nodes += branch.state_mask.count_ones() as usize; - branch_node_masks.get(&path).copied() - } - TrieNode::Extension(_) => { - // There is always exactly one child of an extension node. - result.new_nodes += 1; - None - } - _ => None, - }; - - let node = ProofTrieNode { path, node: proof_node, masks }; - - if is_root { - // Perform sanity check. - if matches!(node.node, TrieNode::EmptyRoot) && proof_nodes_len > 1 { - return Err(SparseStateTrieErrorKind::InvalidRootNode { - path, - node: alloy_rlp::encode(&node.node).into(), - } - .into()) - } - - result.root_node = Some(node); - - continue - } - - result.nodes.push(node); - } - - Ok(result) -} - /// Result of [`filter_revealed_v2_proof_nodes`]. #[derive(Debug, PartialEq, Eq)] struct FilteredV2ProofNodes { /// Root node which was pulled out of the original node set to be handled specially. - root_node: Option, + root_node: Option, /// Filtered proof nodes. Root node is removed. - nodes: Vec, + nodes: Vec, /// Number of new nodes that will be revealed. This includes all children of branch nodes, even /// if they are not in the proof. new_nodes: usize, @@ -1594,19 +1204,13 @@ struct FilteredV2ProofNodes { /// This counts nodes and their children (for branch and extension nodes) to provide /// proper capacity hints for `reserve_nodes`. Used when `skip_proof_node_filtering` is /// enabled and no filtering is performed. -fn estimate_v2_proof_capacity(nodes: &[ProofTrieNode]) -> usize { +fn estimate_v2_proof_capacity(nodes: &[ProofTrieNodeV2]) -> usize { let mut capacity = nodes.len(); for node in nodes { - match &node.node { - TrieNode::Branch(branch) => { - capacity += branch.state_mask.count_ones() as usize; - } - TrieNode::Extension(_) => { - capacity += 1; - } - _ => {} - }; + if let TrieNodeV2::Branch(branch) = &node.node { + capacity += branch.state_mask.count_ones() as usize; + } } capacity @@ -1615,10 +1219,10 @@ fn estimate_v2_proof_capacity(nodes: &[ProofTrieNode]) -> usize { /// Filters V2 proof nodes that are already revealed, separates the root node if present, and /// returns additional information about the number of total, skipped, and new nodes. /// -/// Unlike [`filter_map_revealed_nodes`], V2 proof nodes already have masks included in the -/// `ProofTrieNode` structure, so no separate masks map is needed. +/// V2 proof nodes already have masks included in the `ProofTrieNode` structure, so no separate +/// masks map is needed. fn filter_revealed_v2_proof_nodes( - proof_nodes: Vec, + proof_nodes: Vec, revealed_nodes: &mut HashSet, ) -> SparseStateTrieResult { let mut result = FilteredV2ProofNodes { @@ -1632,7 +1236,7 @@ fn filter_revealed_v2_proof_nodes( // duplicate EmptyRoot nodes may appear (e.g., storage proofs split across chunks for an // account with empty storage). We only error if there's an EmptyRoot alongside real nodes. let non_empty_root_count = - proof_nodes.iter().filter(|n| !matches!(n.node, TrieNode::EmptyRoot)).count(); + proof_nodes.iter().filter(|n| !matches!(n.node, TrieNodeV2::EmptyRoot)).count(); for node in proof_nodes { result.metric_values.total_nodes += 1; @@ -1649,22 +1253,20 @@ fn filter_revealed_v2_proof_nodes( result.new_nodes += 1; // Count children for capacity estimation - match &node.node { - TrieNode::Branch(branch) => { - result.new_nodes += branch.state_mask.count_ones() as usize; - } - TrieNode::Extension(_) => { - result.new_nodes += 1; - } - _ => {} - }; + if let TrieNodeV2::Branch(branch) = &node.node { + result.new_nodes += branch.state_mask.count_ones() as usize; + } if is_root { // Perform sanity check: EmptyRoot is only valid if there are no other real nodes. - if matches!(node.node, TrieNode::EmptyRoot) && non_empty_root_count > 0 { + if matches!(node.node, TrieNodeV2::EmptyRoot) && non_empty_root_count > 0 { return Err(SparseStateTrieErrorKind::InvalidRootNode { path: node.path, - node: alloy_rlp::encode(&node.node).into(), + node: { + let mut buf = Vec::new(); + node.node.encode(&mut buf); + buf.into() + }, } .into()) } @@ -1694,7 +1296,8 @@ mod tests { use reth_trie::{updates::StorageTrieUpdates, HashBuilder, MultiProof, EMPTY_ROOT_HASH}; use reth_trie_common::{ proof::{ProofNodes, ProofRetainer}, - BranchNode, BranchNodeMasks, BranchNodeMasksMap, LeafNode, StorageMultiProof, TrieMask, + BranchNodeMasks, BranchNodeMasksMap, BranchNodeV2, LeafNode, RlpNode, StorageMultiProof, + TrieMask, }; #[test] @@ -1703,11 +1306,11 @@ mod tests { let mut sparse = SparseStateTrie::::default(); let leaf_value = alloy_rlp::encode(TrieAccount::default()); - let leaf_1 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new( + let leaf_1 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new( Nibbles::default(), leaf_value.clone(), ))); - let leaf_2 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new( + let leaf_2 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new( Nibbles::default(), leaf_value.clone(), ))); @@ -1716,9 +1319,11 @@ mod tests { account_subtree: ProofNodes::from_iter([ ( Nibbles::default(), - alloy_rlp::encode(TrieNode::Branch(BranchNode { + alloy_rlp::encode(TrieNodeV2::Branch(BranchNodeV2 { + key: Nibbles::default(), stack: vec![RlpNode::from_rlp(&leaf_1), RlpNode::from_rlp(&leaf_2)], state_mask: TrieMask::new(0b11), + branch_rlp_node: None, })) .into(), ), @@ -1772,11 +1377,11 @@ mod tests { let mut sparse = SparseStateTrie::::default(); let leaf_value = alloy_rlp::encode(TrieAccount::default()); - let leaf_1 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new( + let leaf_1 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new( Nibbles::default(), leaf_value.clone(), ))); - let leaf_2 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new( + let leaf_2 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new( Nibbles::default(), leaf_value.clone(), ))); @@ -1789,9 +1394,11 @@ mod tests { subtree: ProofNodes::from_iter([ ( Nibbles::default(), - alloy_rlp::encode(TrieNode::Branch(BranchNode { + alloy_rlp::encode(TrieNodeV2::Branch(BranchNodeV2 { + key: Nibbles::default(), stack: vec![RlpNode::from_rlp(&leaf_1), RlpNode::from_rlp(&leaf_2)], state_mask: TrieMask::new(0b11), + branch_rlp_node: None, })) .into(), ), @@ -1862,20 +1469,22 @@ mod tests { let mut sparse = SparseStateTrie::::default(); let leaf_value = alloy_rlp::encode(TrieAccount::default()); - let leaf_1_node = TrieNode::Leaf(LeafNode::new(Nibbles::default(), leaf_value.clone())); - let leaf_2_node = TrieNode::Leaf(LeafNode::new(Nibbles::default(), leaf_value.clone())); + let leaf_1_node = TrieNodeV2::Leaf(LeafNode::new(Nibbles::default(), leaf_value.clone())); + let leaf_2_node = TrieNodeV2::Leaf(LeafNode::new(Nibbles::default(), leaf_value.clone())); - let branch_node = TrieNode::Branch(BranchNode { + let branch_node = TrieNodeV2::Branch(BranchNodeV2 { + key: Nibbles::default(), stack: vec![ RlpNode::from_rlp(&alloy_rlp::encode(&leaf_1_node)), RlpNode::from_rlp(&alloy_rlp::encode(&leaf_2_node)), ], state_mask: TrieMask::new(0b11), + branch_rlp_node: None, }); // Create V2 proof nodes with masks already included let v2_proof_nodes = vec![ - ProofTrieNode { + ProofTrieNodeV2 { path: Nibbles::default(), node: branch_node, masks: Some(BranchNodeMasks { @@ -1883,8 +1492,8 @@ mod tests { tree_mask: TrieMask::default(), }), }, - ProofTrieNode { path: Nibbles::from_nibbles([0x0]), node: leaf_1_node, masks: None }, - ProofTrieNode { path: Nibbles::from_nibbles([0x1]), node: leaf_2_node, masks: None }, + ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x0]), node: leaf_1_node, masks: None }, + ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x1]), node: leaf_2_node, masks: None }, ]; // Reveal V2 proof nodes @@ -1923,21 +1532,25 @@ mod tests { let mut sparse = SparseStateTrie::::default(); let storage_value: Vec = alloy_rlp::encode_fixed_size(&U256::from(42)).to_vec(); - let leaf_1_node = TrieNode::Leaf(LeafNode::new(Nibbles::default(), storage_value.clone())); - let leaf_2_node = TrieNode::Leaf(LeafNode::new(Nibbles::default(), storage_value.clone())); + let leaf_1_node = + TrieNodeV2::Leaf(LeafNode::new(Nibbles::default(), storage_value.clone())); + let leaf_2_node = + TrieNodeV2::Leaf(LeafNode::new(Nibbles::default(), storage_value.clone())); - let branch_node = TrieNode::Branch(BranchNode { + let branch_node = TrieNodeV2::Branch(BranchNodeV2 { + key: Nibbles::default(), stack: vec![ RlpNode::from_rlp(&alloy_rlp::encode(&leaf_1_node)), RlpNode::from_rlp(&alloy_rlp::encode(&leaf_2_node)), ], state_mask: TrieMask::new(0b11), + branch_rlp_node: None, }); let v2_proof_nodes = vec![ - ProofTrieNode { path: Nibbles::default(), node: branch_node, masks: None }, - ProofTrieNode { path: Nibbles::from_nibbles([0x0]), node: leaf_1_node, masks: None }, - ProofTrieNode { path: Nibbles::from_nibbles([0x1]), node: leaf_2_node, masks: None }, + ProofTrieNodeV2 { path: Nibbles::default(), node: branch_node, masks: None }, + ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x0]), node: leaf_1_node, masks: None }, + ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x1]), node: leaf_2_node, masks: None }, ]; // Reveal V2 storage proof nodes for account @@ -2132,51 +1745,4 @@ mod tests { } ); } - - #[test] - fn test_filter_map_revealed_nodes() { - let mut revealed_nodes = HashSet::from_iter([Nibbles::from_nibbles([0x0])]); - let leaf = TrieNode::Leaf(LeafNode::new(Nibbles::default(), alloy_rlp::encode([]))); - let leaf_encoded = alloy_rlp::encode(&leaf); - let branch = TrieNode::Branch(BranchNode::new( - vec![RlpNode::from_rlp(&leaf_encoded), RlpNode::from_rlp(&leaf_encoded)], - TrieMask::new(0b11), - )); - let proof_nodes = alloy_trie::proof::DecodedProofNodes::from_iter([ - (Nibbles::default(), branch.clone()), - (Nibbles::from_nibbles([0x0]), leaf.clone()), - (Nibbles::from_nibbles([0x1]), leaf.clone()), - ]); - - let branch_node_masks = BranchNodeMasksMap::default(); - - let decoded = - filter_map_revealed_nodes(proof_nodes, &mut revealed_nodes, &branch_node_masks) - .unwrap(); - - assert_eq!( - decoded, - FilterMappedProofNodes { - root_node: Some(ProofTrieNode { - path: Nibbles::default(), - node: branch, - masks: None, - }), - nodes: vec![ProofTrieNode { - path: Nibbles::from_nibbles([0x1]), - node: leaf, - masks: None, - }], - // Branch, two of its children, one leaf - new_nodes: 4, - // Metric values - metric_values: ProofNodesMetricValues { - // Branch, leaf, leaf - total_nodes: 3, - // Revealed leaf node with path 0x1 - skipped_nodes: 1, - }, - } - ); - } } diff --git a/crates/trie/sparse/src/traits.rs b/crates/trie/sparse/src/traits.rs index 4b6844ce1b..87ea0f5ae7 100644 --- a/crates/trie/sparse/src/traits.rs +++ b/crates/trie/sparse/src/traits.rs @@ -9,7 +9,7 @@ use alloy_primitives::{ }; use alloy_trie::BranchNodeCompact; use reth_execution_errors::SparseTrieResult; -use reth_trie_common::{BranchNodeMasks, Nibbles, ProofTrieNode, TrieNode}; +use reth_trie_common::{BranchNodeMasks, Nibbles, ProofTrieNodeV2, TrieNodeV2}; use crate::provider::TrieNodeProvider; @@ -59,7 +59,7 @@ pub trait SparseTrie: Sized + Debug + Send + Sync { /// May panic if the trie is not new/cleared, and has already revealed nodes. fn set_root( &mut self, - root: TrieNode, + root: TrieNodeV2, masks: Option, retain_updates: bool, ) -> SparseTrieResult<()>; @@ -69,7 +69,7 @@ pub trait SparseTrie: Sized + Debug + Send + Sync { /// See [`Self::set_root`] for more details. fn with_root( mut self, - root: TrieNode, + root: TrieNodeV2, masks: Option, retain_updates: bool, ) -> SparseTrieResult { @@ -111,10 +111,10 @@ pub trait SparseTrie: Sized + Debug + Send + Sync { fn reveal_node( &mut self, path: Nibbles, - node: TrieNode, + node: TrieNodeV2, masks: Option, ) -> SparseTrieResult<()> { - self.reveal_nodes(&mut [ProofTrieNode { path, node, masks }]) + self.reveal_nodes(&mut [ProofTrieNodeV2 { path, node, masks }]) } /// Reveals one or more trie nodes if they have not been revealed before. @@ -135,8 +135,8 @@ pub trait SparseTrie: Sized + Debug + Send + Sync { /// # Note /// /// The implementation may modify the input nodes. A common thing to do is [`std::mem::replace`] - /// each node with [`TrieNode::EmptyRoot`] to avoid cloning. - fn reveal_nodes(&mut self, nodes: &mut [ProofTrieNode]) -> SparseTrieResult<()>; + /// 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. /// diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index ab0f9c9b73..c5f30fdb9a 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -5,7 +5,7 @@ use crate::{ use alloc::{boxed::Box, vec::Vec}; use alloy_primitives::{map::B256Map, B256}; use reth_execution_errors::{SparseTrieErrorKind, SparseTrieResult}; -use reth_trie_common::{BranchNodeMasks, Nibbles, RlpNode, TrieMask, TrieNode}; +use reth_trie_common::{BranchNodeMasks, Nibbles, RlpNode, TrieMask, TrieNode, TrieNodeV2}; use tracing::instrument; /// A sparse trie that is either in a "blind" state (no nodes are revealed, root node hash is @@ -63,7 +63,7 @@ impl RevealableSparseTrie { /// A mutable reference to the underlying [`RevealableSparseTrie`](SparseTrieTrait). pub fn reveal_root( &mut self, - root: TrieNode, + root: TrieNodeV2, masks: Option, retain_updates: bool, ) -> SparseTrieResult<&mut T> { diff --git a/crates/trie/trie/src/proof/mod.rs b/crates/trie/trie/src/proof/mod.rs index f9059b1c3c..2c0f6b8b2d 100644 --- a/crates/trie/trie/src/proof/mod.rs +++ b/crates/trie/trie/src/proof/mod.rs @@ -31,7 +31,7 @@ pub use trie_node::*; /// on the hash builder and follows the same algorithm as the state root calculator. /// See `StateRoot::root` for more info. #[derive(Debug)] -pub struct Proof { +pub struct Proof { /// The factory for traversing trie nodes. trie_cursor_factory: T, /// The factory for hashed cursors. @@ -40,6 +40,8 @@ pub struct Proof { prefix_sets: TriePrefixSetsMut, /// Flag indicating whether to include branch node masks in the proof. collect_branch_node_masks: bool, + /// Added and removed keys for proof retention. + added_removed_keys: Option, } impl Proof { @@ -50,26 +52,31 @@ impl Proof { hashed_cursor_factory: h, prefix_sets: TriePrefixSetsMut::default(), collect_branch_node_masks: false, + added_removed_keys: None, } } +} +impl Proof { /// Set the trie cursor factory. - pub fn with_trie_cursor_factory(self, trie_cursor_factory: TF) -> Proof { + pub fn with_trie_cursor_factory(self, trie_cursor_factory: TF) -> Proof { Proof { trie_cursor_factory, hashed_cursor_factory: self.hashed_cursor_factory, prefix_sets: self.prefix_sets, collect_branch_node_masks: self.collect_branch_node_masks, + added_removed_keys: self.added_removed_keys, } } /// Set the hashed cursor factory. - pub fn with_hashed_cursor_factory(self, hashed_cursor_factory: HF) -> Proof { + pub fn with_hashed_cursor_factory(self, hashed_cursor_factory: HF) -> Proof { Proof { trie_cursor_factory: self.trie_cursor_factory, hashed_cursor_factory, prefix_sets: self.prefix_sets, collect_branch_node_masks: self.collect_branch_node_masks, + added_removed_keys: self.added_removed_keys, } } @@ -85,6 +92,21 @@ impl Proof { self } + /// Configures the proof to retain certain nodes which would otherwise fall outside the target + /// set, when those nodes might be required to calculate the state root when keys have been + /// added or removed to the trie. + /// + /// If None is given then retention of extra proofs is disabled. + pub fn with_added_removed_keys(self, added_removed_keys: Option) -> Proof { + Proof { + trie_cursor_factory: self.trie_cursor_factory, + hashed_cursor_factory: self.hashed_cursor_factory, + prefix_sets: self.prefix_sets, + collect_branch_node_masks: self.collect_branch_node_masks, + added_removed_keys, + } + } + /// Get a reference to the trie cursor factory. pub const fn trie_cursor_factory(&self) -> &T { &self.trie_cursor_factory @@ -96,10 +118,11 @@ impl Proof { } } -impl Proof +impl Proof where T: TrieCursorFactory + Clone, H: HashedCursorFactory + Clone, + K: AsRef, { /// Generate an account proof from intermediate nodes. pub fn account_proof( @@ -126,10 +149,13 @@ where // Create the walker. let mut prefix_set = self.prefix_sets.account_prefix_set.clone(); prefix_set.extend_keys(targets.keys().map(Nibbles::unpack)); - let walker = TrieWalker::<_>::state_trie(trie_cursor, prefix_set.freeze()); + let walker = + TrieWalker::<_, AddedRemovedKeys>::state_trie(trie_cursor, prefix_set.freeze()) + .with_added_removed_keys(self.added_removed_keys.as_ref()); // Create a hash builder to rebuild the root node since it is not available in the database. - let retainer = targets.keys().map(Nibbles::unpack).collect(); + let retainer: ProofRetainer = targets.keys().map(Nibbles::unpack).collect(); + let retainer = retainer.with_added_removed_keys(self.added_removed_keys.as_ref()); let mut hash_builder = HashBuilder::default() .with_proof_retainer(retainer) .with_updates(self.collect_branch_node_masks); diff --git a/crates/trie/trie/src/proof_v2/mod.rs b/crates/trie/trie/src/proof_v2/mod.rs index f2f7755aaf..6e22306bd2 100644 --- a/crates/trie/trie/src/proof_v2/mod.rs +++ b/crates/trie/trie/src/proof_v2/mod.rs @@ -15,7 +15,9 @@ use alloy_primitives::{keccak256, B256, U256}; use alloy_rlp::Encodable; use alloy_trie::{BranchNodeCompact, TrieMask}; use reth_execution_errors::trie::StateProofError; -use reth_trie_common::{BranchNode, BranchNodeMasks, Nibbles, ProofTrieNode, RlpNode, TrieNode}; +use reth_trie_common::{ + BranchNodeMasks, BranchNodeRef, BranchNodeV2, Nibbles, ProofTrieNodeV2, RlpNode, TrieNodeV2, +}; use std::cmp::Ordering; use tracing::{error, instrument, trace}; @@ -86,7 +88,7 @@ pub struct ProofCalculator { cached_branch_stack: Vec<(Nibbles, BranchNodeCompact)>, /// The proofs which will be returned from the calculation. This gets taken at the end of every /// proof call. - retained_proofs: Vec, + retained_proofs: Vec, /// Free-list of re-usable buffers of [`RlpNode`]s, used for encoding branch nodes to RLP. /// /// We are generally able to re-use these buffers across different branch nodes for the @@ -193,7 +195,7 @@ where trace!(target: TRACE_TARGET, ?path, target = ?lower, "should_retain: called"); debug_assert!(self.retained_proofs.last().is_none_or( - |ProofTrieNode { path: last_retained_path, .. }| { + |ProofTrieNodeV2 { path: last_retained_path, .. }| { depth_first::cmp(path, last_retained_path) == Ordering::Greater } ), @@ -268,14 +270,14 @@ where if self.should_retain(targets, &child_path, true) { trace!(target: TRACE_TARGET, ?child_path, "Retaining child"); - // Convert to `ProofTrieNode`, which will be what is retained. + // Convert to `ProofTrieNodeV2`, which will be what is retained. // // If this node is a branch then its `rlp_nodes_buf` will be taken and not returned to // the `rlp_nodes_bufs` free-list. self.rlp_encode_buf.clear(); let proof_node = child.into_proof_trie_node(child_path, &mut self.rlp_encode_buf)?; - // Use the `ProofTrieNode` to encode the `RlpNode`, and then push it onto retained + // Use the `ProofTrieNodeV2` to encode the `RlpNode`, and then push it onto retained // nodes before returning. self.rlp_encode_buf.clear(); proof_node.node.encode(&mut self.rlp_encode_buf); @@ -540,17 +542,19 @@ where self.branch_path.len() - branch.ext_len as usize, ); - // Wrap the `BranchNode` so it can be pushed onto the child stack. - let mut branch_as_child = ProofTrieBranchChild::Branch { - node: BranchNode::new(rlp_nodes_buf, branch.state_mask), - masks: branch.masks, + // Compute hash for the branch node if it has a parent extension. + let rlp_node = if short_key.is_empty() { + None + } else { + self.rlp_encode_buf.clear(); + BranchNodeRef::new(&rlp_nodes_buf, branch.state_mask).encode(&mut self.rlp_encode_buf); + Some(RlpNode::from_rlp(&self.rlp_encode_buf)) }; - // If there is an extension then encode the branch as an `RlpNode` and use it to construct - // the extension in its place - if !short_key.is_empty() { - let branch_rlp_node = self.commit_child(targets, self.branch_path, branch_as_child)?; - branch_as_child = ProofTrieBranchChild::Extension { short_key, child: branch_rlp_node }; + // Wrap the `BranchNodeV2` so it can be pushed onto the child stack. + let branch_as_child = ProofTrieBranchChild::Branch { + node: BranchNodeV2::new(short_key, rlp_nodes_buf, branch.state_mask, rlp_node), + masks: branch.masks, }; self.child_stack.push(branch_as_child); @@ -1264,9 +1268,9 @@ where (true, None) => { // If `child_stack` is empty it means there was no keys at all, retain an empty // root node. - self.retained_proofs.push(ProofTrieNode { + self.retained_proofs.push(ProofTrieNodeV2 { path: Nibbles::new(), // root path - node: TrieNode::EmptyRoot, + node: TrieNodeV2::EmptyRoot, masks: None, }); } @@ -1288,7 +1292,7 @@ where &mut self, value_encoder: &mut VE, targets: &mut [Target], - ) -> Result, StateProofError> { + ) -> Result, StateProofError> { // If there are no targets then nothing could be returned, return early. if targets.is_empty() { trace!(target: TRACE_TARGET, "Empty targets, returning"); @@ -1332,7 +1336,7 @@ where &mut self, value_encoder: &mut VE, targets: &mut [Target], - ) -> Result, StateProofError> { + ) -> Result, StateProofError> { self.trie_cursor.reset(); self.hashed_cursor.reset(); self.proof_inner(value_encoder, targets) @@ -1346,7 +1350,7 @@ where /// This method reuses the internal RLP encode buffer for efficiency. pub fn compute_root_hash( &mut self, - proof_nodes: &[ProofTrieNode], + proof_nodes: &[ProofTrieNodeV2], ) -> Result, StateProofError> { // Find the root node (node at empty path) let root_node = proof_nodes.iter().find(|node| node.path.is_empty()); @@ -1368,7 +1372,10 @@ where /// This method does not accept targets nor retain proofs. Returns the root node which can /// be used to compute the root hash via [`Self::compute_root_hash`]. #[instrument(target = TRACE_TARGET, level = "trace", skip(self, value_encoder))] - pub fn root_node(&mut self, value_encoder: &mut VE) -> Result { + pub fn root_node( + &mut self, + value_encoder: &mut VE, + ) -> Result { // Initialize the variables which track the state of the two cursors. Both indicate the // cursors are unseeked. let mut trie_cursor_state = TrieCursorState::unseeked(); @@ -1434,15 +1441,15 @@ where &mut self, hashed_address: B256, targets: &mut [Target], - ) -> Result, StateProofError> { + ) -> Result, StateProofError> { self.hashed_cursor.set_hashed_address(hashed_address); // Shortcut: check if storage is empty if self.hashed_cursor.is_storage_empty()? { // Return a single EmptyRoot node at the root path - return Ok(vec![ProofTrieNode { + return Ok(vec![ProofTrieNodeV2 { path: Nibbles::default(), - node: TrieNode::EmptyRoot, + node: TrieNodeV2::EmptyRoot, masks: None, }]) } @@ -1464,13 +1471,13 @@ where pub fn storage_root_node( &mut self, hashed_address: B256, - ) -> Result { + ) -> Result { self.hashed_cursor.set_hashed_address(hashed_address); if self.hashed_cursor.is_storage_empty()? { - return Ok(ProofTrieNode { + return Ok(ProofTrieNodeV2 { path: Nibbles::default(), - node: TrieNode::EmptyRoot, + node: TrieNodeV2::EmptyRoot, masks: None, }) } @@ -1642,16 +1649,35 @@ mod tests { }; use alloy_primitives::map::{B256Map, B256Set}; use alloy_rlp::Decodable; + use alloy_trie::proof::AddedRemovedKeys; use itertools::Itertools; use reth_primitives_traits::Account; use reth_trie_common::{ updates::{StorageTrieUpdates, TrieUpdates}, - HashedPostState, MultiProofTargets, TrieNode, + HashedPostState, MultiProofTargets, ProofTrieNode, TrieNode, }; /// Target to use with the `tracing` crate. static TRACE_TARGET: &str = "trie::proof_v2::tests"; + /// Converts legacy proofs to V2 proofs by combining extension nodes with their child branch + /// nodes. + /// + /// In the legacy proof format, extension nodes and branch nodes are separate. In the V2 format, + /// they are combined into a single `BranchNodeV2` where the extension's key becomes the + /// branch's `key` field. + /// + /// Converts legacy proofs (sorted in depth-first order) to V2 format. + /// + /// In depth-first order, children come BEFORE parents. So when we encounter an extension node, + /// its child branch has already been processed and is in the result. We need to pop it and + /// combine it with the extension. + fn convert_legacy_proofs_to_v2(legacy_proofs: &[ProofTrieNode]) -> Vec { + ProofTrieNodeV2::from_sorted_trie_nodes( + legacy_proofs.iter().map(|p| (p.path, p.node.clone(), p.masks)), + ) + } + /// A test harness for comparing `ProofCalculator` and legacy `Proof` implementations. /// /// This harness creates mock cursor factories from a `HashedPostState` and provides @@ -1752,6 +1778,12 @@ mod tests { let proof_legacy_result = Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone()) .with_branch_node_masks(true) + .with_added_removed_keys(Some( + // This will force the HashBuilder to always retain the child branch of all + // extensions. We need this because in V2 extensions and branches are a + // single node type, so child branches are always included with extensions. + AddedRemovedKeys::default().with_assume_added(true), + )) .multiproof(legacy_targets)?; // Helper function to check if a node path matches at least one target @@ -1764,12 +1796,10 @@ mod tests { }) }; - // Decode and sort legacy proof nodes, filtering to only those that match at least one - // target + // Decode and sort legacy proof nodes let proof_legacy_nodes = proof_legacy_result .account_subtree .iter() - .filter(|(path, _)| node_matches_target(path)) .map(|(path, node_enc)| { let mut buf = node_enc.as_ref(); let node = TrieNode::decode(&mut buf) @@ -1788,8 +1818,18 @@ mod tests { .sorted_by(|a, b| depth_first::cmp(&a.path, &b.path)) .collect::>(); + // Convert legacy proofs to V2 proofs by combining extensions with their child branches + let proof_legacy_nodes_v2 = convert_legacy_proofs_to_v2(&proof_legacy_nodes); + + // Filter to only keep nodes which match a target. We do this after conversion so we + // don't keep branches whose extension parents are excluded due to a min_len. + let proof_legacy_nodes_v2 = proof_legacy_nodes_v2 + .into_iter() + .filter(|ProofTrieNodeV2 { path, .. }| node_matches_target(path)) + .collect::>(); + // Basic comparison: both should succeed and produce identical results - pretty_assertions::assert_eq!(proof_legacy_nodes, proof_v2_result); + pretty_assertions::assert_eq!(proof_legacy_nodes_v2, proof_v2_result); // Also test root_node - get a fresh calculator and verify it returns the root node // that hashes to the expected root @@ -1876,7 +1916,7 @@ mod tests { } proptest! { - #![proptest_config(ProptestConfig::with_cases(8000))] + #![proptest_config(ProptestConfig::with_cases(4000))] #[test] /// Tests that ProofCalculator produces valid proofs for randomly generated /// HashedPostState with proof targets. diff --git a/crates/trie/trie/src/proof_v2/node.rs b/crates/trie/trie/src/proof_v2/node.rs index 6acdaae881..5d8cc4185a 100644 --- a/crates/trie/trie/src/proof_v2/node.rs +++ b/crates/trie/trie/src/proof_v2/node.rs @@ -1,10 +1,9 @@ use crate::proof_v2::DeferredValueEncoder; use alloy_rlp::Encodable; -use alloy_trie::nodes::ExtensionNodeRef; use reth_execution_errors::trie::StateProofError; use reth_trie_common::{ - BranchNode, BranchNodeMasks, ExtensionNode, LeafNode, LeafNodeRef, Nibbles, ProofTrieNode, - RlpNode, TrieMask, TrieNode, + BranchNodeMasks, BranchNodeV2, LeafNode, LeafNodeRef, Nibbles, ProofTrieNodeV2, RlpNode, + TrieMask, TrieNodeV2, }; /// A trie node which is the child of a branch in the trie. @@ -17,17 +16,10 @@ pub(crate) enum ProofTrieBranchChild { /// The [`DeferredValueEncoder`] which will encode the leaf's value. value: RF, }, - /// An extension node whose child branch has been converted to an [`RlpNode`] - Extension { - /// The short key of the leaf. - short_key: Nibbles, - /// The [`RlpNode`] of the child branch. - child: RlpNode, - }, /// A branch node whose children have already been flattened into [`RlpNode`]s. Branch { /// The node itself, for use during RLP encoding. - node: BranchNode, + node: BranchNodeV2, /// Bitmasks carried over from cached `BranchNodeCompact` values, if any. masks: Option, }, @@ -36,8 +28,10 @@ pub(crate) enum ProofTrieBranchChild { } impl ProofTrieBranchChild { - /// Converts this child into its RLP node representation. This potentially also returns an - /// `RlpNode` buffer which can be re-used for other [`ProofTrieBranchChild`]s. + /// Converts this child into its RLP node representation. + /// + /// This potentially also returns an `RlpNode` buffer which can be re-used for other + /// [`ProofTrieBranchChild`]s. pub(crate) fn into_rlp( self, buf: &mut Vec, @@ -65,10 +59,6 @@ impl ProofTrieBranchChild { LeafNodeRef::new(&short_key, value_buf).encode(&mut leaf_buf); Ok((RlpNode::from_rlp(&buf[value_enc_len..]), None)) } - Self::Extension { short_key, child } => { - ExtensionNodeRef::new(&short_key, child.as_slice()).encode(buf); - Ok((RlpNode::from_rlp(buf), None)) - } Self::Branch { node: branch_node, .. } => { branch_node.encode(buf); Ok((RlpNode::from_rlp(buf), Some(branch_node.stack))) @@ -77,7 +67,7 @@ impl ProofTrieBranchChild { } } - /// Converts this child into a [`ProofTrieNode`] having the given path. + /// Converts this child into a [`ProofTrieNodeV2`] having the given path. /// /// # Panics /// @@ -86,7 +76,7 @@ impl ProofTrieBranchChild { self, path: Nibbles, buf: &mut Vec, - ) -> Result { + ) -> Result { let (node, masks) = match self { Self::Leaf { short_key, value } => { value.encode(buf)?; @@ -98,24 +88,22 @@ impl ProofTrieBranchChild { // this value, and the passed in buffer can remain with whatever large capacity it // already has. let rlp_val = buf.clone(); - (TrieNode::Leaf(LeafNode::new(short_key, rlp_val)), None) + (TrieNodeV2::Leaf(LeafNode::new(short_key, rlp_val)), None) } - Self::Extension { short_key, child } => { - (TrieNode::Extension(ExtensionNode { key: short_key, child }), None) - } - Self::Branch { node, masks } => (TrieNode::Branch(node), masks), + Self::Branch { node, masks } => (TrieNodeV2::Branch(node), masks), Self::RlpNode(_) => panic!("Cannot call `into_proof_trie_node` on RlpNode"), }; - Ok(ProofTrieNode { node, path, masks }) + Ok(ProofTrieNodeV2 { node, path, masks }) } - /// Returns the short key of the child, if it is a leaf or extension, or empty if its a - /// [`Self::Branch`] or [`Self::RlpNode`]. + /// Returns the short key of the child, if it is a leaf or branch, or empty if its a + /// [`Self::RlpNode`]. pub(crate) fn short_key(&self) -> &Nibbles { match self { - Self::Leaf { short_key, .. } | Self::Extension { short_key, .. } => short_key, - Self::Branch { .. } | Self::RlpNode(_) => { + Self::Leaf { short_key, .. } | + Self::Branch { node: BranchNodeV2 { key: short_key, .. }, .. } => short_key, + Self::RlpNode(_) => { static EMPTY_NIBBLES: Nibbles = Nibbles::new(); &EMPTY_NIBBLES } @@ -134,14 +122,17 @@ impl ProofTrieBranchChild { /// - If the node is a [`Self::Branch`] or [`Self::RlpNode`] pub(crate) fn trim_short_key_prefix(&mut self, len: usize) { match self { - Self::Extension { short_key, child } if short_key.len() == len => { - *self = Self::RlpNode(core::mem::take(child)); - } - Self::Leaf { short_key, .. } | Self::Extension { short_key, .. } => { + Self::Leaf { short_key, .. } => { *short_key = trim_nibbles_prefix(short_key, len); } - Self::Branch { .. } | Self::RlpNode(_) => { - panic!("Cannot call `trim_short_key_prefix` on Branch or RlpNode") + Self::Branch { node: BranchNodeV2 { key, branch_rlp_node, .. }, .. } => { + *key = trim_nibbles_prefix(key, len); + if key.is_empty() { + *branch_rlp_node = None; + } + } + Self::RlpNode(_) => { + panic!("Cannot call `trim_short_key_prefix` on RlpNode") } } } diff --git a/crates/trie/trie/src/trie_cursor/depth_first.rs b/crates/trie/trie/src/trie_cursor/depth_first.rs index b9cef85e3f..5803077086 100644 --- a/crates/trie/trie/src/trie_cursor/depth_first.rs +++ b/crates/trie/trie/src/trie_cursor/depth_first.rs @@ -6,36 +6,9 @@ use tracing::trace; /// Compares two Nibbles in depth-first order. /// -/// In depth-first ordering: -/// - Descendants come before their ancestors (children before parents) -/// - Siblings are ordered lexicographically -/// -/// # Example -/// -/// ```text -/// 0x11 comes before 0x1 (child before parent) -/// 0x12 comes before 0x1 (child before parent) -/// 0x11 comes before 0x12 (lexicographical among siblings) -/// 0x1 comes before 0x21 (lexicographical among siblings) -/// Result: 0x11, 0x12, 0x1, 0x21 -/// ``` +/// See [`reth_trie_common::depth_first_cmp`] for details. pub fn cmp(a: &Nibbles, b: &Nibbles) -> Ordering { - // If the two are of equal length, then compare them lexicographically - if a.len() == b.len() { - return a.cmp(b) - } - - // If one is a prefix of the other, then the other comes first - let common_prefix_len = a.common_prefix_length(b); - if a.len() == common_prefix_len { - return Ordering::Greater - } else if b.len() == common_prefix_len { - return Ordering::Less - } - - // Otherwise the nibble after the prefix determines the ordering. We know that neither is empty - // at this point, otherwise the previous if/else block would have caught it. - a.get_unchecked(common_prefix_len).cmp(&b.get_unchecked(common_prefix_len)) + reth_trie_common::depth_first_cmp(a, b) } /// An iterator that traverses trie nodes in depth-first post-order.