use crate::{
prefix_set::{PrefixSetMut, TriePrefixSetsMut},
Nibbles,
};
use alloy_primitives::{
keccak256,
map::{hash_map, B256HashMap, B256HashSet, HashMap, HashSet},
Address, B256, U256,
};
use itertools::Itertools;
use rayon::prelude::{IntoParallelIterator, ParallelIterator};
use reth_primitives::Account;
use reth_trie_common::KeyHasher;
use revm::db::{states::CacheAccount, AccountStatus, BundleAccount};
use std::borrow::Cow;
#[derive(PartialEq, Eq, Clone, Default, Debug)]
pub struct HashedPostState {
pub accounts: B256HashMap<Option<Account>>,
pub storages: B256HashMap<HashedStorage>,
}
impl HashedPostState {
pub fn from_bundle_state<'a, KH: KeyHasher>(
state: impl IntoParallelIterator<Item = (&'a Address, &'a BundleAccount)>,
) -> Self {
let hashed = state
.into_par_iter()
.map(|(address, account)| {
let hashed_address = KH::hash_key(address);
let hashed_account = account.info.as_ref().map(Into::into);
let hashed_storage = HashedStorage::from_plain_storage(
account.status,
account.storage.iter().map(|(slot, value)| (slot, &value.present_value)),
);
(hashed_address, (hashed_account, hashed_storage))
})
.collect::<Vec<(B256, (Option<Account>, HashedStorage))>>();
let mut accounts = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
let mut storages = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
for (address, (account, storage)) in hashed {
accounts.insert(address, account);
storages.insert(address, storage);
}
Self { accounts, storages }
}
pub fn from_cache_state<'a, KH: KeyHasher>(
state: impl IntoParallelIterator<Item = (&'a Address, &'a CacheAccount)>,
) -> Self {
let hashed = state
.into_par_iter()
.map(|(address, account)| {
let hashed_address = KH::hash_key(address);
let hashed_account = account.account.as_ref().map(|a| (&a.info).into());
let hashed_storage = HashedStorage::from_plain_storage(
account.status,
account.account.as_ref().map(|a| a.storage.iter()).into_iter().flatten(),
);
(hashed_address, (hashed_account, hashed_storage))
})
.collect::<Vec<(B256, (Option<Account>, HashedStorage))>>();
let mut accounts = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
let mut storages = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
for (address, (account, storage)) in hashed {
accounts.insert(address, account);
storages.insert(address, storage);
}
Self { accounts, storages }
}
pub fn from_hashed_storage(hashed_address: B256, storage: HashedStorage) -> Self {
Self {
accounts: HashMap::default(),
storages: HashMap::from_iter([(hashed_address, storage)]),
}
}
pub fn with_accounts(
mut self,
accounts: impl IntoIterator<Item = (B256, Option<Account>)>,
) -> Self {
self.accounts = HashMap::from_iter(accounts);
self
}
pub fn with_storages(
mut self,
storages: impl IntoIterator<Item = (B256, HashedStorage)>,
) -> Self {
self.storages = HashMap::from_iter(storages);
self
}
pub fn is_empty(&self) -> bool {
self.accounts.is_empty() && self.storages.is_empty()
}
pub fn construct_prefix_sets(&self) -> TriePrefixSetsMut {
let mut account_prefix_set = PrefixSetMut::with_capacity(self.accounts.len());
let mut destroyed_accounts = HashSet::default();
for (hashed_address, account) in &self.accounts {
account_prefix_set.insert(Nibbles::unpack(hashed_address));
if account.is_none() {
destroyed_accounts.insert(*hashed_address);
}
}
let mut storage_prefix_sets =
HashMap::with_capacity_and_hasher(self.storages.len(), Default::default());
for (hashed_address, hashed_storage) in &self.storages {
account_prefix_set.insert(Nibbles::unpack(hashed_address));
storage_prefix_sets.insert(*hashed_address, hashed_storage.construct_prefix_set());
}
TriePrefixSetsMut { account_prefix_set, storage_prefix_sets, destroyed_accounts }
}
pub fn extend(&mut self, other: Self) {
self.extend_inner(Cow::Owned(other));
}
pub fn extend_ref(&mut self, other: &Self) {
self.extend_inner(Cow::Borrowed(other));
}
fn extend_inner(&mut self, other: Cow<'_, Self>) {
self.accounts.extend(other.accounts.iter().map(|(&k, &v)| (k, v)));
self.storages.reserve(other.storages.len());
match other {
Cow::Borrowed(other) => {
self.extend_storages(other.storages.iter().map(|(k, v)| (*k, Cow::Borrowed(v))))
}
Cow::Owned(other) => {
self.extend_storages(other.storages.into_iter().map(|(k, v)| (k, Cow::Owned(v))))
}
}
}
fn extend_storages<'a>(
&mut self,
storages: impl IntoIterator<Item = (B256, Cow<'a, HashedStorage>)>,
) {
for (hashed_address, storage) in storages {
match self.storages.entry(hashed_address) {
hash_map::Entry::Vacant(entry) => {
entry.insert(storage.into_owned());
}
hash_map::Entry::Occupied(mut entry) => {
entry.get_mut().extend(&storage);
}
}
}
}
pub fn into_sorted(self) -> HashedPostStateSorted {
let mut updated_accounts = Vec::new();
let mut destroyed_accounts = HashSet::default();
for (hashed_address, info) in self.accounts {
if let Some(info) = info {
updated_accounts.push((hashed_address, info));
} else {
destroyed_accounts.insert(hashed_address);
}
}
updated_accounts.sort_unstable_by_key(|(address, _)| *address);
let accounts = HashedAccountsSorted { accounts: updated_accounts, destroyed_accounts };
let storages = self
.storages
.into_iter()
.map(|(hashed_address, storage)| (hashed_address, storage.into_sorted()))
.collect();
HashedPostStateSorted { accounts, storages }
}
}
#[derive(PartialEq, Eq, Clone, Debug, Default)]
pub struct HashedStorage {
pub wiped: bool,
pub storage: B256HashMap<U256>,
}
impl HashedStorage {
pub fn new(wiped: bool) -> Self {
Self { wiped, storage: HashMap::default() }
}
pub fn from_iter(wiped: bool, iter: impl IntoIterator<Item = (B256, U256)>) -> Self {
Self { wiped, storage: HashMap::from_iter(iter) }
}
pub fn from_plain_storage<'a>(
status: AccountStatus,
storage: impl IntoIterator<Item = (&'a U256, &'a U256)>,
) -> Self {
Self::from_iter(
status.was_destroyed(),
storage.into_iter().map(|(key, value)| (keccak256(B256::from(*key)), *value)),
)
}
pub fn construct_prefix_set(&self) -> PrefixSetMut {
if self.wiped {
PrefixSetMut::all()
} else {
let mut prefix_set = PrefixSetMut::with_capacity(self.storage.len());
for hashed_slot in self.storage.keys() {
prefix_set.insert(Nibbles::unpack(hashed_slot));
}
prefix_set
}
}
pub fn extend(&mut self, other: &Self) {
if other.wiped {
self.wiped = true;
self.storage.clear();
}
self.storage.extend(other.storage.iter().map(|(&k, &v)| (k, v)));
}
pub fn into_sorted(self) -> HashedStorageSorted {
let mut non_zero_valued_slots = Vec::new();
let mut zero_valued_slots = HashSet::default();
for (hashed_slot, value) in self.storage {
if value.is_zero() {
zero_valued_slots.insert(hashed_slot);
} else {
non_zero_valued_slots.push((hashed_slot, value));
}
}
non_zero_valued_slots.sort_unstable_by_key(|(key, _)| *key);
HashedStorageSorted { non_zero_valued_slots, zero_valued_slots, wiped: self.wiped }
}
}
#[derive(PartialEq, Eq, Clone, Default, Debug)]
pub struct HashedPostStateSorted {
pub(crate) accounts: HashedAccountsSorted,
pub(crate) storages: B256HashMap<HashedStorageSorted>,
}
impl HashedPostStateSorted {
pub const fn new(
accounts: HashedAccountsSorted,
storages: B256HashMap<HashedStorageSorted>,
) -> Self {
Self { accounts, storages }
}
pub const fn accounts(&self) -> &HashedAccountsSorted {
&self.accounts
}
pub const fn account_storages(&self) -> &B256HashMap<HashedStorageSorted> {
&self.storages
}
}
#[derive(Clone, Eq, PartialEq, Default, Debug)]
pub struct HashedAccountsSorted {
pub(crate) accounts: Vec<(B256, Account)>,
pub(crate) destroyed_accounts: B256HashSet,
}
impl HashedAccountsSorted {
pub fn accounts_sorted(&self) -> impl Iterator<Item = (B256, Option<Account>)> {
self.accounts
.iter()
.map(|(address, account)| (*address, Some(*account)))
.chain(self.destroyed_accounts.iter().map(|address| (*address, None)))
.sorted_by_key(|entry| *entry.0)
}
}
#[derive(Clone, Eq, PartialEq, Debug)]
pub struct HashedStorageSorted {
pub(crate) non_zero_valued_slots: Vec<(B256, U256)>,
pub(crate) zero_valued_slots: B256HashSet,
pub(crate) wiped: bool,
}
impl HashedStorageSorted {
pub const fn is_wiped(&self) -> bool {
self.wiped
}
pub fn storage_slots_sorted(&self) -> impl Iterator<Item = (B256, U256)> {
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)
}
}
#[cfg(test)]
mod tests {
use super::*;
use alloy_primitives::Bytes;
use reth_trie_common::KeccakKeyHasher;
use revm::{
db::{
states::{plain_account::PlainStorage, StorageSlot},
PlainAccount, StorageWithOriginalValues,
},
primitives::{AccountInfo, Bytecode},
};
#[test]
fn hashed_state_wiped_extension() {
let hashed_address = B256::default();
let hashed_slot = B256::with_last_byte(64);
let hashed_slot2 = B256::with_last_byte(65);
let original_slot_value = U256::from(123);
let mut hashed_state = HashedPostState::default().with_storages([(
hashed_address,
HashedStorage::from_iter(
false,
[(hashed_slot, original_slot_value), (hashed_slot2, original_slot_value)],
),
)]);
let updated_slot_value = U256::from(321);
let extension = HashedPostState::default().with_storages([(
hashed_address,
HashedStorage::from_iter(false, [(hashed_slot, updated_slot_value)]),
)]);
hashed_state.extend(extension);
let account_storage = hashed_state.storages.get(&hashed_address);
assert_eq!(
account_storage.and_then(|st| st.storage.get(&hashed_slot)),
Some(&updated_slot_value)
);
assert_eq!(
account_storage.and_then(|st| st.storage.get(&hashed_slot2)),
Some(&original_slot_value)
);
assert_eq!(account_storage.map(|st| st.wiped), Some(false));
let wiped_extension =
HashedPostState::default().with_storages([(hashed_address, HashedStorage::new(true))]);
hashed_state.extend(wiped_extension);
let account_storage = hashed_state.storages.get(&hashed_address);
assert_eq!(account_storage.map(|st| st.storage.is_empty()), Some(true));
assert_eq!(account_storage.map(|st| st.wiped), Some(true));
hashed_state.extend(HashedPostState::default().with_storages([(
hashed_address,
HashedStorage::from_iter(false, [(hashed_slot, original_slot_value)]),
)]));
let account_storage = hashed_state.storages.get(&hashed_address);
assert_eq!(
account_storage.and_then(|st| st.storage.get(&hashed_slot)),
Some(&original_slot_value)
);
assert_eq!(account_storage.and_then(|st| st.storage.get(&hashed_slot2)), None);
assert_eq!(account_storage.map(|st| st.wiped), Some(true));
hashed_state.extend(HashedPostState::default().with_storages([(
hashed_address,
HashedStorage::from_iter(false, [(hashed_slot2, updated_slot_value)]),
)]));
let account_storage = hashed_state.storages.get(&hashed_address);
assert_eq!(
account_storage.and_then(|st| st.storage.get(&hashed_slot)),
Some(&original_slot_value)
);
assert_eq!(
account_storage.and_then(|st| st.storage.get(&hashed_slot2)),
Some(&updated_slot_value)
);
assert_eq!(account_storage.map(|st| st.wiped), Some(true));
}
#[test]
fn test_hashed_post_state_from_bundle_state() {
let address = Address::random();
let account_info = AccountInfo {
balance: U256::from(123),
nonce: 42,
code_hash: B256::random(),
code: Some(Bytecode::LegacyRaw(Bytes::from(vec![1, 2]))),
};
let mut storage = StorageWithOriginalValues::default();
storage.insert(
U256::from(1),
StorageSlot { present_value: U256::from(4), ..Default::default() },
);
let account = BundleAccount {
status: AccountStatus::Changed,
info: Some(account_info.clone()),
storage,
original_info: None,
};
let state = vec![(&address, &account)];
let hashed_state = HashedPostState::from_bundle_state::<KeccakKeyHasher>(state);
assert_eq!(hashed_state.accounts.len(), 1);
assert_eq!(hashed_state.storages.len(), 1);
assert_eq!(
*hashed_state.accounts.get(&keccak256(address)).unwrap(),
Some(account_info.into())
);
}
#[test]
fn test_hashed_post_state_from_cache_state() {
let address = Address::random();
let account_info = AccountInfo {
balance: U256::from(500),
nonce: 5,
code_hash: B256::random(),
code: None,
};
let mut storage = PlainStorage::default();
storage.insert(U256::from(1), U256::from(35636));
let account = CacheAccount {
account: Some(PlainAccount { info: account_info.clone(), storage }),
status: AccountStatus::Changed,
};
let state = vec![(&address, &account)];
let hashed_state = HashedPostState::from_cache_state::<KeccakKeyHasher>(state);
assert_eq!(hashed_state.accounts.len(), 1);
assert_eq!(hashed_state.storages.len(), 1);
assert_eq!(
*hashed_state.accounts.get(&keccak256(address)).unwrap(),
Some(account_info.into())
);
}
#[test]
fn test_hashed_post_state_with_accounts() {
let address_1 = Address::random();
let address_2 = Address::random();
let account_info_1 = AccountInfo {
balance: U256::from(1000),
nonce: 1,
code_hash: B256::random(),
code: None,
};
let account_1 = (keccak256(address_1), Some(account_info_1.into()));
let account_2 = (keccak256(address_2), None);
let hashed_state = HashedPostState::default().with_accounts(vec![account_1, account_2]);
assert_eq!(hashed_state.accounts.len(), 2);
assert!(hashed_state.accounts.contains_key(&keccak256(address_1)));
assert!(hashed_state.accounts.contains_key(&keccak256(address_2)));
}
#[test]
fn test_hashed_post_state_with_storages() {
let address_1 = Address::random();
let address_2 = Address::random();
let storage_1 = (keccak256(address_1), HashedStorage::new(false));
let storage_2 = (keccak256(address_2), HashedStorage::new(true));
let hashed_state = HashedPostState::default().with_storages(vec![storage_1, storage_2]);
assert_eq!(hashed_state.storages.len(), 2);
assert!(hashed_state.storages.contains_key(&keccak256(address_1)));
assert!(hashed_state.storages.contains_key(&keccak256(address_2)));
}
#[test]
fn test_hashed_post_state_is_empty() {
let empty_state = HashedPostState::default();
assert!(empty_state.is_empty());
let non_empty_state = HashedPostState::default()
.with_accounts(vec![(keccak256(Address::random()), Some(Account::default()))]);
assert!(!non_empty_state.is_empty());
}
}