use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use alloy_primitives::{Address, B256, U256};
use revm::primitives::hardfork::SpecId;
use revm::state::{AccountInfo, Bytecode};
use crate::access_set::StorageAccessList;
pub(crate) struct BaseState {
pub(crate) accounts: HashMap<Address, AccountInfo>,
pub(crate) storage: HashMap<Address, Arc<HashMap<U256, U256>>>,
pub(crate) code_by_hash: HashMap<B256, Bytecode>,
}
pub struct EvmSnapshot {
pub(crate) base: Arc<BaseState>,
pub(crate) overlay_accounts: HashMap<Address, AccountInfo>,
pub(crate) overlay_storage: HashMap<Address, HashMap<U256, U256>>,
pub(crate) overlay_code_by_hash: HashMap<B256, Bytecode>,
pub(crate) storage_cleared: HashSet<Address>,
pub(crate) accounts_not_existing: HashSet<Address>,
pub(crate) block_hashes: HashMap<u64, B256>,
pub(crate) block_number: Option<u64>,
pub(crate) basefee: Option<u64>,
pub(crate) coinbase: Option<Address>,
pub(crate) prevrandao: Option<B256>,
pub(crate) gas_limit: Option<u64>,
pub(crate) chain_id: u64,
pub(crate) timestamp: Option<u64>,
pub(crate) spec_id: SpecId,
pub(crate) shared_memory_capacity: usize,
}
impl EvmSnapshot {
pub const fn chain_id(&self) -> u64 {
self.chain_id
}
pub const fn block_number(&self) -> Option<u64> {
self.block_number
}
pub const fn basefee(&self) -> Option<u64> {
self.basefee
}
pub const fn coinbase(&self) -> Option<Address> {
self.coinbase
}
pub const fn prevrandao(&self) -> Option<B256> {
self.prevrandao
}
pub const fn gas_limit(&self) -> Option<u64> {
self.gas_limit
}
pub const fn timestamp(&self) -> Option<u64> {
self.timestamp
}
pub fn resident_read_set(&self) -> StorageAccessList {
let mut resident = StorageAccessList::default();
resident.accounts.extend(self.base.accounts.keys().copied());
resident
.accounts
.extend(self.overlay_accounts.keys().copied());
resident
.accounts
.extend(self.accounts_not_existing.iter().copied());
resident
.code_hashes
.extend(self.base.code_by_hash.keys().copied());
resident
.code_hashes
.extend(self.overlay_code_by_hash.keys().copied());
for (address, slots) in &self.base.storage {
resident
.slots
.extend(slots.keys().copied().map(|slot| (*address, slot)));
}
for (address, slots) in &self.overlay_storage {
resident
.slots
.extend(slots.keys().copied().map(|slot| (*address, slot)));
}
resident
.block_numbers
.extend(self.block_hashes.keys().copied());
resident
}
pub fn missing_read_set(&self, required: &StorageAccessList) -> StorageAccessList {
StorageAccessList {
accounts: required
.accounts
.iter()
.copied()
.filter(|address| {
!self.accounts_not_existing.contains(address)
&& self.account_info(*address).is_none()
})
.collect(),
code_hashes: required
.code_hashes
.iter()
.copied()
.filter(|hash| self.code(*hash).is_none())
.collect(),
slots: required
.slots
.iter()
.copied()
.filter(|(address, slot)| self.storage_value(*address, *slot).is_none())
.collect(),
block_numbers: required
.block_numbers
.iter()
.copied()
.filter(|number| !self.block_hashes.contains_key(number))
.collect(),
}
}
pub(crate) fn account_info(&self, address: Address) -> Option<&AccountInfo> {
if self.accounts_not_existing.contains(&address) {
return None;
}
self.overlay_accounts
.get(&address)
.or_else(|| self.base.accounts.get(&address))
}
pub fn storage_value(&self, address: Address, slot: U256) -> Option<U256> {
if let Some(account_storage) = self.overlay_storage.get(&address) {
if let Some(value) = account_storage.get(&slot) {
return Some(*value);
}
if self.storage_cleared.contains(&address) {
return Some(U256::ZERO);
}
}
self.base
.storage
.get(&address)
.and_then(|s| s.get(&slot).copied())
}
pub(crate) fn code(&self, code_hash: B256) -> Option<&Bytecode> {
self.overlay_code_by_hash
.get(&code_hash)
.or_else(|| self.base.code_by_hash.get(&code_hash))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn empty_base() -> Arc<BaseState> {
Arc::new(BaseState {
accounts: HashMap::new(),
storage: HashMap::new(),
code_by_hash: HashMap::new(),
})
}
#[test]
fn test_snapshot_is_send_sync() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<EvmSnapshot>();
assert_send_sync::<Arc<EvmSnapshot>>();
}
#[test]
fn test_empty_snapshot() {
let snap = EvmSnapshot {
base: empty_base(),
overlay_accounts: HashMap::new(),
overlay_storage: HashMap::new(),
overlay_code_by_hash: HashMap::new(),
storage_cleared: HashSet::new(),
accounts_not_existing: HashSet::new(),
block_hashes: HashMap::new(),
block_number: Some(100),
basefee: Some(1000),
coinbase: None,
prevrandao: None,
gas_limit: None,
chain_id: 42161,
timestamp: None,
spec_id: SpecId::CANCUN,
shared_memory_capacity: 64_000,
};
assert_eq!(snap.chain_id, 42161);
assert_eq!(snap.block_number, Some(100));
}
}