fix(witness): collect witness using sparse trie (#13072)

crates/evm/execution-errors/src/trie.rs

@@ -118,7 +118,7 @@ pub enum SparseTrieError {
 }
 
 /// Trie witness errors.
-#[derive(Error, PartialEq, Eq, Clone, Debug)]
+#[derive(Error, Debug)]
 pub enum TrieWitnessError {
     /// Error gather proofs.
     #[error(transparent)]
@@ -126,15 +126,12 @@ pub enum TrieWitnessError {
     /// RLP decoding error.
     #[error(transparent)]
     Rlp(#[from] alloy_rlp::Error),
+    /// Sparse state trie error.
+    #[error(transparent)]
+    Sparse(#[from] SparseStateTrieError),
     /// Missing account.
     #[error("missing account {_0}")]
     MissingAccount(B256),
-    /// Missing target node.
-    #[error("target node missing from proof {_0:?}")]
-    MissingTargetNode(Nibbles),
-    /// Unexpected empty root.
-    #[error("unexpected empty root: {_0:?}")]
-    UnexpectedEmptyRoot(Nibbles),
 }
 
 impl From<TrieWitnessError> for ProviderError {

crates/trie/sparse/src/state.rs

@@ -36,10 +36,10 @@ pub struct SparseStateTrie<F: BlindedProviderFactory = DefaultBlindedProviderFac
 impl Default for SparseStateTrie {
     fn default() -> Self {
         Self {
+            provider_factory: Default::default(),
             state: Default::default(),
             storages: Default::default(),
             revealed: Default::default(),
-            provider_factory: Default::default(),
             retain_updates: false,
             account_rlp_buf: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE),
         }

crates/trie/trie/src/proof/blinded.rs

@@ -20,6 +20,17 @@ pub struct ProofBlindedProviderFactory<T, H> {
     prefix_sets: Arc<TriePrefixSetsMut>,
 }
 
+impl<T, H> ProofBlindedProviderFactory<T, H> {
+    /// Create new proof-based blinded provider factory.
+    pub const fn new(
+        trie_cursor_factory: T,
+        hashed_cursor_factory: H,
+        prefix_sets: Arc<TriePrefixSetsMut>,
+    ) -> Self {
+        Self { trie_cursor_factory, hashed_cursor_factory, prefix_sets }
+    }
+}
+
 impl<T, H> BlindedProviderFactory for ProofBlindedProviderFactory<T, H>
 where
     T: TrieCursorFactory + Clone,
@@ -57,6 +68,17 @@ pub struct ProofBlindedAccountProvider<T, H> {
     prefix_sets: Arc<TriePrefixSetsMut>,
 }
 
+impl<T, H> ProofBlindedAccountProvider<T, H> {
+    /// Create new proof-based blinded account node provider.
+    pub const fn new(
+        trie_cursor_factory: T,
+        hashed_cursor_factory: H,
+        prefix_sets: Arc<TriePrefixSetsMut>,
+    ) -> Self {
+        Self { trie_cursor_factory, hashed_cursor_factory, prefix_sets }
+    }
+}
+
 impl<T, H> BlindedProvider for ProofBlindedAccountProvider<T, H>
 where
     T: TrieCursorFactory + Clone,
@@ -89,6 +111,18 @@ pub struct ProofBlindedStorageProvider<T, H> {
     account: B256,
 }
 
+impl<T, H> ProofBlindedStorageProvider<T, H> {
+    /// Create new proof-based blinded storage node provider.
+    pub const fn new(
+        trie_cursor_factory: T,
+        hashed_cursor_factory: H,
+        prefix_sets: Arc<TriePrefixSetsMut>,
+        account: B256,
+    ) -> Self {
+        Self { trie_cursor_factory, hashed_cursor_factory, prefix_sets, account }
+    }
+}
+
 impl<T, H> BlindedProvider for ProofBlindedStorageProvider<T, H>
 where
     T: TrieCursorFactory + Clone,

crates/trie/trie/src/witness.rs

@@ -1,23 +1,25 @@
 use crate::{
     hashed_cursor::{HashedCursor, HashedCursorFactory},
     prefix_set::TriePrefixSetsMut,
-    proof::{Proof, StorageProof},
+    proof::{Proof, ProofBlindedProviderFactory},
     trie_cursor::TrieCursorFactory,
-    HashedPostState, TRIE_ACCOUNT_RLP_MAX_SIZE,
+    HashedPostState,
 };
-use alloy_consensus::EMPTY_ROOT_HASH;
 use alloy_primitives::{
     keccak256,
-    map::{HashMap, HashSet},
+    map::{Entry, HashMap, HashSet},
     Bytes, B256,
 };
-use alloy_rlp::{BufMut, Decodable, Encodable};
+use itertools::Itertools;
-use itertools::{Either, Itertools};
+use reth_execution_errors::{
-use reth_execution_errors::{StateProofError, TrieWitnessError};
+    SparseStateTrieError, SparseTrieError, StateProofError, TrieWitnessError,
-use reth_trie_common::{
-    BranchNode, HashBuilder, Nibbles, StorageMultiProof, TrieAccount, TrieNode, CHILD_INDEX_RANGE,
 };
-use std::collections::BTreeMap;
+use reth_trie_common::Nibbles;
+use reth_trie_sparse::{
+    blinded::{BlindedProvider, BlindedProviderFactory},
+    SparseStateTrie,
+};
+use std::sync::{mpsc, Arc};
 
 /// State transition witness for the trie.
 #[derive(Debug)]
@@ -90,108 +92,75 @@ where
         }
 
         let proof_targets = self.get_proof_targets(&state)?;
-        let mut account_multiproof =
+        let multiproof =
             Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone())
                 .with_prefix_sets_mut(self.prefix_sets.clone())
                 .multiproof(proof_targets.clone())?;
 
-        // Attempt to compute state root from proofs and gather additional
+        // Record all nodes from multiproof in the witness
-        // information for the witness.
+        for account_node in multiproof.account_subtree.values() {
-        let mut account_rlp = Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE);
+            if let Entry::Vacant(entry) = self.witness.entry(keccak256(account_node.as_ref())) {
-        let mut account_trie_nodes = BTreeMap::default();
+                entry.insert(account_node.clone());
-        for (hashed_address, hashed_slots) in proof_targets {
+            }
-            let storage_multiproof = account_multiproof
+        }
-                .storages
+        for storage_node in multiproof.storages.values().flat_map(|s| s.subtree.values()) {
-                .remove(&hashed_address)
+            if let Entry::Vacant(entry) = self.witness.entry(keccak256(storage_node.as_ref())) {
-                .unwrap_or_else(StorageMultiProof::empty);
+                entry.insert(storage_node.clone());
+            }
+        }
 
-            // Gather and record account trie nodes.
+        let (tx, rx) = mpsc::channel();
-            let account = state
+        let proof_provider_factory = ProofBlindedProviderFactory::new(
-                .accounts
+            self.trie_cursor_factory,
-                .get(&hashed_address)
+            self.hashed_cursor_factory,
-                .ok_or(TrieWitnessError::MissingAccount(hashed_address))?;
+            Arc::new(self.prefix_sets),
-            let value =
+        );
-                (account.is_some() || storage_multiproof.root != EMPTY_ROOT_HASH).then(|| {
+        let mut sparse_trie =
-                    account_rlp.clear();
+            SparseStateTrie::new(WitnessBlindedProviderFactory::new(proof_provider_factory, tx));
-                    TrieAccount::from((account.unwrap_or_default(), storage_multiproof.root))
+        sparse_trie.reveal_multiproof(proof_targets.clone(), multiproof)?;
-                        .encode(&mut account_rlp as &mut dyn BufMut);
-                    account_rlp.clone()
-                });
-            let key = Nibbles::unpack(hashed_address);
-            account_trie_nodes.extend(target_nodes(
-                key.clone(),
-                value,
-                Some(&mut self.witness),
-                account_multiproof
-                    .account_subtree
-                    .matching_nodes_iter(&key)
-                    .sorted_by(|a, b| a.0.cmp(b.0)),
-            )?);
 
-            // Gather and record storage trie nodes for this account.
+        // Attempt to update state trie to gather additional information for the witness.
-            let mut storage_trie_nodes = BTreeMap::default();
+        for (hashed_address, hashed_slots) in
+            proof_targets.into_iter().sorted_unstable_by_key(|(ha, _)| *ha)
+        {
+            // Update storage trie first.
             let storage = state.storages.get(&hashed_address);
-            for hashed_slot in hashed_slots {
+            let storage_trie = sparse_trie
-                let slot_nibbles = Nibbles::unpack(hashed_slot);
+                .storage_trie_mut(&hashed_address)
-                let slot_value = storage
+                .ok_or(SparseStateTrieError::Sparse(SparseTrieError::Blind))?;
+            for hashed_slot in hashed_slots.into_iter().sorted_unstable() {
+                let storage_nibbles = Nibbles::unpack(hashed_slot);
+                let maybe_leaf_value = storage
                     .and_then(|s| s.storage.get(&hashed_slot))
                     .filter(|v| !v.is_zero())
                     .map(|v| alloy_rlp::encode_fixed_size(v).to_vec());
-                storage_trie_nodes.extend(target_nodes(
+
-                    slot_nibbles.clone(),
+                if let Some(value) = maybe_leaf_value {
-                    slot_value,
+                    storage_trie
-                    Some(&mut self.witness),
+                        .update_leaf(storage_nibbles, value)
-                    storage_multiproof
+                        .map_err(SparseStateTrieError::Sparse)?;
-                        .subtree
+                } else {
-                        .matching_nodes_iter(&slot_nibbles)
+                    storage_trie
-                        .sorted_by(|a, b| a.0.cmp(b.0)),
+                        .remove_leaf(&storage_nibbles)
-                )?);
+                        .map_err(SparseStateTrieError::Sparse)?;
+                }
             }
 
-            next_root_from_proofs(storage_trie_nodes, |key: Nibbles| {
+            // Calculate storage root after updates.
-                // Right pad the target with 0s.
+            storage_trie.root();
-                let mut padded_key = key.pack();
-                padded_key.resize(32, 0);
-                let target_key = B256::from_slice(&padded_key);
-                let storage_prefix_set = self
-                    .prefix_sets
-                    .storage_prefix_sets
-                    .get(&hashed_address)
-                    .cloned()
-                    .unwrap_or_default();
-                let proof = StorageProof::new_hashed(
-                    self.trie_cursor_factory.clone(),
-                    self.hashed_cursor_factory.clone(),
-                    hashed_address,
-                )
-                .with_prefix_set_mut(storage_prefix_set)
-                .storage_multiproof(HashSet::from_iter([target_key]))?;
 
-                // The subtree only contains the proof for a single target.
+            let account = state
-                let node =
+                .accounts
-                    proof.subtree.get(&key).ok_or(TrieWitnessError::MissingTargetNode(key))?;
+                .get(&hashed_address)
-                self.witness.insert(keccak256(node.as_ref()), node.clone()); // record in witness
+                .ok_or(TrieWitnessError::MissingAccount(hashed_address))?
-                Ok(node.clone())
+                .unwrap_or_default();
-            })?;
+            sparse_trie.update_account(hashed_address, account)?;
+
+            while let Ok(node) = rx.try_recv() {
+                self.witness.insert(keccak256(&node), node);
+            }
         }
 
-        next_root_from_proofs(account_trie_nodes, |key: Nibbles| {
-            // Right pad the target with 0s.
-            let mut padded_key = key.pack();
-            padded_key.resize(32, 0);
-            let targets = HashMap::from_iter([(B256::from_slice(&padded_key), HashSet::default())]);
-            let proof =
-                Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone())
-                    .with_prefix_sets_mut(self.prefix_sets.clone())
-                    .multiproof(targets)?;
-
-            // The subtree only contains the proof for a single target.
-            let node =
-                proof.account_subtree.get(&key).ok_or(TrieWitnessError::MissingTargetNode(key))?;
-            self.witness.insert(keccak256(node.as_ref()), node.clone()); // record in witness
-            Ok(node.clone())
-        })?;
-
         Ok(self.witness)
     }
 
@@ -225,141 +194,65 @@ where
     }
 }
 
-/// Decodes and unrolls all nodes from the proof. Returns only sibling nodes
+#[derive(Debug)]
-/// in the path of the target and the final leaf node with updated value.
+struct WitnessBlindedProviderFactory<F> {
-pub fn target_nodes<'b>(
+    /// Blinded node provider factory.
-    key: Nibbles,
+    provider_factory: F,
-    value: Option<Vec<u8>>,
+    /// Sender for forwarding fetched blinded node.
-    mut witness: Option<&mut HashMap<B256, Bytes>>,
+    tx: mpsc::Sender<Bytes>,
-    proof: impl IntoIterator<Item = (&'b Nibbles, &'b Bytes)>,
-) -> Result<BTreeMap<Nibbles, Either<B256, Vec<u8>>>, TrieWitnessError> {
-    let mut trie_nodes = BTreeMap::default();
-    let mut proof_iter = proof.into_iter().enumerate().peekable();
-    while let Some((idx, (path, encoded))) = proof_iter.next() {
-        // Record the node in witness.
-        if let Some(witness) = witness.as_mut() {
-            witness.insert(keccak256(encoded.as_ref()), encoded.clone());
-        }
-
-        let mut next_path = path.clone();
-        match TrieNode::decode(&mut &encoded[..])? {
-            TrieNode::Branch(branch) => {
-                next_path.push(key[path.len()]);
-                let children = branch_node_children(path.clone(), &branch);
-                for (child_path, value) in children {
-                    if !key.starts_with(&child_path) {
-                        let value = if value.len() < B256::len_bytes() {
-                            Either::Right(value.to_vec())
-                        } else {
-                            Either::Left(B256::from_slice(&value[1..]))
-                        };
-                        trie_nodes.insert(child_path, value);
-                    }
-                }
-            }
-            TrieNode::Extension(extension) => {
-                next_path.extend_from_slice(&extension.key);
-            }
-            TrieNode::Leaf(leaf) => {
-                next_path.extend_from_slice(&leaf.key);
-                if next_path != key {
-                    trie_nodes
-                        .insert(next_path.clone(), Either::Right(leaf.value.as_slice().to_vec()));
-                }
-            }
-            TrieNode::EmptyRoot => {
-                if idx != 0 || proof_iter.peek().is_some() {
-                    return Err(TrieWitnessError::UnexpectedEmptyRoot(next_path))
-                }
-            }
-        };
-    }
-
-    if let Some(value) = value {
-        trie_nodes.insert(key, Either::Right(value));
-    }
-
-    Ok(trie_nodes)
 }
 
-/// Computes the next root hash of a trie by processing a set of trie nodes and
+impl<F> WitnessBlindedProviderFactory<F> {
-/// their provided values.
+    const fn new(provider_factory: F, tx: mpsc::Sender<Bytes>) -> Self {
-pub fn next_root_from_proofs(
+        Self { provider_factory, tx }
-    trie_nodes: BTreeMap<Nibbles, Either<B256, Vec<u8>>>,
-    mut trie_node_provider: impl FnMut(Nibbles) -> Result<Bytes, TrieWitnessError>,
-) -> Result<B256, TrieWitnessError> {
-    // Ignore branch child hashes in the path of leaves or lower child hashes.
-    let mut keys = trie_nodes.keys().peekable();
-    let mut ignored = HashSet::<Nibbles>::default();
-    while let Some(key) = keys.next() {
-        if keys.peek().is_some_and(|next| next.starts_with(key)) {
-            ignored.insert(key.clone());
-        }
     }
-
-    let mut hash_builder = HashBuilder::default();
-    let mut trie_nodes = trie_nodes.into_iter().filter(|e| !ignored.contains(&e.0)).peekable();
-    while let Some((path, value)) = trie_nodes.next() {
-        match value {
-            Either::Left(branch_hash) => {
-                let parent_branch_path = path.slice(..path.len() - 1);
-                if hash_builder.key.starts_with(&parent_branch_path) ||
-                    trie_nodes.peek().is_some_and(|next| next.0.starts_with(&parent_branch_path))
-                {
-                    hash_builder.add_branch(path, branch_hash, false);
-                } else {
-                    // Parent is a branch node that needs to be turned into an extension node.
-                    let mut path = path.clone();
-                    loop {
-                        let node = trie_node_provider(path.clone())?;
-                        match TrieNode::decode(&mut &node[..])? {
-                            TrieNode::Branch(branch) => {
-                                let children = branch_node_children(path, &branch);
-                                for (child_path, value) in children {
-                                    if value.len() < B256::len_bytes() {
-                                        hash_builder.add_leaf(child_path, value);
-                                    } else {
-                                        let hash = B256::from_slice(&value[1..]);
-                                        hash_builder.add_branch(child_path, hash, false);
-                                    }
-                                }
-                                break
-                            }
-                            TrieNode::Leaf(leaf) => {
-                                let mut child_path = path;
-                                child_path.extend_from_slice(&leaf.key);
-                                hash_builder.add_leaf(child_path, &leaf.value);
-                                break
-                            }
-                            TrieNode::Extension(ext) => {
-                                path.extend_from_slice(&ext.key);
-                            }
-                            TrieNode::EmptyRoot => {
-                                return Err(TrieWitnessError::UnexpectedEmptyRoot(path))
-                            }
-                        }
-                    }
-                }
-            }
-            Either::Right(leaf_value) => {
-                hash_builder.add_leaf(path, &leaf_value);
-            }
-        }
-    }
-    Ok(hash_builder.root())
 }
 
-/// Returned branch node children with keys in order.
+impl<F> BlindedProviderFactory for WitnessBlindedProviderFactory<F>
-fn branch_node_children(prefix: Nibbles, node: &BranchNode) -> Vec<(Nibbles, &[u8])> {
+where
-    let mut children = Vec::with_capacity(node.state_mask.count_ones() as usize);
+    F: BlindedProviderFactory,
-    let mut stack_ptr = node.as_ref().first_child_index();
+    F::AccountNodeProvider: BlindedProvider<Error = SparseTrieError>,
-    for index in CHILD_INDEX_RANGE {
+    F::StorageNodeProvider: BlindedProvider<Error = SparseTrieError>,
-        if node.state_mask.is_bit_set(index) {
+{
-            let mut child_path = prefix.clone();
+    type AccountNodeProvider = WitnessBlindedProvider<F::AccountNodeProvider>;
-            child_path.push(index);
+    type StorageNodeProvider = WitnessBlindedProvider<F::StorageNodeProvider>;
-            children.push((child_path, &node.stack[stack_ptr][..]));
+
-            stack_ptr += 1;
+    fn account_node_provider(&self) -> Self::AccountNodeProvider {
-        }
+        let provider = self.provider_factory.account_node_provider();
+        WitnessBlindedProvider::new(provider, self.tx.clone())
+    }
+
+    fn storage_node_provider(&self, account: B256) -> Self::StorageNodeProvider {
+        let provider = self.provider_factory.storage_node_provider(account);
+        WitnessBlindedProvider::new(provider, self.tx.clone())
+    }
+}
+
+#[derive(Debug)]
+struct WitnessBlindedProvider<P> {
+    /// Proof-based blinded.
+    provider: P,
+    /// Sender for forwarding fetched blinded node.
+    tx: mpsc::Sender<Bytes>,
+}
+
+impl<P> WitnessBlindedProvider<P> {
+    const fn new(provider: P, tx: mpsc::Sender<Bytes>) -> Self {
+        Self { provider, tx }
+    }
+}
+
+impl<P> BlindedProvider for WitnessBlindedProvider<P>
+where
+    P: BlindedProvider<Error = SparseTrieError>,
+{
+    type Error = P::Error;
+
+    fn blinded_node(&mut self, path: Nibbles) -> Result<Option<Bytes>, Self::Error> {
+        let maybe_node = self.provider.blinded_node(path)?;
+        if let Some(node) = &maybe_node {
+            self.tx.send(node.clone()).map_err(|error| SparseTrieError::Other(Box::new(error)))?;
+        }
+        Ok(maybe_node)
     }
-    children
 }