perf(engine): sparse trie calculation for state root task (#12843)

crates/engine/tree/Cargo.toml

@@ -13,41 +13,43 @@ workspace = true
 [dependencies]
 # reth
 reth-beacon-consensus.workspace = true
-reth-blockchain-tree.workspace = true
 reth-blockchain-tree-api.workspace = true
+reth-blockchain-tree.workspace = true
 reth-chain-state.workspace = true
-reth-consensus.workspace = true
 reth-chainspec.workspace = true
+reth-consensus.workspace = true
 reth-engine-primitives.workspace = true
 reth-errors.workspace = true
 reth-evm.workspace = true
 reth-network-p2p.workspace = true
-reth-payload-builder.workspace = true
 reth-payload-builder-primitives.workspace = true
+reth-payload-builder.workspace = true
 reth-payload-primitives.workspace = true
 reth-payload-validator.workspace = true
-reth-primitives.workspace = true
 reth-primitives-traits.workspace = true
+reth-primitives.workspace = true
 reth-provider.workspace = true
 reth-prune.workspace = true
 reth-revm.workspace = true
 reth-stages-api.workspace = true
 reth-tasks.workspace = true
-reth-trie.workspace = true
 reth-trie-parallel.workspace = true
+reth-trie-sparse.workspace = true
+reth-trie.workspace = true
 
 # alloy
-alloy-primitives.workspace = true
-alloy-eips.workspace = true
-alloy-rpc-types-engine.workspace = true
 alloy-consensus.workspace = true
+alloy-eips.workspace = true
+alloy-primitives.workspace = true
+alloy-rlp.workspace = true
+alloy-rpc-types-engine.workspace = true
 
 revm-primitives.workspace = true
 
 # common
 futures.workspace = true
-tokio = { workspace = true, features = ["macros", "sync"] }
 thiserror.workspace = true
+tokio = { workspace = true, features = ["macros", "sync"] }
 
 # metrics
 metrics.workspace = true
@@ -64,20 +66,21 @@ reth-tracing = { workspace = true, optional = true }
 
 [dev-dependencies]
 # reth
-reth-db = { workspace = true, features = ["test-utils"] }
 reth-chain-state = { workspace = true, features = ["test-utils"] }
+reth-chainspec.workspace = true
+reth-db = { workspace = true, features = ["test-utils"] }
 reth-ethereum-engine-primitives.workspace = true
 reth-evm = { workspace = true, features = ["test-utils"] }
 reth-exex-types.workspace = true
 reth-network-p2p = { workspace = true, features = ["test-utils"] }
-reth-prune.workspace = true
 reth-prune-types.workspace = true
+reth-prune.workspace = true
 reth-rpc-types-compat.workspace = true
 reth-stages = { workspace = true, features = ["test-utils"] }
 reth-static-file.workspace = true
 reth-tracing.workspace = true
-reth-chainspec.workspace = true
 
+# alloy
 alloy-rlp.workspace = true
 
 assert_matches.workspace = true
@@ -90,23 +93,23 @@ harness = false
 
 [features]
 test-utils = [
-    "reth-db/test-utils",
+    "reth-blockchain-tree/test-utils",
     "reth-chain-state/test-utils",
+    "reth-chainspec/test-utils",
+    "reth-consensus/test-utils",
+    "reth-db/test-utils",
+    "reth-evm/test-utils",
     "reth-network-p2p/test-utils",
+    "reth-payload-builder/test-utils",
+    "reth-primitives-traits/test-utils",
+    "reth-primitives/test-utils",
+    "reth-provider/test-utils",
     "reth-prune-types",
+    "reth-prune-types?/test-utils",
+    "reth-revm/test-utils",
+    "reth-stages-api/test-utils",
     "reth-stages/test-utils",
     "reth-static-file",
     "reth-tracing",
-    "reth-blockchain-tree/test-utils",
-    "reth-chainspec/test-utils",
-    "reth-consensus/test-utils",
-    "reth-evm/test-utils",
-    "reth-payload-builder/test-utils",
-    "reth-primitives/test-utils",
-    "reth-revm/test-utils",
-    "reth-stages-api/test-utils",
-    "reth-provider/test-utils",
     "reth-trie/test-utils",
-    "reth-prune-types?/test-utils",
-    "reth-primitives-traits/test-utils",
 ]

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

@@ -1,15 +1,27 @@
 //! State root task related functionality.
 
+use alloy_primitives::map::FbHashMap;
+use alloy_rlp::{BufMut, Encodable};
 use reth_provider::providers::ConsistentDbView;
-use reth_trie::{updates::TrieUpdates, TrieInput};
+use reth_trie::{
+    updates::TrieUpdates, HashedPostState, MultiProof, Nibbles, TrieAccount, TrieInput,
+    EMPTY_ROOT_HASH,
+};
 use reth_trie_parallel::root::ParallelStateRootError;
-use revm_primitives::{EvmState, B256};
-use std::sync::{
-    mpsc::{self, Receiver, RecvError},
-    Arc,
+use reth_trie_sparse::{SparseStateTrie, SparseStateTrieResult};
+use revm_primitives::{map::FbHashSet, EvmState, B256};
+use std::{
+    sync::{
+        mpsc::{self, Receiver, RecvError},
+        Arc,
+    },
+    time::{Duration, Instant},
 };
 use tracing::debug;
 
+/// The level below which the sparse trie hashes are calculated in [`update_sparse_trie`].
+const SPARSE_TRIE_INCREMENTAL_LEVEL: usize = 2;
+
 /// Result of the state root calculation
 pub(crate) type StateRootResult = Result<(B256, TrieUpdates), ParallelStateRootError>;
 
@@ -133,6 +145,74 @@ where
     }
 }
 
