From 6b088bd8819c9e1d66efc0a050c8c98e37ffd457 Mon Sep 17 00:00:00 2001 From: Alexey Shekhirin Date: Mon, 25 Nov 2024 14:15:12 +0000 Subject: [PATCH] perf(engine): sparse trie calculation for state root task (#12843) --- Cargo.lock | 1 + Cargo.toml | 1 + crates/engine/tree/Cargo.toml | 51 ++++++++-------- crates/engine/tree/src/tree/root.rs | 90 +++++++++++++++++++++++++++-- crates/trie/common/src/proofs.rs | 32 ++++++++-- 5 files changed, 142 insertions(+), 33 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a3ccf46b3..c8d4d2b51 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7281,6 +7281,7 @@ dependencies = [ "reth-tracing", "reth-trie", "reth-trie-parallel", + "reth-trie-sparse", "revm-primitives", "thiserror 1.0.69", "tokio", diff --git a/Cargo.toml b/Cargo.toml index ad17ea4ad..113d0661f 100644 --- a/Cargo.toml +++ b/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 } diff --git a/crates/engine/tree/Cargo.toml b/crates/engine/tree/Cargo.toml index 70be84a9f..5242268b1 100644 --- a/crates/engine/tree/Cargo.toml +++ b/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", ] diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs index 45cf5a780..27f835ec7 100644 --- a/crates/engine/tree/src/tree/root.rs +++ b/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, + multiproof: MultiProof, + targets: FbHashMap<32, FbHashSet<32>>, + state: HashedPostState, +) -> SparseStateTrieResult<(Box, 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::*; diff --git a/crates/trie/common/src/proofs.rs b/crates/trie/common/src/proofs.rs index d0a5cd220..8e014f6d8 100644 --- a/crates/trie/common/src/proofs.rs +++ b/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, + ) -> 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::>(); // Inspect the last node in the proof. If it's a leaf node with matching suffix,