Merge remote-tracking branch 'origin/edr/main' into edr/feat/send-tra…

…nsaction

crates/edr_evm/benches/state/database_commit.rs

@@ -57,7 +57,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/edr_evm/benches/state/util.rs

@@ -4,11 +4,9 @@ use std::sync::Arc;
 
 use criterion::{BatchSize, BenchmarkId, Criterion};
 use edr_eth::{Address, Bytes, U256};
-use edr_evm::state::{StateError, SyncState, TrieState};
-#[cfg(all(test, feature = "test-remote"))]
-use edr_evm::{state::ForkState, HashMap, RandomHashGenerator};
 #[cfg(all(test, feature = "test-remote"))]
-use parking_lot::Mutex;
+use edr_evm::state::ForkState;
+use edr_evm::state::{StateError, SyncState, TrieState};
 use revm::primitives::{AccountInfo, Bytecode, KECCAK_EMPTY};
 use tempfile::TempDir;
 #[cfg(all(test, feature = "test-remote"))]
@@ -43,7 +41,9 @@ impl EdrStates {
 
         #[cfg(all(test, feature = "test-remote"))]
         let fork = {
-            use edr_eth::remote::RpcClient;
+            use edr_eth::remote::{BlockSpec, RpcClient};
+            use edr_evm::RandomHashGenerator;
+            use parking_lot::Mutex;
 
             let rpc_client = Arc::new(RpcClient::new(
                 &std::env::var_os("ALCHEMY_URL")
@@ -53,15 +53,18 @@ impl EdrStates {
                 cache_dir.path().to_path_buf(),
             ));
 
-            runtime
-                .block_on(ForkState::new(
-                    runtime.handle().clone(),
-                    rpc_client,
-                    Arc::new(Mutex::new(RandomHashGenerator::with_seed("seed"))),
-                    fork_block_number,
-                    HashMap::default(),
-                ))
-                .expect("Failed to construct ForkedState")
+            let block = runtime
+                .block_on(rpc_client.get_block_by_number(BlockSpec::Number(fork_block_number)))
+                .expect("failed to retrieve block by number")
+                .expect("block should exist");
+
+            ForkState::new(
+                runtime.handle().clone(),
+                rpc_client,
+                Arc::new(Mutex::new(RandomHashGenerator::with_seed("seed"))),
+                fork_block_number,
+                block.state_root,
+            )
         };
 
         Self {

crates/edr_evm/src/block/builder.rs

@@ -14,9 +14,8 @@ use edr_eth::{
 use revm::{
     db::DatabaseComponentError,
     primitives::{
-        Account, AccountInfo, AccountStatus, BlobExcessGasAndPrice, BlockEnv, CfgEnv, EVMError,
-        ExecutionResult, HashMap, InvalidHeader, InvalidTransaction, Output, ResultAndState,
-        SpecId,
+        AccountInfo, BlobExcessGasAndPrice, BlockEnv, CfgEnv, EVMError, ExecutionResult,
+        InvalidHeader, InvalidTransaction, Output, ResultAndState, SpecId,
     },
 };
 
@@ -130,7 +129,7 @@ impl BlockBuilder {
             header,
             callers: Vec::new(),
             transactions: Vec::new(),
-            state_diff: StateDiff::new(),
+            state_diff: StateDiff::default(),
             receipts: Vec::new(),
             parent_gas_limit,
         })
@@ -210,7 +209,8 @@ impl BlockBuilder {
             state: state_diff,
         } = run_transaction(evm, inspector)?;
 
-        self.state_diff.extend(state_diff.clone());
+        self.state_diff.apply_diff(state_diff.clone());
+
         state.commit(state_diff);
 
         self.header.gas_used += result.gas_used();
@@ -314,14 +314,7 @@ impl BlockBuilder {
                 .basic(address)?
                 .expect("Account must exist after modification");
 
-            self.state_diff.insert(
-                address,
-                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/edr_evm/src/blockchain.rs

@@ -4,7 +4,7 @@ mod remote;
 /// Storage data structures for a blockchain
 pub mod storage;
 
-use std::{fmt::Debug, sync::Arc};
+use std::{collections::BTreeMap, fmt::Debug, ops::Bound::Included, sync::Arc};
 
 use edr_eth::{
     receipt::BlockReceipt, remote::RpcClientError, spec::HardforkActivations, B256, U256,
@@ -17,7 +17,7 @@ pub use self::{
     local::{CreationError as LocalCreationError, LocalBlockchain},
 };
 use crate::{
-    state::{StateDiff, SyncState},
+    state::{StateDiff, StateOverride, SyncState},
     Block, LocalBlock, SyncBlock,
 };
 
@@ -131,10 +131,16 @@ 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.
     fn state_at_block_number(
         &self,
         block_number: u64,
+        // Block number -> state overrides
+        state_overrides: &BTreeMap<u64, StateOverride>,
     ) -> Result<Box<dyn SyncState<Self::StateError>>, Self::BlockchainError>;
 
     /// Retrieves the total difficulty at the block with the provided hash.
@@ -192,14 +198,34 @@ where
 fn compute_state_at_block<BlockT: Block + Clone>(
     state: &mut dyn DatabaseCommit,
     local_storage: &ReservableSparseBlockchainStorage<BlockT>,
-    block_number: u64,
+    first_local_block_number: u64,
+    last_local_block_number: u64,
+    state_overrides: &BTreeMap<u64, 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");
+        .state_diffs_until_block(last_local_block_number)
+        .unwrap_or_default();
 
-    for state_diff in state_diffs {
-        state.commit(state_diff.clone());
+    let mut overriden_state_diffs: BTreeMap<u64, StateDiff> = state_diffs
+        .iter()
+        .map(|(block_number, state_diff)| (*block_number, state_diff.clone()))
+        .collect();
+
+    for (block_number, state_override) in state_overrides.range((
+        Included(&first_local_block_number),
+        Included(&last_local_block_number),
+    )) {
+        overriden_state_diffs
+            .entry(*block_number)
+            .and_modify(|state_diff| {
+                state_diff.apply_diff(state_override.diff.as_inner().clone());
+            })
+            .or_insert_with(|| state_override.diff.clone());
+    }
+
+    for (_block_number, state_diff) in overriden_state_diffs {
+        state.commit(state_diff.into());
     }
 }
 

crates/edr_evm/src/blockchain/forked.rs

@@ -1,16 +1,16 @@
-use std::{num::NonZeroU64, sync::Arc};
+use std::{collections::BTreeMap, num::NonZeroU64, sync::Arc};
 
 use edr_eth::{
     block::{largest_safe_block_number, safe_block_depth, LargestSafeBlockNumberArgs},
     receipt::BlockReceipt,
     remote::{RpcClient, RpcClientError},
     spec::{chain_hardfork_activations, chain_name, HardforkActivations},
-    Address, B256, U256,
+    B256, U256,
 };
 use parking_lot::Mutex;
 use revm::{
     db::BlockHashRef,
-    primitives::{AccountInfo, HashMap, SpecId},
+    primitives::{HashMap, SpecId},
 };
 use tokio::runtime;
 
@@ -19,7 +19,7 @@ use super::{
     validate_next_block, Blockchain, BlockchainError, BlockchainMut,
 };
 use crate::{
-    state::{ForkState, StateDiff, StateError, SyncState},
+    state::{ForkState, StateDiff, StateError, StateOverride, SyncState},
     Block, LocalBlock, RandomHashGenerator, SyncBlock,
 };
 
@@ -55,8 +55,7 @@ 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,
+    state_root_generator: Arc<Mutex<RandomHashGenerator>>,
     fork_block_number: u64,
     chain_id: u64,
     network_id: u64,
@@ -73,7 +72,6 @@ impl ForkedBlockchain {
         rpc_client: RpcClient,
         fork_block_number: Option<u64>,
         state_root_generator: Arc<Mutex<RandomHashGenerator>>,
-        account_overrides: HashMap<Address, AccountInfo>,
         hardfork_activation_overrides: HashMap<u64, HardforkActivations>,
     ) -> Result<Self, CreationError> {
         let (chain_id, network_id, latest_block_number) = tokio::join!(
@@ -142,19 +140,11 @@ impl ForkedBlockchain {
         }
 
         let rpc_client = Arc::new(rpc_client);
-        let fork_state = ForkState::new(
-            runtime.clone(),
-            rpc_client.clone(),
-            state_root_generator,
-            fork_block_number,
-            account_overrides,
-        )
-        .await?;
 
         Ok(Self {
             local_storage: ReservableSparseBlockchainStorage::empty(fork_block_number),
             remote: RemoteBlockchain::new(rpc_client, runtime),
-            fork_state,
+            state_root_generator,
             fork_block_number,
             chain_id,
             network_id,
@@ -349,33 +339,50 @@ impl Blockchain for ForkedBlockchain {
     fn state_at_block_number(
         &self,
         block_number: u64,
+        state_overrides: &BTreeMap<u64, StateOverride>,
     ) -> Result<Box<dyn SyncState<Self::StateError>>, Self::BlockchainError> {
         if block_number > self.last_block_number() {
             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
-
-                tokio::task::block_in_place(move || {
-                    self.runtime().block_on(ForkState::new(
-                        self.runtime().clone(),
-                        self.remote.client().clone(),
-                        self.fork_state.state_root_generator().clone(),
-                        block_number,
-                        HashMap::new(),
-                    ))
-                })?
-            }
-            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);
-                state
-            }
+        let state_root = if let Some(state_override) = state_overrides.get(&block_number) {
+            state_override.state_root
+        } else {
+            self.block_by_number(block_number)?
+                .expect("The block is validated to exist")
+                .header()
+                .state_root
         };
 
+        let mut state = ForkState::new(
+            self.runtime().clone(),
+            self.remote.client().clone(),
+            self.state_root_generator.clone(),
+            block_number,
+            state_root,
+        );
+
+        let (first_block_number, last_block_number) =
+            match block_number.cmp(&self.fork_block_number) {
+                // Only override the state at the forked block
+                std::cmp::Ordering::Less => (block_number, block_number),
+                // Override blocks between the forked block and the requested block
+                std::cmp::Ordering::Equal | std::cmp::Ordering::Greater => {
+                    (self.fork_block_number, block_number)
+                }
+            };
+
+        compute_state_at_block(
+            &mut state,
+            &self.local_storage,
+            first_block_number,
+            last_block_number,
+            state_overrides,
+        );
+
+        // Override the state root in case the local state was modified
+        state.set_state_root(state_root);
+
         Ok(Box::new(state))
     }