anvil/eth/backend/mem/

fork_db.rs

use crate::eth::backend::db::{
    Db, MaybeForkedDatabase, MaybeFullDatabase, SerializableAccountRecord, SerializableBlock,
    SerializableHistoricalStates, SerializableState, SerializableTransaction, StateDb,
    cache_block_hash,
};
use alloy_network::Network;
use alloy_primitives::{Address, B256, U256, map::AddressMap};
use alloy_rpc_types::BlockId;
use foundry_evm::{
    backend::{
        BlockchainDb, DatabaseResult, RevertStateSnapshotAction, SharedBackend, StateSnapshot,
    },
    fork::database::ForkDbStateSnapshot,
};
use revm::{
    context::BlockEnv,
    database::{Database, DbAccount},
    state::AccountInfo,
};

pub use foundry_evm::fork::database::ForkedDatabase;

impl<N: Network> MaybeFullDatabase for SharedBackend<N> {
    fn clear_into_state_snapshot(&mut self) -> StateSnapshot {
        StateSnapshot::default()
    }

    fn read_as_state_snapshot(&self) -> StateSnapshot {
        StateSnapshot::default()
    }

    fn clear(&mut self) {}

    fn init_from_state_snapshot(&mut self, _state_snapshot: StateSnapshot) {}
}

impl<N: Network> Db for ForkedDatabase<N> {
    fn insert_account(&mut self, address: Address, account: AccountInfo) {
        self.database_mut().insert_account(address, account)
    }

    fn set_storage_at(&mut self, address: Address, slot: B256, val: B256) -> DatabaseResult<()> {
        // this ensures the account is loaded first
        let _ = Database::basic(self, address)?;
        self.database_mut().set_storage_at(address, slot, val)
    }

    fn insert_block_hash(&mut self, number: U256, hash: B256) {
        cache_block_hash(&mut self.inner().block_hashes().write(), number, hash);
    }

    fn set_block_hashes(&mut self, block_hashes: Vec<(U256, B256)>) {
        *self.inner().block_hashes().write() = block_hashes.into_iter().collect();
    }

    fn dump_state(
        &self,
        at: BlockEnv,
        best_number: u64,
        blocks: Vec<SerializableBlock>,
        transactions: Vec<SerializableTransaction>,
        historical_states: Option<SerializableHistoricalStates>,
    ) -> DatabaseResult<Option<SerializableState>> {
        let mut db = self.database().clone();
        let accounts = self
            .database()
            .cache
            .accounts
            .clone()
            .into_iter()
            .map(|(k, v)| -> DatabaseResult<_> {
                let code = if let Some(code) = v.info.code {
                    code
                } else {
                    db.code_by_hash(v.info.code_hash)?
                };
                Ok((
                    k,
                    SerializableAccountRecord {
                        nonce: v.info.nonce,
                        balance: v.info.balance,
                        code: code.original_bytes(),
                        storage: v.storage.into_iter().map(|(k, v)| (k.into(), v.into())).collect(),
                    },
                ))
            })
            .collect::<Result<_, _>>()?;
        Ok(Some(SerializableState {
            block: Some(at),
            accounts,
            best_block_number: Some(best_number),
            blocks,
            transactions,
            historical_states,
        }))
    }

    fn snapshot_state(&mut self) -> U256 {
        self.insert_state_snapshot()
    }

    fn revert_state(&mut self, id: U256, action: RevertStateSnapshotAction) -> bool {
        self.revert_state_snapshot(id, action)
    }

    fn current_state(&self) -> StateDb {
        StateDb::new(self.create_state_snapshot())
    }
}

impl<N: Network> MaybeFullDatabase for ForkedDatabase<N> {
    fn maybe_as_full_db(&self) -> Option<&AddressMap<DbAccount>> {
        Some(&self.database().cache.accounts)
    }

    fn maybe_full_db(&self) -> Option<AddressMap<DbAccount>> {
        None
    }

    fn clear_into_state_snapshot(&mut self) -> StateSnapshot {
        let db = self.inner().db();
        let accounts = std::mem::take(&mut *db.accounts.write());
        let storage = std::mem::take(&mut *db.storage.write());
        let block_hashes = std::mem::take(&mut *db.block_hashes.write());
        StateSnapshot { accounts, storage, block_hashes }
    }

    fn read_as_state_snapshot(&self) -> StateSnapshot {
        let db = self.inner().db();
        let accounts = db.accounts.read().clone();
        let storage = db.storage.read().clone();
        let block_hashes = db.block_hashes.read().clone();
        StateSnapshot { accounts, storage, block_hashes }
    }

    fn clear(&mut self) {
        self.flush_cache();
        self.clear_into_state_snapshot();
    }

    fn init_from_state_snapshot(&mut self, state_snapshot: StateSnapshot) {
        let db = self.inner().db();
        let StateSnapshot { accounts, storage, block_hashes } = state_snapshot;
        *db.accounts.write() = accounts;
        *db.storage.write() = storage;
        *db.block_hashes.write() = block_hashes;
    }
}

impl<N: Network> MaybeFullDatabase for ForkDbStateSnapshot<N> {
    fn maybe_as_full_db(&self) -> Option<&AddressMap<DbAccount>> {
        Some(&self.local.cache.accounts)
    }

    fn maybe_full_db(&self) -> Option<AddressMap<DbAccount>> {
        None
    }

    fn clear_into_state_snapshot(&mut self) -> StateSnapshot {
        let mut state_snapshot = std::mem::take(&mut self.state_snapshot);
        let local_state_snapshot = self.local.clear_into_state_snapshot();
        state_snapshot.accounts.extend(local_state_snapshot.accounts);
        state_snapshot.storage.extend(local_state_snapshot.storage);
        state_snapshot.block_hashes.extend(local_state_snapshot.block_hashes);
        state_snapshot
    }

    fn read_as_state_snapshot(&self) -> StateSnapshot {
        let mut state_snapshot = self.state_snapshot.clone();
        let local_state_snapshot = self.local.read_as_state_snapshot();
        state_snapshot.accounts.extend(local_state_snapshot.accounts);
        state_snapshot.storage.extend(local_state_snapshot.storage);
        state_snapshot.block_hashes.extend(local_state_snapshot.block_hashes);
        state_snapshot
    }

    fn clear(&mut self) {
        std::mem::take(&mut self.state_snapshot);
        self.local.clear()
    }

    fn init_from_state_snapshot(&mut self, state_snapshot: StateSnapshot) {
        self.state_snapshot = state_snapshot;
    }
}

impl<N: Network> MaybeForkedDatabase for ForkedDatabase<N> {
    fn maybe_reset(&mut self, urls: Vec<String>, block_number: BlockId) -> Result<(), String> {
        self.reset(urls, block_number)
    }

    fn maybe_flush_cache(&self) -> Result<(), String> {
        self.flush_cache();
        Ok(())
    }

    fn maybe_inner(&self) -> Result<&BlockchainDb, String> {
        Ok(self.inner())
    }
}