foundry_evm_core/fork/

database.rs

//! A revm database that forks off a remote client

use crate::{
    backend::{RevertStateSnapshotAction, StateSnapshot},
    state_snapshot::StateSnapshots,
};
use alloy_network::Network;
use alloy_primitives::{Address, B256, U256};
use alloy_rpc_types::BlockId;
use foundry_fork_db::{BlockchainDb, DatabaseError, ForkBlockEnv, SharedBackend};
use parking_lot::Mutex;
use revm::{
    Database, DatabaseCommit,
    bytecode::Bytecode,
    context::BlockEnv,
    database::{AccountState, CacheDB, DatabaseRef},
    primitives::AddressMap,
    state::{Account, AccountInfo},
};
use std::sync::Arc;

/// a [revm::Database] that's forked off another client
///
/// The `backend` is used to retrieve (missing) data, which is then fetched from the remote
/// endpoint. The inner in-memory database holds this storage and will be used for write operations.
/// This database uses the `backend` for read and the `db` for write operations. But note the
/// `backend` will also write (missing) data to the `db` in the background
#[derive(Clone, Debug)]
pub struct ForkedDatabase<N: Network, B: ForkBlockEnv = BlockEnv> {
    /// Responsible for fetching missing data.
    ///
    /// This is responsible for getting data.
    backend: SharedBackend<N, B>,
    /// Cached Database layer, ensures that changes are not written to the database that
    /// exclusively stores the state of the remote client.
    ///
    /// This separates Read/Write operations
    ///   - reads from the `SharedBackend as DatabaseRef` writes to the internal cache storage.
    cache_db: CacheDB<SharedBackend<N, B>>,
    /// Contains all the data already fetched.
    ///
    /// This exclusively stores the _unchanged_ remote client state.
    db: BlockchainDb<B>,
    /// Holds the state snapshots of a blockchain.
    state_snapshots: Arc<Mutex<StateSnapshots<ForkDbStateSnapshot<N, B>>>>,
}

impl<N: Network, B: ForkBlockEnv> ForkedDatabase<N, B> {
    /// Creates a new instance of this DB
    pub fn new(backend: SharedBackend<N, B>, db: BlockchainDb<B>) -> Self {
        Self {
            cache_db: CacheDB::new(backend.clone()),
            backend,
            db,
            state_snapshots: Arc::new(Mutex::new(Default::default())),
        }
    }

    pub const fn database(&self) -> &CacheDB<SharedBackend<N, B>> {
        &self.cache_db
    }

    pub const fn database_mut(&mut self) -> &mut CacheDB<SharedBackend<N, B>> {
        &mut self.cache_db
    }

    pub const fn state_snapshots(&self) -> &Arc<Mutex<StateSnapshots<ForkDbStateSnapshot<N, B>>>> {
        &self.state_snapshots
    }

    /// Reset the fork to a fresh forked state, and optionally update the fork config
    pub fn reset(
        &mut self,
        _urls: Vec<String>,
        block_number: impl Into<BlockId>,
    ) -> Result<(), String> {
        self.backend.set_pinned_block(block_number).map_err(|err| err.to_string())?;

        // TODO need to find a way to update generic provider via url

        // wipe the storage retrieved from remote
        self.inner().db().clear();
        // create a fresh `CacheDB`, effectively wiping modified state
        self.cache_db = CacheDB::new(self.backend.clone());
        trace!(target: "backend::forkdb", "Cleared database");
        Ok(())
    }

    /// Flushes the cache to disk if configured
    pub fn flush_cache(&self) {
        self.db.cache().flush()
    }

    /// Returns the database that holds the remote state
    pub const fn inner(&self) -> &BlockchainDb<B> {
        &self.db
    }

    pub fn create_state_snapshot(&self) -> ForkDbStateSnapshot<N, B> {
        let db = self.db.db();
        let state_snapshot = StateSnapshot {
            accounts: db.accounts.read().clone(),
            storage: db.storage.read().clone(),
            block_hashes: db.block_hashes.read().clone(),
            account_states: None,
        };
        ForkDbStateSnapshot { local: self.cache_db.clone(), state_snapshot }
    }

    pub fn insert_state_snapshot(&self) -> U256 {
        let state_snapshot = self.create_state_snapshot();
        let mut state_snapshots = self.state_snapshots().lock();
        let id = state_snapshots.insert(state_snapshot);
        trace!(target: "backend::forkdb", "Created new snapshot {}", id);
        id
    }

