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

Cargo.toml

@@ -419,6 +419,7 @@ reth-trie = { path = "crates/trie/trie" }
 reth-trie-common = { path = "crates/trie/common" }
 reth-trie-db = { path = "crates/trie/db" }
 reth-trie-parallel = { path = "crates/trie/parallel" }
+reth-trie-sparse = { path = "crates/trie/sparse" }
 
 # revm
 revm = { version = "18.0.0", features = ["std"], default-features = false }

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-consensus.workspace = true
 alloy-eips.workspace = true
+alloy-primitives.workspace = true
+alloy-rlp.workspace = true
 alloy-rpc-types-engine.workspace = true
-alloy-consensus.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-chain-state/test-utils",
-    "reth-network-p2p/test-utils",
-    "reth-prune-types",
-    "reth-stages/test-utils",
-    "reth-static-file",
-    "reth-tracing",
     "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-provider/test-utils",
+    "reth-stages/test-utils",
+    "reth-static-file",
+    "reth-tracing",
     "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::*;

crates/trie/common/src/proofs.rs

@@ -26,6 +26,31 @@ pub struct MultiProof {
 }
 
 impl MultiProof {
+    /// Return the account proof nodes for the given account path.
+    pub fn account_proof_nodes(&self, path: &Nibbles) -> Vec<(Nibbles, Bytes)> {
+        self.account_subtree.matching_nodes_sorted(path)
+    }
+
+    /// Return the storage proof nodes for the given storage slots of the account path.
+    pub fn storage_proof_nodes(
+        &self,
+        hashed_address: B256,
+        slots: impl IntoIterator<Item = B256>,
+    ) -> Vec<(B256, Vec<(Nibbles, Bytes)>)> {
+        self.storages
+            .get(&hashed_address)
+            .map(|storage_mp| {
+                slots
+                    .into_iter()
+                    .map(|slot| {
+                        let nibbles = Nibbles::unpack(slot);
+                        (slot, storage_mp.subtree.matching_nodes_sorted(&nibbles))
+                    })
+                    .collect()
+            })
+            .unwrap_or_default()
+    }
+
     /// Construct the account proof from the multiproof.
     pub fn account_proof(
         &self,
@@ -37,10 +62,9 @@ impl MultiProof {
 
         // Retrieve the account proof.
         let proof = self
-            .account_subtree
-            .matching_nodes_iter(&nibbles)
-            .sorted_by(|a, b| a.0.cmp(b.0))
-            .map(|(_, node)| node.clone())
+            .account_proof_nodes(&nibbles)
+            .into_iter()
+            .map(|(_, node)| node)
             .collect::<Vec<_>>();
 
         // Inspect the last node in the proof. If it's a leaf node with matching suffix,