diff --git a/crates/trie/common/src/hashed_state.rs b/crates/trie/common/src/hashed_state.rs index 50d9f20af0b..015c322bd96 100644 --- a/crates/trie/common/src/hashed_state.rs +++ b/crates/trie/common/src/hashed_state.rs @@ -8,7 +8,7 @@ use crate::{ use alloc::{borrow::Cow, vec::Vec}; use alloy_primitives::{ keccak256, - map::{hash_map, B256Map, B256Set, HashMap, HashSet}, + map::{hash_map, B256Map, HashMap, HashSet}, Address, B256, U256, }; use itertools::Itertools; @@ -323,17 +323,20 @@ impl HashedPostState { /// Converts hashed post state into [`HashedPostStateSorted`]. pub fn into_sorted(self) -> HashedPostStateSorted { - let mut updated_accounts = Vec::new(); - let mut destroyed_accounts = HashSet::default(); + let mut accounts = Vec::new(); + let mut destroyed_accounts = Vec::new(); for (hashed_address, info) in self.accounts { if let Some(info) = info { - updated_accounts.push((hashed_address, info)); + accounts.push((hashed_address, Some(info))); } else { - destroyed_accounts.insert(hashed_address); + accounts.push((hashed_address, None)); + destroyed_accounts.push(hashed_address); } } - updated_accounts.sort_unstable_by_key(|(address, _)| *address); - let accounts = HashedAccountsSorted { accounts: updated_accounts, destroyed_accounts }; + accounts.sort_unstable_by_key(|(address, _)| *address); + destroyed_accounts.sort_unstable(); + + let accounts = HashedAccountsSorted { accounts, destroyed_accounts }; let storages = self .storages @@ -352,17 +355,20 @@ impl HashedPostState { /// This allows us to reuse the allocated space. This allocates new space for the sorted hashed /// post state, like `into_sorted`. pub fn drain_into_sorted(&mut self) -> HashedPostStateSorted { - let mut updated_accounts = Vec::new(); - let mut destroyed_accounts = HashSet::default(); + let mut accounts = Vec::new(); + let mut destroyed_accounts = Vec::new(); for (hashed_address, info) in self.accounts.drain() { if let Some(info) = info { - updated_accounts.push((hashed_address, info)); + accounts.push((hashed_address, Some(info))); } else { - destroyed_accounts.insert(hashed_address); + accounts.push((hashed_address, None)); + destroyed_accounts.push(hashed_address); } } - updated_accounts.sort_unstable_by_key(|(address, _)| *address); - let accounts = HashedAccountsSorted { accounts: updated_accounts, destroyed_accounts }; + accounts.sort_unstable_by_key(|(address, _)| *address); + destroyed_accounts.sort_unstable(); + + let accounts = HashedAccountsSorted { accounts, destroyed_accounts }; let storages = self .storages @@ -442,18 +448,20 @@ impl HashedStorage { /// Converts hashed storage into [`HashedStorageSorted`]. pub fn into_sorted(self) -> HashedStorageSorted { - let mut non_zero_valued_slots = Vec::new(); - let mut zero_valued_slots = HashSet::default(); + let mut slots = Vec::new(); + let mut zero_valued_slots = Vec::new(); for (hashed_slot, value) in self.storage { if value.is_zero() { - zero_valued_slots.insert(hashed_slot); + slots.push((hashed_slot, None)); + zero_valued_slots.push(hashed_slot); } else { - non_zero_valued_slots.push((hashed_slot, value)); + slots.push((hashed_slot, Some(value))); } } - non_zero_valued_slots.sort_unstable_by_key(|(key, _)| *key); + slots.sort_unstable_by_key(|(key, _)| *key); + zero_valued_slots.sort_unstable(); - HashedStorageSorted { non_zero_valued_slots, zero_valued_slots, wiped: self.wiped } + HashedStorageSorted { slots, zero_valued_slots, wiped: self.wiped } } } @@ -495,30 +503,28 @@ impl AsRef for HashedPostStateSorted { /// Sorted account state optimized for iterating during state trie calculation. #[derive(Clone, Eq, PartialEq, Default, Debug)] pub struct HashedAccountsSorted { - /// Sorted collection of hashed addresses and their account info. - pub accounts: Vec<(B256, Account)>, - /// Set of destroyed account keys. - pub destroyed_accounts: B256Set, + /// Sorted collection of hashed addresses and their account info (Some for updates, None for + /// deletions). + pub accounts: Vec<(B256, Option)>, + /// Sorted collection of destroyed account keys for efficient lookups. + pub destroyed_accounts: Vec, } impl HashedAccountsSorted { - /// Returns a sorted iterator over updated accounts. - pub fn accounts_sorted(&self) -> impl Iterator)> { - self.accounts - .iter() - .map(|(address, account)| (*address, Some(*account))) - .chain(self.destroyed_accounts.iter().map(|address| (*address, None))) - .sorted_by_key(|entry| *entry.0) + /// Returns a sorted iterator over updated accounts (Some) and destroyed accounts (None). + pub fn accounts_sorted(&self) -> impl Iterator)> + '_ { + self.accounts.iter().copied() } } /// Sorted hashed storage optimized for iterating during state trie calculation. #[derive(Clone, Eq, PartialEq, Debug)] pub struct HashedStorageSorted { - /// Sorted hashed storage slots with non-zero value. - pub non_zero_valued_slots: Vec<(B256, U256)>, - /// Slots that have been zero valued. - pub zero_valued_slots: B256Set, + /// Sorted hashed storage slots with their values (Some for non-zero updates, None for + /// zero/deletions). + pub slots: Vec<(B256, Option)>, + /// Sorted slots that have been zero-valued for efficient lookups. + pub zero_valued_slots: Vec, /// Flag indicating whether the storage was wiped or not. pub wiped: bool, } @@ -531,11 +537,7 @@ impl HashedStorageSorted { /// Returns a sorted iterator over updated storage slots. pub fn storage_slots_sorted(&self) -> impl Iterator { - self.non_zero_valued_slots - .iter() - .map(|(hashed_slot, value)| (*hashed_slot, *value)) - .chain(self.zero_valued_slots.iter().map(|hashed_slot| (*hashed_slot, U256::ZERO))) - .sorted_by_key(|entry| *entry.0) + self.slots.iter().map(|(slot, value)| (*slot, value.unwrap_or(U256::ZERO))) } } diff --git a/crates/trie/trie/src/hashed_cursor/post_state.rs b/crates/trie/trie/src/hashed_cursor/post_state.rs index e81aa4af22a..0d6c4675322 100644 --- a/crates/trie/trie/src/hashed_cursor/post_state.rs +++ b/crates/trie/trie/src/hashed_cursor/post_state.rs @@ -1,6 +1,6 @@ use super::{HashedCursor, HashedCursorFactory, HashedStorageCursor}; use crate::forward_cursor::ForwardInMemoryCursor; -use alloy_primitives::{map::B256Set, B256, U256}; +use alloy_primitives::{B256, U256}; use reth_primitives_traits::Account; use reth_storage_errors::db::DatabaseError; use reth_trie_common::{HashedAccountsSorted, HashedPostStateSorted, HashedStorageSorted}; @@ -50,10 +50,10 @@ where pub struct HashedPostStateAccountCursor<'a, C> { /// The database cursor. cursor: C, - /// Forward-only in-memory cursor over accounts. - post_state_cursor: ForwardInMemoryCursor<'a, B256, Account>, - /// Reference to the collection of account keys that were destroyed. - destroyed_accounts: &'a B256Set, + /// Forward-only in-memory cursor over accounts (Some for updates, None for deletions). + post_state_cursor: ForwardInMemoryCursor<'a, B256, Option>, + /// Reference to the sorted collection of destroyed account keys. + destroyed_accounts: &'a Vec, /// The last hashed account that was returned by the cursor. /// De facto, this is a current cursor position. last_account: Option, @@ -76,18 +76,20 @@ where /// This function only checks the post state, not the database, because the latter does not /// store destroyed accounts. fn is_account_cleared(&self, account: &B256) -> bool { - self.destroyed_accounts.contains(account) + self.destroyed_accounts.binary_search(account).is_ok() } fn seek_inner(&mut self, key: B256) -> Result, DatabaseError> { // Take the next account from the post state with the key greater than or equal to the // sought key. let post_state_entry = self.post_state_cursor.seek(&key); + // Filter to only valid updates (Some). + let post_state_account = post_state_entry.and_then(|(k, acc)| acc.map(|a| (k, a))); // It's an exact match, return the account from post state without looking up in the // database. - if post_state_entry.is_some_and(|entry| entry.0 == key) { - return Ok(post_state_entry) + if post_state_account.is_some_and(|(k, _)| k == key) { + return Ok(post_state_account); } // It's not an exact match, reposition to the first greater or equal account that wasn't @@ -98,12 +100,14 @@ where } // Compare two entries and return the lowest. - Ok(Self::compare_entries(post_state_entry, db_entry)) + Ok(Self::compare_entries(post_state_account, db_entry)) } fn next_inner(&mut self, last_account: B256) -> Result, DatabaseError> { // Take the next account from the post state with the key greater than the last sought key. let post_state_entry = self.post_state_cursor.first_after(&last_account); + // Filter to only valid updates (Some). + let post_state_account = post_state_entry.and_then(|(k, acc)| acc.map(|a| (k, a))); // If post state was given precedence or account was cleared, move the cursor forward. let mut db_entry = self.cursor.seek(last_account)?; @@ -114,7 +118,7 @@ where } // Compare two entries and return the lowest. - Ok(Self::compare_entries(post_state_entry, db_entry)) + Ok(Self::compare_entries(post_state_account, db_entry)) } /// Return the account with the lowest hashed account key. @@ -184,10 +188,11 @@ where pub struct HashedPostStateStorageCursor<'a, C> { /// The database cursor. cursor: C, - /// Forward-only in-memory cursor over non zero-valued account storage slots. - post_state_cursor: Option>, - /// Reference to the collection of storage slot keys that were cleared. - cleared_slots: Option<&'a B256Set>, + /// Forward-only in-memory cursor over storage slots (Some for non-zero updates, None for + /// zero/deletions). + post_state_cursor: Option>>, + /// Reference to the sorted collection of zero-valued slot keys. + cleared_slots: Option<&'a Vec>, /// Flag indicating whether database storage was wiped. storage_wiped: bool, /// The last slot that has been returned by the cursor. @@ -201,8 +206,7 @@ where { /// Create new instance of [`HashedPostStateStorageCursor`] for the given hashed address. pub fn new(cursor: C, post_state_storage: Option<&'a HashedStorageSorted>) -> Self { - let post_state_cursor = - post_state_storage.map(|s| ForwardInMemoryCursor::new(&s.non_zero_valued_slots)); + let post_state_cursor = post_state_storage.map(|s| ForwardInMemoryCursor::new(&s.slots)); let cleared_slots = post_state_storage.map(|s| &s.zero_valued_slots); let storage_wiped = post_state_storage.is_some_and(|s| s.wiped); Self { cursor, post_state_cursor, cleared_slots, storage_wiped, last_slot: None } @@ -211,18 +215,20 @@ where /// Check if the slot was zeroed out in the post state. /// The database is not checked since it already has no zero-valued slots. fn is_slot_zero_valued(&self, slot: &B256) -> bool { - self.cleared_slots.is_some_and(|s| s.contains(slot)) + self.cleared_slots.as_ref().is_some_and(|s| s.binary_search(slot).is_ok()) } /// Find the storage entry in post state or database that's greater or equal to provided subkey. fn seek_inner(&mut self, subkey: B256) -> Result, DatabaseError> { // Attempt to find the account's storage in post state. let post_state_entry = self.post_state_cursor.as_mut().and_then(|c| c.seek(&subkey)); + // Filter to only valid updates (Some). + let post_state_slot = post_state_entry.and_then(|(k, v)| v.map(|val| (k, val))); // If database storage was wiped or it's an exact match, // return the storage slot from post state without looking up in the database. - if self.storage_wiped || post_state_entry.is_some_and(|entry| entry.0 == subkey) { - return Ok(post_state_entry) + if self.storage_wiped || post_state_slot.is_some_and(|(k, _)| k == subkey) { + return Ok(post_state_slot); } // It's not an exact match and storage was not wiped, @@ -233,7 +239,7 @@ where } // Compare two entries and return the lowest. - Ok(Self::compare_entries(post_state_entry, db_entry)) + Ok(Self::compare_entries(post_state_slot, db_entry)) } /// Find the storage entry that is right after current cursor position. @@ -241,10 +247,12 @@ where // Attempt to find the account's storage in post state. let post_state_entry = self.post_state_cursor.as_mut().and_then(|c| c.first_after(&last_slot)); + // Filter to only valid updates (Some). + let post_state_slot = post_state_entry.and_then(|(k, v)| v.map(|val| (k, val))); // Return post state entry immediately if database was wiped. if self.storage_wiped { - return Ok(post_state_entry) + return Ok(post_state_slot); } // If post state was given precedence, move the cursor forward. @@ -258,7 +266,7 @@ where } // Compare two entries and return the lowest. - Ok(Self::compare_entries(post_state_entry, db_entry)) + Ok(Self::compare_entries(post_state_slot, db_entry)) } /// Return the storage entry with the lowest hashed storage key (hashed slot).