+/// Updates the sparse trie with the given proofs and state, and returns the updated trie and the
+/// time it took.
+#[allow(dead_code)]
+fn update_sparse_trie(
+    mut trie: Box<SparseStateTrie>,
+    multiproof: MultiProof,
+    targets: FbHashMap<32, FbHashSet<32>>,
+    state: HashedPostState,
+) -> SparseStateTrieResult<(Box<SparseStateTrie>, Duration)> {
+    let started_at = Instant::now();
+
+    // Reveal new accounts and storage slots.
+    for (address, slots) in targets {
+        let path = Nibbles::unpack(address);
+        trie.reveal_account(address, multiproof.account_proof_nodes(&path))?;
+
+        let storage_proofs = multiproof.storage_proof_nodes(address, slots);
+
+        for (slot, proof) in storage_proofs {
+            trie.reveal_storage_slot(address, slot, proof)?;
+        }
+    }
+
+    // Update storage slots with new values and calculate storage roots.
+    let mut storage_roots = FbHashMap::default();
+    for (address, storage) in state.storages {
+        if storage.wiped {
+            trie.wipe_storage(address)?;
+            storage_roots.insert(address, EMPTY_ROOT_HASH);
+        }
+
+        for (slot, value) in storage.storage {
+            let slot_path = Nibbles::unpack(slot);
+            trie.update_storage_leaf(
+                address,
+                slot_path,
+                alloy_rlp::encode_fixed_size(&value).to_vec(),
+            )?;
+        }
+
+        storage_roots.insert(address, trie.storage_root(address).unwrap());
+    }
+
+    // Update accounts with new values and include updated storage roots
+    for (address, account) in state.accounts {
+        let path = Nibbles::unpack(address);
+
+        if let Some(account) = account {
+            let storage_root = storage_roots
+                .remove(&address)
+                .map(Some)
+                .unwrap_or_else(|| trie.storage_root(address))
+                .unwrap_or(EMPTY_ROOT_HASH);
+
+            let mut encoded = Vec::with_capacity(128);
+            TrieAccount::from((account, storage_root)).encode(&mut encoded as &mut dyn BufMut);
+            trie.update_account_leaf(path, encoded)?;
+        } else {
+            trie.remove_account_leaf(&path)?;
+        }
+    }
+
+    trie.calculate_below_level(SPARSE_TRIE_INCREMENTAL_LEVEL);
+    let elapsed = started_at.elapsed();
+
+    Ok((trie, elapsed))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;