refactor: store state overrides per block

crates/rethnet_evm/benches/state/database_commit.rs

@@ -59,7 +59,7 @@ fn bench_database_commit(c: &mut Criterion) {
                 code_hash: code.map_or(KECCAK_EMPTY, |code| code.hash_slow()),
             },
             storage,
-            status: AccountStatus::default(),
+            status: AccountStatus::Touched,
         };
 
         account.mark_touch();

crates/rethnet_evm/src/block/builder.rs

@@ -14,8 +14,8 @@ use rethnet_eth::{
 use revm::{
     db::DatabaseComponentError,
     primitives::{
-        Account, AccountInfo, AccountStatus, BlockEnv, CfgEnv, EVMError, ExecutionResult, HashMap,
-        InvalidTransaction, Output, ResultAndState, SpecId,
+        AccountInfo, BlockEnv, CfgEnv, EVMError, ExecutionResult, InvalidTransaction, Output,
+        ResultAndState, SpecId,
     },
 };
 
@@ -127,7 +127,7 @@ impl BlockBuilder {
             header,
             callers: Vec::new(),
             transactions: Vec::new(),
-            state_diff: StateDiff::new(),
+            state_diff: StateDiff::default(),
             receipts: Vec::new(),
             parent_gas_limit,
         })
@@ -200,16 +200,7 @@ impl BlockBuilder {
             state: state_diff,
         } = run_transaction(evm, inspector)?;
 
-        for (address, account_diff) in &state_diff {
-            self.state_diff
-                .entry(*address)
-                .and_modify(|account| {
-                    account.info = account_diff.info.clone();
-                    account.status.insert(account_diff.status);
-                    account.storage.extend(account_diff.storage.clone());
-                })
-                .or_insert(account_diff.clone());
-        }
+        self.state_diff.apply_diff(state_diff.clone());
 
         state.commit(state_diff);
 