    /// Removes the snapshot from the tracked snapshot and sets it as the current state
    pub fn revert_state_snapshot(&mut self, id: U256, action: RevertStateSnapshotAction) -> bool {
        let state_snapshot = { self.state_snapshots().lock().remove(id) };
        if let Some(state_snapshot) = state_snapshot {
            if action.is_keep() {
                self.state_snapshots().lock().insert_at(state_snapshot.clone(), id);
            }
            let ForkDbStateSnapshot {
                local,
                state_snapshot: StateSnapshot { accounts, storage, block_hashes, .. },
            } = state_snapshot;
            let db = self.inner().db();
            {
                let mut accounts_lock = db.accounts.write();
                accounts_lock.clear();
                accounts_lock.extend(accounts);
            }
            {
                let mut storage_lock = db.storage.write();
                storage_lock.clear();
                storage_lock.extend(storage);
            }
            {
                let mut block_hashes_lock = db.block_hashes.write();
                block_hashes_lock.clear();
                block_hashes_lock.extend(block_hashes);
            }

            self.cache_db = local;

            trace!(target: "backend::forkdb", "Reverted snapshot {}", id);
            true
        } else {
            warn!(target: "backend::forkdb", "No snapshot to revert for {}", id);
            false
        }
    }
}

impl<N: Network, B: ForkBlockEnv> Database for ForkedDatabase<N, B> {
    type Error = DatabaseError;

    fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
        // Note: this will always return Some, since the `SharedBackend` will always load the
        // account, this differs from `<CacheDB as Database>::basic`, See also
        // [MemDb::ensure_loaded](crate::backend::MemDb::ensure_loaded)
        Database::basic(&mut self.cache_db, address)
    }

    fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
        Database::code_by_hash(&mut self.cache_db, code_hash)
    }

    fn storage(&mut self, address: Address, index: U256) -> Result<U256, Self::Error> {
        Database::storage(&mut self.cache_db, address, index)
    }

    fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
        Database::block_hash(&mut self.cache_db, number)
    }
}

impl<N: Network, B: ForkBlockEnv> DatabaseRef for ForkedDatabase<N, B> {
    type Error = DatabaseError;

    fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
        self.cache_db.basic_ref(address)
    }

    fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
        self.cache_db.code_by_hash_ref(code_hash)
    }

    fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
        DatabaseRef::storage_ref(&self.cache_db, address, index)
    }

    fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
        self.cache_db.block_hash_ref(number)
    }
}

impl<N: Network, B: ForkBlockEnv> DatabaseCommit for ForkedDatabase<N, B> {
    fn commit(&mut self, changes: AddressMap<Account>) {
        self.database_mut().commit(changes)
    }
}

/// Represents a snapshot of the database
///
/// This mimics `revm::CacheDB`
#[derive(Clone, Debug)]
pub struct ForkDbStateSnapshot<N: Network, B: ForkBlockEnv = BlockEnv> {
    pub local: CacheDB<SharedBackend<N, B>>,
    pub state_snapshot: StateSnapshot,
}

impl<N: Network, B: ForkBlockEnv> ForkDbStateSnapshot<N, B> {
    /// Lookup storage in `state_snapshot`, then fall back to the backend (remote RPC).
    fn storage_from_snapshot_or_backend(
        &self,
        address: Address,
        index: U256,
    ) -> Result<U256, DatabaseError> {
        // Check state_snapshot.storage first (data fetched by SharedBackend / disk cache).
        if let Some(val) = self.state_snapshot.storage.get(&address).and_then(|s| s.get(&index)) {
            return Ok(*val);
        }
        // Fall back to the underlying backend (SharedBackend → remote RPC).
        DatabaseRef::storage_ref(&self.local, address, index)
    }
}

// This `DatabaseRef` implementation works similar to `CacheDB` which prioritizes modified elements,
// and uses another db as fallback
// We prioritize stored changed accounts/storage
impl<N: Network, B: ForkBlockEnv> DatabaseRef for ForkDbStateSnapshot<N, B> {
    type Error = DatabaseError;

    fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
        match self.local.cache.accounts.get(&address) {
            Some(account) => Ok(Some(account.info.clone())),
            None => {
                let mut acc = self.state_snapshot.accounts.get(&address).cloned();

                if acc.is_none() {
                    acc = self.local.basic_ref(address)?;
                }
                Ok(acc)
            }
        }
    }

    fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
        self.local.code_by_hash_ref(code_hash)
    }

    fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
        match self.local.cache.accounts.get(&address) {
            Some(account) => match account.storage.get(&index) {
                Some(entry) => Ok(*entry),
                None if matches!(
                    account.account_state,
                    AccountState::NotExisting | AccountState::StorageCleared
                ) =>
                {
                    Ok(U256::ZERO)
                }
                None => self.storage_from_snapshot_or_backend(address, index),
            },
            None => self.storage_from_snapshot_or_backend(address, index),
        }
    }

    fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
        match self.state_snapshot.block_hashes.get(&U256::from(number)).copied() {
            None => self.local.block_hash_ref(number),
            Some(block_hash) => Ok(block_hash),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::backend::BlockchainDbMeta;
    use alloy_network::AnyNetwork;
    use alloy_provider::{ProviderBuilder, mock::Asserter};
    use foundry_common::provider::get_http_provider;
    use revm::database::DbAccount;

    fn fork_snapshot_with_closed_backend() -> ForkDbStateSnapshot<AnyNetwork> {
        let provider =
            ProviderBuilder::<_, _, AnyNetwork>::default().connect_mocked_client(Asserter::new());
        let db = BlockchainDb::new(
            BlockchainDbMeta::new(BlockEnv::default(), "http://localhost".to_string()),
            None,
        );
        let (backend, handler) = SharedBackend::new(provider, db, None);
        drop(handler);
        ForkDbStateSnapshot {
            local: CacheDB::new(backend),
            state_snapshot: StateSnapshot::default(),
        }
    }

    #[test]
    fn fork_db_state_snapshot_respects_local_storage_clearing() {
        let address = Address::random();
        let slot = U256::from(42);
        for account_state in [AccountState::NotExisting, AccountState::StorageCleared] {
            let mut snapshot = fork_snapshot_with_closed_backend();
            snapshot.state_snapshot.storage.entry(address).or_default().insert(slot, U256::from(7));
            snapshot.local.cache.accounts.insert(
                address,
                DbAccount { account_state: account_state.clone(), ..Default::default() },
            );

            assert_eq!(snapshot.storage_ref(address, slot).unwrap(), U256::ZERO);
            assert_eq!(snapshot.storage_ref(address, slot + U256::ONE).unwrap(), U256::ZERO);

            if account_state.is_storage_cleared() {
                snapshot
                    .local
                    .cache
                    .accounts
                    .get_mut(&address)
                    .unwrap()
                    .storage
                    .insert(slot, U256::from(9));
                assert_eq!(snapshot.storage_ref(address, slot).unwrap(), U256::from(9));
            }
        }
    }

    /// Demonstrates that `Database::basic` for `ForkedDatabase` will always return the
    /// `AccountInfo`
    #[tokio::test(flavor = "multi_thread")]
    async fn fork_db_insert_basic_default() {
        let rpc = foundry_test_utils::rpc::next_http_rpc_endpoint();
        let provider = get_http_provider(rpc.clone());
        let meta = BlockchainDbMeta::new(BlockEnv::default(), rpc);

        let db = BlockchainDb::new(meta, None);

        let backend = SharedBackend::spawn_backend(Arc::new(provider), db.clone(), None).await;

        let mut db = ForkedDatabase::new(backend, db);
        let address = Address::random();

        let info = Database::basic(&mut db, address).unwrap();
        assert!(info.is_some());
        let mut info = info.unwrap();
        info.balance = U256::from(500u64);

        // insert the modified account info
        db.database_mut().insert_account_info(address, info.clone());

        let loaded = Database::basic(&mut db, address).unwrap();
        assert!(loaded.is_some());
        assert_eq!(loaded.unwrap(), info);
    }

    /// Reads cached remote storage without consulting the backend, including when the local
    /// account is present but does not override the slot.
    #[test]
    fn fork_db_state_snapshot_reads_storage_from_snapshot() {
        let mut snapshot = fork_snapshot_with_closed_backend();
        let address = Address::random();
        let slot = U256::from(42u64);
        let expected = U256::from(0xdeadbeefu64);

        snapshot.state_snapshot.storage.entry(address).or_default().insert(slot, expected);

        assert_eq!(snapshot.storage_ref(address, slot).unwrap(), expected);
        snapshot.local.cache.accounts.insert(address, DbAccount::default());
        assert_eq!(snapshot.storage_ref(address, slot).unwrap(), expected);
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn fork_db_revert_invalidates_newer_snapshots() {
        let rpc = foundry_test_utils::rpc::next_http_rpc_endpoint();
        let provider = get_http_provider(rpc.clone());
        let meta = BlockchainDbMeta::new(BlockEnv::default(), rpc);
        let db = BlockchainDb::new(meta, None);
        let backend = SharedBackend::spawn_backend(Arc::new(provider), db.clone(), None).await;
        let mut db = ForkedDatabase::new(backend, db);

        let first = db.insert_state_snapshot();
        let second = db.insert_state_snapshot();

        assert!(db.revert_state_snapshot(first, RevertStateSnapshotAction::RevertRemove));
        assert!(!db.revert_state_snapshot(second, RevertStateSnapshotAction::RevertRemove));
    }
}