feat(engine, trie): fetch blinded nodes when removing a sparse trie leaf

crates/engine/tree/src/tree/root.rs

@@ -5,13 +5,16 @@ use reth_provider::{
     providers::ConsistentDbView, BlockReader, DBProvider, DatabaseProviderFactory,
 };
 use reth_trie::{
-    proof::Proof, updates::TrieUpdates, HashedPostState, HashedStorage, MultiProof, Nibbles,
-    TrieInput,
+    hashed_cursor::HashedPostStateCursorFactory,
+    proof::{Proof, StorageProof},
+    trie_cursor::InMemoryTrieCursorFactory,
+    updates::{TrieUpdates, TrieUpdatesSorted},
+    HashedPostState, HashedPostStateSorted, HashedStorage, MultiProof, Nibbles, TrieInput,
 };
-use reth_trie_db::DatabaseProof;
+use reth_trie_db::{DatabaseHashedCursorFactory, DatabaseProof, DatabaseTrieCursorFactory};
 use reth_trie_parallel::root::ParallelStateRootError;
 use reth_trie_sparse::{SparseStateTrie, SparseStateTrieResult};
-use revm_primitives::{keccak256, EvmState, B256};
+use revm_primitives::{keccak256, map::DefaultHashBuilder, EvmState, B256};
 use std::{
     collections::BTreeMap,
     sync::{
@@ -317,10 +320,26 @@ where
 
         // TODO(alexey): store proof targets in `ProofSequecner` to avoid recomputing them
         let targets = get_proof_targets(&state, &HashSet::default());
+        let provider_ro = self.config.consistent_view.provider_ro();
+        let nodes_sorted = self.config.input.nodes.clone().into_sorted();
+        let state_sorted = self.config.input.state.clone().into_sorted();
 
         let tx = self.tx.clone();
         rayon::spawn(move || {
-            let result = update_sparse_trie(trie, multiproof, targets, state);
+            let Ok(provider_ro) = provider_ro else {
+                error!(target: "engine::root", "Failed to get read-only provider for sparse trie update");
+                return;
+            };
+
+            let result = update_sparse_trie::<Factory>(
+                trie,
+                multiproof,
+                targets,
+                state,
+                provider_ro,
+                nodes_sorted,
+                state_sorted,
+            );
             match result {
                 Ok((trie, elapsed)) => {
                     trace!(
@@ -476,11 +495,14 @@ fn get_proof_targets(
 
 /// Updates the sparse trie with the given proofs and state, and returns the updated trie and the
 /// time it took.
-fn update_sparse_trie(
+fn update_sparse_trie<Factory: DatabaseProviderFactory>(
     mut trie: Box<SparseStateTrie>,
     multiproof: MultiProof,
     targets: HashMap<B256, HashSet<B256>>,
     state: HashedPostState,
+    provider: Factory::Provider,
+    input_nodes_sorted: TrieUpdatesSorted,
+    input_state_sorted: HashedPostStateSorted,
 ) -> SparseStateTrieResult<(Box<SparseStateTrie>, Duration)> {
     let started_at = Instant::now();
 
@@ -496,8 +518,29 @@ fn update_sparse_trie(
         for (slot, value) in storage.storage {
             let slot_nibbles = Nibbles::unpack(slot);
             if value.is_zero() {
-                // TODO: handle blinded node error
-                trie.remove_storage_leaf(address, &slot_nibbles)?;
+                trie.remove_storage_leaf(address, &slot_nibbles, |path| {
+                    // Right pad the target with 0s.
+                    let mut padded_key = path.pack();
+                    padded_key.resize(32, 0);
+                    let mut targets = HashSet::with_hasher(DefaultHashBuilder::default());
+                    targets.insert(B256::from_slice(&padded_key));
+                    let proof = StorageProof::new_hashed(
+                        InMemoryTrieCursorFactory::new(
+                            DatabaseTrieCursorFactory::new(provider.tx_ref()),
+                            &input_nodes_sorted,
+                        ),
+                        HashedPostStateCursorFactory::new(
+                            DatabaseHashedCursorFactory::new(provider.tx_ref()),
+                            &input_state_sorted,
+                        ),
+                        address,
+                    )
+                    .storage_multiproof(targets)
+                    .unwrap();
+
+                    // The subtree only contains the proof for a single target.
+                    proof.subtree.get(&path).cloned()
+                })?;
             } else {
                 trie.update_storage_leaf(
                     address,
@@ -512,7 +555,31 @@ fn update_sparse_trie(
 
     // Update accounts with new values
     for (address, account) in state.accounts {
-        trie.update_account(address, account.unwrap_or_default())?;
+        trie.update_account(address, account.unwrap_or_default(), |path| {
+            // Right pad the target with 0s.
+            let mut padded_key = path.pack();
+            padded_key.resize(32, 0);
+            let mut targets = HashMap::with_hasher(DefaultHashBuilder::default());
+            targets.insert(
+                B256::from_slice(&padded_key),
+                HashSet::with_hasher(DefaultHashBuilder::default()),
+            );
+            let proof = Proof::new(
+                InMemoryTrieCursorFactory::new(
+                    DatabaseTrieCursorFactory::new(provider.tx_ref()),
+                    &input_nodes_sorted,
+                ),
+                HashedPostStateCursorFactory::new(
+                    DatabaseHashedCursorFactory::new(provider.tx_ref()),
+                    &input_state_sorted,
+                ),
+            )
+            .multiproof(targets)
+            .unwrap();
+
+            // The subtree only contains the proof for a single target.
+            proof.account_subtree.get(&path).cloned()
+        })?;
     }
 
     trie.calculate_below_level(SPARSE_TRIE_INCREMENTAL_LEVEL);

crates/trie/sparse/src/state.rs

@@ -209,7 +209,12 @@ impl SparseStateTrie {
     /// the storage root based on update storage trie or look it up from existing leaf value.
     ///
     /// If the new account info and storage trie are empty, the account leaf will be removed.
-    pub fn update_account(&mut self, address: B256, account: Account) -> SparseStateTrieResult<()> {
+    pub fn update_account(
+        &mut self,
+        address: B256,
+        account: Account,
+        fetch_node: impl FnMut(Nibbles) -> Option<Bytes>,
+    ) -> SparseStateTrieResult<()> {
         let nibbles = Nibbles::unpack(address);
         let storage_root = if let Some(storage_trie) = self.storages.get_mut(&address) {
             storage_trie.root().ok_or(SparseTrieError::Blind)?
@@ -228,7 +233,7 @@ impl SparseStateTrie {
         };
 
         if account.is_empty() && storage_root == EMPTY_ROOT_HASH {
-            self.remove_account_leaf(&nibbles)
+            self.remove_account_leaf(&nibbles, fetch_node)
         } else {
             self.account_rlp_buf.clear();
             TrieAccount::from((account, storage_root)).encode(&mut self.account_rlp_buf);
@@ -247,8 +252,12 @@ impl SparseStateTrie {
     }
 
     /// Remove the account leaf node.
-    pub fn remove_account_leaf(&mut self, path: &Nibbles) -> SparseStateTrieResult<()> {
-        self.state.remove_leaf(path)?;
+    pub fn remove_account_leaf(
+        &mut self,
+        path: &Nibbles,
+        fetch_node: impl FnMut(Nibbles) -> Option<Bytes>,
+    ) -> SparseStateTrieResult<()> {
+        self.state.remove_leaf(path, fetch_node)?;
         Ok(())
     }
 
@@ -268,8 +277,9 @@ impl SparseStateTrie {
         &mut self,
         address: B256,
         slot: &Nibbles,
+        fetch_node: impl FnMut(Nibbles) -> Option<Bytes>,
     ) -> SparseStateTrieResult<()> {
-        self.storages.entry(address).or_default().remove_leaf(slot)?;
+        self.storages.entry(address).or_default().remove_leaf(slot, fetch_node)?;
         Ok(())
     }