Trie hash optimizations (#5827)

This commit is contained in:
Vitaly Drogan
2023-12-26 11:43:10 +02:00
committed by GitHub
parent cbf6324594
commit abc168efa6
13 changed files with 123 additions and 75 deletions

7
Cargo.lock generated
View File

@@ -4607,9 +4607,9 @@ dependencies = [
[[package]]
name = "nybbles"
version = "0.1.1"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "47dddada2357f8e7786f4f4d837db7bdddec02c7c3e5da7840d92c70390f6dc0"
checksum = "836816c354fb2c09622b54545a6f98416147346b13cc7eba5f92fab6b3042c93"
dependencies = [
"alloy-rlp",
"arbitrary",
@@ -6240,6 +6240,7 @@ dependencies = [
name = "reth-primitives"
version = "0.1.0-alpha.13"
dependencies = [
"ahash",
"alloy-primitives",
"alloy-rlp",
"alloy-trie",
@@ -6289,6 +6290,7 @@ dependencies = [
name = "reth-provider"
version = "0.1.0-alpha.13"
dependencies = [
"ahash",
"alloy-rlp",
"assert_matches",
"auto_impl",
@@ -6665,6 +6667,7 @@ dependencies = [
name = "reth-trie"
version = "0.1.0-alpha.13"
dependencies = [
"ahash",
"alloy-rlp",
"auto_impl",
"criterion",

View File

@@ -179,6 +179,7 @@ metrics = "0.21.1" # Needed for `metrics-macro` to resolve the crate using `::me
hex-literal = "0.4"
once_cell = "1.17"
syn = "2.0"
ahash = "0.8.6"
# proc-macros
proc-macro2 = "1.0"

View File

@@ -21,7 +21,7 @@ alloy-primitives = { workspace = true, features = ["rand", "rlp"] }
alloy-rlp = { workspace = true, features = ["arrayvec"] }
alloy-trie = { version = "0.1", features = ["serde"] }
ethers-core = { workspace = true, default-features = false, optional = true }
nybbles = { version = "0.1", features = ["serde", "rlp"] }
nybbles = { version = "0.1.2", features = ["serde", "rlp"] }
# crypto
secp256k1 = { workspace = true, features = ["global-context", "recovery"] }
@@ -50,6 +50,7 @@ sucds = "~0.6"
tempfile.workspace = true
thiserror.workspace = true
zstd = { version = "0.12", features = ["experimental"] }
ahash.workspace = true
# `test-utils` feature
hash-db = { version = "~0.15", optional = true }

View File

@@ -36,6 +36,7 @@ pin-project.workspace = true
parking_lot.workspace = true
dashmap = { version = "5.5", features = ["inline"] }
strum.workspace = true
ahash.workspace = true
# test-utils
alloy-rlp = { workspace = true, optional = true }

View File

@@ -11,6 +11,7 @@ use crate::{
PruneCheckpointWriter, StageCheckpointReader, StorageReader, TransactionVariant,
TransactionsProvider, TransactionsProviderExt, WithdrawalsProvider,
};
use ahash::{AHashMap, AHashSet};
use itertools::{izip, Itertools};
use reth_db::{
common::KeyValue,
@@ -49,7 +50,7 @@ use reth_trie::{
};
use revm::primitives::{BlockEnv, CfgEnv, SpecId};
use std::{
collections::{hash_map, BTreeMap, BTreeSet, HashMap, HashSet},
collections::{hash_map, BTreeMap, BTreeSet, HashMap},
fmt::Debug,
ops::{Bound, Deref, DerefMut, Range, RangeBounds, RangeInclusive},
sync::{mpsc, Arc},
@@ -2037,8 +2038,8 @@ impl<TX: DbTxMut + DbTx> HashingWriter for DatabaseProvider<TX> {
) -> ProviderResult<()> {
// Initialize prefix sets.
let mut account_prefix_set = PrefixSetMut::default();
let mut storage_prefix_set: HashMap<B256, PrefixSetMut> = HashMap::default();
let mut destroyed_accounts = HashSet::default();
let mut storage_prefix_set: AHashMap<B256, PrefixSetMut> = AHashMap::default();
let mut destroyed_accounts = AHashSet::default();
let mut durations_recorder = metrics::DurationsRecorder::default();
@@ -2227,8 +2228,8 @@ impl<TX: DbTxMut + DbTx> BlockExecutionWriter for DatabaseProvider<TX> {
// Initialize prefix sets.
let mut account_prefix_set = PrefixSetMut::default();
let mut storage_prefix_set: HashMap<B256, PrefixSetMut> = HashMap::default();
let mut destroyed_accounts = HashSet::default();
let mut storage_prefix_set: AHashMap<B256, PrefixSetMut> = AHashMap::default();
let mut destroyed_accounts = AHashSet::default();
// Unwind account hashes. Add changed accounts to account prefix set.
let hashed_addresses = self.unwind_account_hashing(range.clone())?;

View File

@@ -26,6 +26,7 @@ tracing.workspace = true
thiserror.workspace = true
derive_more = "0.99"
auto_impl = "1"
ahash.workspace = true
# test-utils
triehash = { version = "0.8", optional = true }

View File

@@ -1,12 +1,12 @@
use super::{HashedAccountCursor, HashedCursorFactory, HashedStorageCursor};
use crate::prefix_set::{PrefixSet, PrefixSetMut};
use ahash::{AHashMap, AHashSet};
use reth_db::{
cursor::{DbCursorRO, DbDupCursorRO},
tables,
transaction::DbTx,
};
use reth_primitives::{trie::Nibbles, Account, StorageEntry, B256, U256};
use std::collections::{HashMap, HashSet};
/// The post state account storage with hashed slots.
#[derive(Debug, Clone, Eq, PartialEq)]
@@ -14,7 +14,7 @@ pub struct HashedStorage {
/// Hashed storage slots with non-zero.
non_zero_valued_storage: Vec<(B256, U256)>,
/// Slots that have been zero valued.
zero_valued_slots: HashSet<B256>,
zero_valued_slots: AHashSet<B256>,
/// Whether the storage was wiped or not.
wiped: bool,
/// Whether the storage entries were sorted or not.
@@ -26,7 +26,7 @@ impl HashedStorage {
pub fn new(wiped: bool) -> Self {
Self {
non_zero_valued_storage: Vec::new(),
zero_valued_slots: HashSet::new(),
zero_valued_slots: AHashSet::new(),
wiped,
sorted: true, // empty is sorted
}
@@ -72,9 +72,9 @@ pub struct HashedPostState {
/// Map of hashed addresses to account info.
accounts: Vec<(B256, Account)>,
/// Set of destroyed accounts.
destroyed_accounts: HashSet<B256>,
destroyed_accounts: AHashSet<B256>,
/// Map of hashed addresses to hashed storage.
storages: HashMap<B256, HashedStorage>,
storages: AHashMap<B256, HashedStorage>,
/// Whether the account and storage entries were sorted or not.
sorted: bool,
}
@@ -83,8 +83,8 @@ impl Default for HashedPostState {
fn default() -> Self {
Self {
accounts: Vec::new(),
destroyed_accounts: HashSet::new(),
storages: HashMap::new(),
destroyed_accounts: AHashSet::new(),
storages: AHashMap::new(),
sorted: true, // empty is sorted
}
}
@@ -139,17 +139,17 @@ impl HashedPostState {
}
/// Returns all destroyed accounts.
pub fn destroyed_accounts(&self) -> HashSet<B256> {
pub fn destroyed_accounts(&self) -> AHashSet<B256> {
self.destroyed_accounts.clone()
}
/// Construct (PrefixSet)[PrefixSet] from hashed post state.
/// The prefix sets contain the hashed account and storage keys that have been changed in the
/// post state.
pub fn construct_prefix_sets(&self) -> (PrefixSet, HashMap<B256, PrefixSet>) {
pub fn construct_prefix_sets(&self) -> (PrefixSet, AHashMap<B256, PrefixSet>) {
// Initialize prefix sets.
let mut account_prefix_set = PrefixSetMut::default();
let mut storage_prefix_set: HashMap<B256, PrefixSetMut> = HashMap::default();
let mut storage_prefix_set: AHashMap<B256, PrefixSetMut> = AHashMap::default();
// Populate account prefix set.
for (hashed_address, _) in &self.accounts {

View File

@@ -88,7 +88,7 @@ where
self.current_walker_key_checked = true;
if self.walker.can_skip_current_node {
return Ok(Some(AccountNode::Branch(TrieBranchNode::new(
key,
key.clone(),
self.walker.hash().unwrap(),
self.walker.children_are_in_trie(),
))))
@@ -97,7 +97,7 @@ where
}
if let Some((hashed_address, account)) = self.current_hashed_entry.take() {
if self.walker.key().map_or(false, |key| key < Nibbles::unpack(hashed_address)) {
if self.walker.key().map_or(false, |key| key < &Nibbles::unpack(hashed_address)) {
self.current_walker_key_checked = false;
continue
}
@@ -178,7 +178,7 @@ where
self.current_walker_key_checked = true;
if self.walker.can_skip_current_node {
return Ok(Some(StorageNode::Branch(TrieBranchNode::new(
key,
key.clone(),
self.walker.hash().unwrap(),
self.walker.children_are_in_trie(),
))))
@@ -188,7 +188,7 @@ where
if let Some(StorageEntry { key: hashed_key, value }) = self.current_hashed_entry.take()
{
if self.walker.key().map_or(false, |key| key < Nibbles::unpack(hashed_key)) {
if self.walker.key().map_or(false, |key| key < &Nibbles::unpack(hashed_key)) {
self.current_walker_key_checked = false;
continue
}

View File

@@ -1,4 +1,5 @@
use super::PrefixSetMut;
use ahash::{AHashMap, AHashSet};
use derive_more::Deref;
use reth_db::{
cursor::DbCursorRO,
@@ -8,10 +9,7 @@ use reth_db::{
DatabaseError,
};
use reth_primitives::{keccak256, trie::Nibbles, BlockNumber, StorageEntry, B256};
use std::{
collections::{HashMap, HashSet},
ops::RangeInclusive,
};
use std::ops::RangeInclusive;
/// Loaded prefix sets.
#[derive(Debug, Default)]
@@ -19,9 +17,9 @@ pub struct LoadedPrefixSets {
/// The account prefix set
pub account_prefix_set: PrefixSetMut,
/// The mapping of hashed account key to the corresponding storage prefix set
pub storage_prefix_sets: HashMap<B256, PrefixSetMut>,
pub storage_prefix_sets: AHashMap<B256, PrefixSetMut>,
/// The account keys of destroyed accounts
pub destroyed_accounts: HashSet<B256>,
pub destroyed_accounts: AHashSet<B256>,
}
/// A wrapper around a database transaction that loads prefix sets within a given block range.

View File

@@ -118,20 +118,18 @@ impl PrefixSet {
/// Returns `true` if any of the keys in the set has the given prefix or
/// if the given prefix is a prefix of any key in the set.
#[inline]
pub fn contains<T: Into<Nibbles>>(&mut self, prefix: T) -> bool {
let prefix = prefix.into();
while self.index > 0 && self.keys[self.index] > prefix {
pub fn contains(&mut self, prefix: &Nibbles) -> bool {
while self.index > 0 && &self.keys[self.index] > prefix {
self.index -= 1;
}
for (idx, key) in self.keys[self.index..].iter().enumerate() {
if key.has_prefix(&prefix) {
if key.has_prefix(prefix) {
self.index += idx;
return true
}
if key > &prefix {
if key > prefix {
self.index += idx;
return false
}

View File

@@ -8,6 +8,7 @@ use crate::{
walker::TrieWalker,
StateRootError, StorageRootError,
};
use ahash::{AHashMap, AHashSet};
use alloy_rlp::{BufMut, Encodable};
use reth_db::transaction::DbTx;
use reth_primitives::{
@@ -16,10 +17,7 @@ use reth_primitives::{
trie::{HashBuilder, Nibbles, TrieAccount},
Address, BlockNumber, B256,
};
use std::{
collections::{HashMap, HashSet},
ops::RangeInclusive,
};
use std::ops::RangeInclusive;
use tracing::{debug, trace};
/// StateRoot is used to compute the root node of a state trie.
@@ -33,9 +31,9 @@ pub struct StateRoot<T, H> {
pub changed_account_prefixes: PrefixSet,
/// A map containing storage changes with the hashed address as key and a set of storage key
/// prefixes as the value.
pub changed_storage_prefixes: HashMap<B256, PrefixSet>,
pub changed_storage_prefixes: AHashMap<B256, PrefixSet>,
/// A map containing keys of accounts that were destroyed.
pub destroyed_accounts: HashSet<B256>,
pub destroyed_accounts: AHashSet<B256>,
/// Previous intermediate state.
previous_state: Option<IntermediateStateRootState>,
/// The number of updates after which the intermediate progress should be returned.
@@ -50,13 +48,13 @@ impl<T, H> StateRoot<T, H> {
}
/// Set the changed storage prefixes.
pub fn with_changed_storage_prefixes(mut self, prefixes: HashMap<B256, PrefixSet>) -> Self {
pub fn with_changed_storage_prefixes(mut self, prefixes: AHashMap<B256, PrefixSet>) -> Self {
self.changed_storage_prefixes = prefixes;
self
}
/// Set the destroyed accounts.
pub fn with_destroyed_accounts(mut self, accounts: HashSet<B256>) -> Self {
pub fn with_destroyed_accounts(mut self, accounts: AHashSet<B256>) -> Self {
self.destroyed_accounts = accounts;
self
}
@@ -113,8 +111,8 @@ impl<'a, TX: DbTx> StateRoot<&'a TX, &'a TX> {
trie_cursor_factory: tx,
hashed_cursor_factory: tx,
changed_account_prefixes: PrefixSetMut::default().freeze(),
changed_storage_prefixes: HashMap::default(),
destroyed_accounts: HashSet::default(),
changed_storage_prefixes: AHashMap::default(),
destroyed_accounts: AHashSet::default(),
previous_state: None,
threshold: 100_000,
}
@@ -519,7 +517,12 @@ mod tests {
Account, Address, StorageEntry, B256, U256,
};
use reth_provider::{test_utils::create_test_provider_factory, DatabaseProviderRW};
use std::{collections::BTreeMap, ops::Mul, str::FromStr, sync::Arc};
use std::{
collections::{BTreeMap, HashMap},
ops::Mul,
str::FromStr,
sync::Arc,
};
fn insert_account(
tx: &impl DbTxMut,

View File

@@ -9,9 +9,11 @@ pub struct CursorSubNode {
/// The key of the current node.
pub key: Nibbles,
/// The index of the next child to visit.
pub nibble: i8,
nibble: i8,
/// The node itself.
pub node: Option<BranchNodeCompact>,
/// Full key
full_key: Nibbles,
}
impl Default for CursorSubNode {
@@ -39,7 +41,9 @@ impl From<StoredSubNode> for CursorSubNode {
Some(n) => n as i8,
None => -1,
};
Self { key: Nibbles::from_nibbles_unchecked(value.key), nibble, node: value.node }
let key = Nibbles::from_nibbles_unchecked(value.key);
let full_key = full_key(key.clone(), nibble);
Self { key, nibble, node: value.node, full_key }
}
}
@@ -60,16 +64,13 @@ impl CursorSubNode {
}
_ => -1,
};
CursorSubNode { key, node, nibble }
let full_key = full_key(key.clone(), nibble);
CursorSubNode { key, node, nibble, full_key }
}
/// Returns the full key of the current node.
pub fn full_key(&self) -> Nibbles {
let mut out = self.key.clone();
if self.nibble >= 0 {
out.push(self.nibble as u8);
}
out
pub fn full_key(&self) -> &Nibbles {
&self.full_key
}
/// Returns `true` if the state flag is set for the current nibble.
@@ -117,4 +118,47 @@ impl CursorSubNode {
None
}
}
/// Returns the next child index to visit.
#[inline]
pub fn nibble(&self) -> i8 {
self.nibble
}
/// Increments the nibble index.
#[inline]
pub fn inc_nibble(&mut self) {
self.nibble += 1;
update_full_key(&mut self.full_key, self.nibble - 1, self.nibble);
}
/// Sets the nibble index.
#[inline]
pub fn set_nibble(&mut self, nibble: i8) {
let old_nibble = self.nibble;
self.nibble = nibble;
update_full_key(&mut self.full_key, old_nibble, self.nibble);
}
}
#[inline]
fn full_key(mut key: Nibbles, nibble: i8) -> Nibbles {
if nibble >= 0 {
key.push(nibble as u8);
}
key
}
#[inline]
fn update_full_key(key: &mut Nibbles, old_nibble: i8, new_nibble: i8) {
if new_nibble >= 0 {
if old_nibble >= 0 {
let last_index = key.len() - 1;
key.set_at(last_index, new_nibble as u8);
} else {
key.push(new_nibble as u8);
}
} else if old_nibble >= 0 {
key.pop();
}
}

View File

@@ -101,7 +101,7 @@ impl<C: TrieCursor> TrieWalker<C> {
if !self.can_skip_current_node && self.children_are_in_trie() {
// If we can't skip the current node and the children are in the trie,
// either consume the next node or move to the next sibling.
match last.nibble {
match last.nibble() {
-1 => self.move_to_next_sibling(true)?,
_ => self.consume_node()?,
}
@@ -115,7 +115,7 @@ impl<C: TrieCursor> TrieWalker<C> {
}
// Return the current key.
Ok(self.key())
Ok(self.key().cloned())
}
/// Retrieves the current root node from the DB, seeking either the exact node or the next one.
@@ -146,12 +146,12 @@ impl<C: TrieCursor> TrieWalker<C> {
// We need to sync the stack with the trie structure when consuming a new node. This is
// necessary for proper traversal and accurately representing the trie in the stack.
if !key.is_empty() && !self.stack.is_empty() {
self.stack[0].nibble = key[0] as i8;
self.stack[0].set_nibble(key[0] as i8);
}
// Create a new CursorSubNode and push it to the stack.
let subnode = CursorSubNode::new(key, Some(node));
let nibble = subnode.nibble;
let nibble = subnode.nibble();
self.stack.push(subnode);
self.update_skip_node();
@@ -175,24 +175,24 @@ impl<C: TrieCursor> TrieWalker<C> {
// Check if the walker needs to backtrack to the previous level in the trie during its
// traversal.
if subnode.nibble >= 15 || (subnode.nibble < 0 && !allow_root_to_child_nibble) {
if subnode.nibble() >= 15 || (subnode.nibble() < 0 && !allow_root_to_child_nibble) {
self.stack.pop();
self.move_to_next_sibling(false)?;
return Ok(())
}
subnode.nibble += 1;
subnode.inc_nibble();
if subnode.node.is_none() {
return self.consume_node()
}
// Find the next sibling with state.
while subnode.nibble < 16 {
while subnode.nibble() < 16 {
if subnode.state_flag() {
return Ok(())
}
subnode.nibble += 1;
subnode.inc_nibble();
}
// Pop the current node and move to the next sibling.
@@ -203,7 +203,7 @@ impl<C: TrieCursor> TrieWalker<C> {
}
/// Returns the current key in the trie.
pub fn key(&self) -> Option<Nibbles> {
pub fn key(&self) -> Option<&Nibbles> {
self.stack.last().map(|n| n.full_key())
}
@@ -220,7 +220,6 @@ impl<C: TrieCursor> TrieWalker<C> {
/// Returns the next unprocessed key in the trie.
pub fn next_unprocessed_key(&self) -> Option<B256> {
self.key()
.as_ref()
.and_then(|key| {
if self.can_skip_current_node {
key.increment().map(|inc| inc.pack())
@@ -235,10 +234,8 @@ impl<C: TrieCursor> TrieWalker<C> {
}
fn update_skip_node(&mut self) {
self.can_skip_current_node = if let Some(key) = self.key() {
let contains_prefix = self.changes.contains(key);
let hash_flag = self.stack.last().unwrap().hash_flag();
!contains_prefix && hash_flag
self.can_skip_current_node = if let Some(node) = self.stack.last() {
!self.changes.contains(node.full_key()) && node.hash_flag()
} else {
false
};
@@ -361,10 +358,10 @@ mod tests {
// No changes
let mut cursor = TrieWalker::new(&mut trie, Default::default());
assert_eq!(cursor.key(), Some(Nibbles::from_nibbles_unchecked([]))); // root
assert_eq!(cursor.key().cloned(), Some(Nibbles::from_nibbles_unchecked([]))); // root
assert!(cursor.can_skip_current_node); // due to root_hash
cursor.advance().unwrap(); // skips to the end of trie
assert_eq!(cursor.key(), None);
assert_eq!(cursor.key().cloned(), None);
// We insert something that's not part of the existing trie/prefix.
let mut changed = PrefixSetMut::default();
@@ -372,17 +369,17 @@ mod tests {
let mut cursor = TrieWalker::new(&mut trie, changed.freeze());
// Root node
assert_eq!(cursor.key(), Some(Nibbles::from_nibbles_unchecked([])));
assert_eq!(cursor.key().cloned(), Some(Nibbles::from_nibbles_unchecked([])));
// Should not be able to skip state due to the changed values
assert!(!cursor.can_skip_current_node);
cursor.advance().unwrap();
assert_eq!(cursor.key(), Some(Nibbles::from_nibbles_unchecked([0x2])));
assert_eq!(cursor.key().cloned(), Some(Nibbles::from_nibbles_unchecked([0x2])));
cursor.advance().unwrap();
assert_eq!(cursor.key(), Some(Nibbles::from_nibbles_unchecked([0x2, 0x1])));
assert_eq!(cursor.key().cloned(), Some(Nibbles::from_nibbles_unchecked([0x2, 0x1])));
cursor.advance().unwrap();
assert_eq!(cursor.key(), Some(Nibbles::from_nibbles_unchecked([0x4])));
assert_eq!(cursor.key().cloned(), Some(Nibbles::from_nibbles_unchecked([0x4])));
cursor.advance().unwrap();
assert_eq!(cursor.key(), None); // the end of trie
assert_eq!(cursor.key().cloned(), None); // the end of trie
}
}