@@ -311,16 +302,7 @@ impl BlockBuilder {
                 .basic(address)?
                 .expect("Account must exist after modification");
 
-            self.state_diff
-                .entry(address)
-                .and_modify(|account| {
-                    account.info.balance = account_info.balance;
-                })
-                .or_insert(Account {
-                    info: account_info,
-                    storage: HashMap::new(),
-                    status: AccountStatus::Touched,
-                });
+            self.state_diff.apply_account_change(address, account_info);
         }
 
         if let Some(gas_limit) = self.parent_gas_limit {

crates/rethnet_evm/src/blockchain.rs

@@ -10,10 +10,14 @@ use async_trait::async_trait;
 use rethnet_eth::{
     receipt::BlockReceipt, remote::RpcClientError, spec::HardforkActivations, B256, U256,
 };
-use revm::{db::BlockHashRef, primitives::SpecId, DatabaseCommit};
+use revm::{
+    db::BlockHashRef,
+    primitives::{HashMap, SpecId},
+    DatabaseCommit,
+};
 
 use crate::{
-    state::{StateDiff, SyncState},
+    state::{StateDiff, StateOverride, SyncState},
     Block, LocalBlock, SyncBlock,
 };
 
@@ -128,10 +132,15 @@ pub trait Blockchain {
     /// Retrieves the hardfork specification used for new blocks.
     fn spec_id(&self) -> SpecId;
 
-    /// Retrieves the state at a given block
+    /// Retrieves the state at a given block.
+    ///
+    /// The state overrides are applied after the block they are associated with.
+    /// The specified override of a nonce may be ignored to maintain validity.
     async fn state_at_block_number(
         &self,
         block_number: &U256,
+        // Block number -> state overrides
+        state_overrides: &HashMap<U256, StateOverride>,
     ) -> Result<Box<dyn SyncState<Self::StateError>>, Self::BlockchainError>;
 
     /// Retrieves the total difficulty at the block with the provided hash.
@@ -197,13 +206,35 @@ fn compute_state_at_block<BlockT: Block + Clone>(
     state: &mut dyn DatabaseCommit,
     local_storage: &ReservableSparseBlockchainStorage<BlockT>,
     block_number: &U256,
+    state_overrides: &HashMap<U256, StateOverride>,
 ) {
+    // If we're dealing with a local block, apply their state diffs
     let state_diffs = local_storage
         .state_diffs_until_block(block_number)
-        .expect("The block is validated to exist");
+        .unwrap_or_default();
+
+    // state diffs are already sorted by block number
+    let mut overriden_state_diffs: Vec<(&U256, StateDiff)> = state_diffs
+        .iter()
+        .map(|(block_number, state_diff)| (block_number, state_diff.clone()))
+        .collect();
+
+    // Override states (in sorted order)
+    for (block_number, overrides) in state_overrides {
+        let index = overriden_state_diffs
+            .binary_search_by_key(&block_number, |(block_number, _)| *block_number);
+        match index {
+            Ok(index) => overriden_state_diffs[index]
+                .1
+                .apply_diff(overrides.diff.clone().into()),
+            Err(index) => {
+                overriden_state_diffs.insert(index, (block_number, overrides.diff.clone()))
+            }
+        }
+    }
 
-    for state_diff in state_diffs {
-        state.commit(state_diff.clone());
+    for (_block_number, state_diff) in overriden_state_diffs {
+        state.commit(state_diff.into());
     }
 }
 

crates/rethnet_evm/src/blockchain/forked.rs

@@ -6,18 +6,18 @@ use parking_lot::Mutex;
 use rethnet_eth::block::LargestSafeBlockNumberArgs;
 use rethnet_eth::receipt::BlockReceipt;
 use rethnet_eth::spec::chain_hardfork_activations;
-use rethnet_eth::Address;
+
 use rethnet_eth::{
     block::{largest_safe_block_number, safe_block_depth},
     remote::{RpcClient, RpcClientError},
     spec::{chain_name, HardforkActivations},
     B256, U256,
 };
-use revm::primitives::{AccountInfo, HashMap};
+use revm::primitives::HashMap;
 use revm::{db::BlockHashRef, primitives::SpecId};
 use tokio::runtime;
 
-use crate::state::{ForkState, StateDiff, StateError, SyncState};
+use crate::state::{ForkState, StateDiff, StateError, StateOverride, SyncState};
 use crate::{Block, LocalBlock, RandomHashGenerator, SyncBlock};
 
 use super::compute_state_at_block;
@@ -58,9 +58,8 @@ pub struct ForkedBlockchain {
     local_storage: ReservableSparseBlockchainStorage<Arc<dyn SyncBlock<Error = BlockchainError>>>,
     // We can force caching here because we only fork from a safe block number.
     remote: RemoteBlockchain<Arc<dyn SyncBlock<Error = BlockchainError>>, true>,
-    // The state at the time of forking
-    fork_state: ForkState,
     runtime: runtime::Handle,
+    state_root_generator: Arc<Mutex<RandomHashGenerator>>,
     fork_block_number: U256,
     chain_id: U256,
     _network_id: U256,
@@ -77,7 +76,6 @@ impl ForkedBlockchain {
         rpc_client: RpcClient,
         fork_block_number: Option<U256>,
         state_root_generator: Arc<Mutex<RandomHashGenerator>>,
-        account_overrides: HashMap<Address, AccountInfo>,
         hardfork_activation_overrides: HashMap<U256, HardforkActivations>,
     ) -> Result<Self, CreationError> {
         let (chain_id, network_id, latest_block_number) = tokio::join!(
@@ -148,33 +146,11 @@ impl ForkedBlockchain {
         let rpc_client = Arc::new(rpc_client);
         let remote = RemoteBlockchain::new(rpc_client.clone());
 
-        let fork_state = if account_overrides.is_empty() {
-            let block: Arc<dyn SyncBlock<Error = BlockchainError>> =
-                remote.block_by_number(&fork_block_number).await?;
-
-            ForkState::new(
-                runtime.clone(),
-                rpc_client,
-                state_root_generator,
-                fork_block_number,
-                block.header().state_root,
-            )
-        } else {
-            ForkState::with_overrides(
-                runtime.clone(),
-                rpc_client,
-                state_root_generator,
-                fork_block_number,
-                account_overrides,
-            )
-            .await?
-        };
-
         Ok(Self {
             local_storage: ReservableSparseBlockchainStorage::empty(fork_block_number),
             remote,
-            fork_state,
             runtime,
+            state_root_generator,
             fork_block_number,
             chain_id,
             _network_id: network_id,
@@ -348,41 +324,36 @@ impl Blockchain for ForkedBlockchain {
     async fn state_at_block_number(
         &self,
         block_number: &U256,
+        state_overrides: &HashMap<U256, StateOverride>,
     ) -> Result<Box<dyn SyncState<Self::StateError>>, Self::BlockchainError> {
         if *block_number > self.last_block_number().await {
             return Err(BlockchainError::UnknownBlockNumber);
         }
 
-        let state = match block_number.cmp(&self.fork_block_number) {
-            std::cmp::Ordering::Less => {
-                // We don't apply account overrides to pre-fork states
-                let block = self.remote.block_by_number(block_number).await?;
-
-                ForkState::new(
-                    self.runtime.clone(),
-                    self.remote.client().clone(),
-                    self.fork_state.state_root_generator().clone(),
-                    *block_number,
-                    block.header().state_root,
-                )
-            }
-            std::cmp::Ordering::Equal => self.fork_state.clone(),
-            std::cmp::Ordering::Greater => {
-                let mut state = self.fork_state.clone();
-                compute_state_at_block(&mut state, &self.local_storage, block_number);
-
-                let state_root = self
-                    .local_storage
-                    .block_by_number(block_number)
-                    .expect("The block is validated to exist")
-                    .header()
-                    .state_root;
+        let state_root = if let Some(state_override) = state_overrides.get(block_number) {
+            state_override.state_root
+        } else {
+            self.block_by_number(block_number)
+                .await?
+                .expect("The block is validated to exist")
+                .header()
+                .state_root
+        };
 
-                state.set_state_root(state_root);
+        let mut state = ForkState::new(
+            self.runtime.clone(),
+            self.remote.client().clone(),
+            self.state_root_generator.clone(),
+            *block_number,
+            state_root,
+        );
 
-                state
-            }
-        };
+        compute_state_at_block(
+            &mut state,
+            &self.local_storage,
+            block_number,
+            state_overrides,
+        );
 
         Ok(Box::new(state))
     }

crates/rethnet_evm/src/blockchain/local.rs

@@ -7,9 +7,11 @@ use std::{
 
 use async_trait::async_trait;
 use rethnet_eth::{block::PartialHeader, trie::KECCAK_NULL_RLP, Bytes, B256, B64, U256};
+use revm::primitives::HashMap;
+use revm::DatabaseCommit;
 use revm::{db::BlockHashRef, primitives::SpecId};
 
-use crate::state::{StateDebug, StateDiff, StateError, SyncState, TrieState};
+use crate::state::{StateDebug, StateDiff, StateError, StateOverride, SyncState, TrieState};
 use crate::{Block, LocalBlock, SyncBlock};
 
 use super::compute_state_at_block;
@@ -42,7 +44,6 @@ pub enum InsertBlockError {
 #[derive(Debug)]
 pub struct LocalBlockchain {
     storage: ReservableSparseBlockchainStorage<Arc<dyn SyncBlock<Error = BlockchainError>>>,
-    genesis_state: TrieState,
     chain_id: U256,
     spec_id: SpecId,
 }
@@ -51,7 +52,7 @@ impl LocalBlockchain {
     /// Constructs a new instance using the provided arguments to build a genesis block.
     #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub fn new(
-        genesis_state: TrieState,
+        genesis_diff: StateDiff,
         chain_id: U256,
         spec_id: SpecId,
         gas_limit: U256,
@@ -61,6 +62,9 @@ impl LocalBlockchain {
     ) -> Result<Self, CreationError> {
         const EXTRA_DATA: &[u8] = b"\x12\x34";
 
+        let mut genesis_state = TrieState::default();
+        genesis_state.commit(genesis_diff.clone().into());
+
         let partial_header = PartialHeader {
             state_root: genesis_state
                 .state_root()
@@ -104,7 +108,7 @@ impl LocalBlockchain {
         Ok(unsafe {
             Self::with_genesis_block_unchecked(
                 LocalBlock::empty(partial_header, spec_id),
-                genesis_state,
+                genesis_diff,
                 chain_id,
                 spec_id,
             )
@@ -115,7 +119,7 @@ impl LocalBlockchain {
     #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub fn with_genesis_block(
         genesis_block: LocalBlock,
-        genesis_state: TrieState,
+        genesis_diff: StateDiff,
         chain_id: U256,
         spec_id: SpecId,
     ) -> Result<Self, InsertBlockError> {
@@ -133,7 +137,7 @@ impl LocalBlockchain {
         }
 
         Ok(unsafe {
-            Self::with_genesis_block_unchecked(genesis_block, genesis_state, chain_id, spec_id)
+            Self::with_genesis_block_unchecked(genesis_block, genesis_diff, chain_id, spec_id)
         })
     }
 
@@ -145,19 +149,21 @@ impl LocalBlockchain {
     #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub unsafe fn with_genesis_block_unchecked(
         genesis_block: LocalBlock,
-        genesis_state: TrieState,
+        genesis_diff: StateDiff,
         chain_id: U256,
         spec_id: SpecId,
     ) -> Self {
         let genesis_block: Arc<dyn SyncBlock<Error = BlockchainError>> = Arc::new(genesis_block);
 
         let total_difficulty = genesis_block.header().difficulty;
-        let storage =
-            ReservableSparseBlockchainStorage::with_genesis_block(genesis_block, total_difficulty);
+        let storage = ReservableSparseBlockchainStorage::with_genesis_block(
+            genesis_block,
+            genesis_diff,
+            total_difficulty,
+        );
 
         Self {
             storage,
-            genesis_state,
             chain_id,
             spec_id,
         }
@@ -243,15 +249,14 @@ impl Blockchain for LocalBlockchain {
     async fn state_at_block_number(
         &self,
         block_number: &U256,
+        state_overrides: &HashMap<U256, StateOverride>,
     ) -> Result<Box<dyn SyncState<Self::StateError>>, Self::BlockchainError> {
         if *block_number > self.last_block_number().await {
             return Err(BlockchainError::UnknownBlockNumber);
         }
 
-        let mut state = self.genesis_state.clone();
-        if *block_number > U256::ZERO {
-            compute_state_at_block(&mut state, &self.storage, block_number);
-        }
+        let mut state = TrieState::default();
+        compute_state_at_block(&mut state, &self.storage, block_number, state_overrides);
 
         Ok(Box::new(state))
     }