Share block hashes within InMemoryStorage-s

crates/pevm/benches/gigagas.rs

@@ -2,7 +2,7 @@
 
 // TODO: More fancy benchmarks & plots.
 
-use std::{num::NonZeroUsize, thread};
+use std::{num::NonZeroUsize, sync::Arc, thread};
 
 use alloy_primitives::{Address, U160, U256};
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
@@ -34,7 +34,7 @@ const GIGA_GAS: u64 = 1_000_000_000;
 static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
 /// Runs a benchmark for executing a set of transactions on a given blockchain state.
-pub fn bench(c: &mut Criterion, name: &str, storage: InMemoryStorage<'_>, txs: Vec<TxEnv>) {
+pub fn bench(c: &mut Criterion, name: &str, storage: InMemoryStorage, txs: Vec<TxEnv>) {
     let concurrency_level = thread::available_parallelism().unwrap_or(NonZeroUsize::MIN);
     let chain = PevmEthereum::mainnet();
     let spec_id = SpecId::LATEST;
@@ -76,9 +76,10 @@ pub fn bench_raw_transfers(c: &mut Criterion) {
     let storage = InMemoryStorage::new(
         std::iter::once(MINER_ADDRESS)
             .chain(START_ADDRESS..START_ADDRESS + block_size)
-            .map(common::mock_account),
-        None,
-        [],
+            .map(common::mock_account)
+            .collect(),
+        Default::default(),
+        Default::default(),
     );
     bench(
         c,
@@ -108,7 +109,7 @@ pub fn bench_erc20(c: &mut Criterion) {
     bench(
         c,
         "Independent ERC20",
-        InMemoryStorage::new(state, Some(&bytecodes), []),
+        InMemoryStorage::new(state, Arc::new(bytecodes), Default::default()),
         txs,
     );
 }
@@ -128,7 +129,7 @@ pub fn bench_uniswap(c: &mut Criterion) {
     bench(
         c,
         "Independent Uniswap",
-        InMemoryStorage::new(final_state, Some(&final_bytecodes), []),
+        InMemoryStorage::new(final_state, Arc::new(final_bytecodes), Default::default()),
         final_txs,
     );
 }

crates/pevm/src/storage/in_memory.rs

@@ -1,34 +1,35 @@
 use std::fmt::Debug;
+use std::sync::Arc;
 
 use alloy_primitives::{keccak256, Address, B256, U256};
 
 use super::{BlockHashes, Bytecodes, ChainState, EvmCode};
-use crate::{AccountBasic, EvmAccount, Storage};
+use crate::{AccountBasic, Storage};
 
 /// A storage that stores chain data in memory.
-#[derive(Debug, Default, Clone)]
-pub struct InMemoryStorage<'a> {
+#[derive(Debug, Clone, Default)]
+pub struct InMemoryStorage {
     accounts: ChainState,
-    bytecodes: Option<&'a Bytecodes>,
-    block_hashes: BlockHashes,
+    bytecodes: Arc<Bytecodes>,
+    block_hashes: Arc<BlockHashes>,
 }
 
-impl<'a> InMemoryStorage<'a> {
+impl InMemoryStorage {
     /// Construct a new [`InMemoryStorage`]
-    pub fn new(
-        accounts: impl IntoIterator<Item = (Address, EvmAccount)>,
-        bytecodes: Option<&'a Bytecodes>,
-        block_hashes: impl IntoIterator<Item = (u64, B256)>,
+    pub const fn new(
+        accounts: ChainState,
+        bytecodes: Arc<Bytecodes>,
+        block_hashes: Arc<BlockHashes>,
     ) -> Self {
-        InMemoryStorage {
-            accounts: accounts.into_iter().collect(),
+        Self {
+            accounts,
             bytecodes,
-            block_hashes: block_hashes.into_iter().collect(),
+            block_hashes,
         }
     }
 }
 
-impl Storage for InMemoryStorage<'_> {
+impl Storage for InMemoryStorage {
     // TODO: More proper error handling
     type Error = u8;
 
@@ -47,10 +48,7 @@ impl Storage for InMemoryStorage<'_> {
     }
 
     fn code_by_hash(&self, code_hash: &B256) -> Result<Option<EvmCode>, Self::Error> {
-        Ok(match self.bytecodes {
-            Some(bytecodes) => bytecodes.get(code_hash).cloned(),
-            None => None,
-        })
+        Ok(self.bytecodes.get(code_hash).cloned())
     }
 
     fn has_storage(&self, address: &Address) -> Result<bool, Self::Error> {

crates/pevm/tests/beneficiary.rs

@@ -12,7 +12,11 @@ const BLOCK_SIZE: usize = 100_000;
 fn test_beneficiary(get_address: fn(usize) -> Address) {
     common::test_execute_revm(
         // Mock the beneficiary account (`Address:ZERO`) and the next `BLOCK_SIZE` user accounts.
-        InMemoryStorage::new((0..=BLOCK_SIZE).map(common::mock_account), None, []),
+        InMemoryStorage::new(
+            (0..=BLOCK_SIZE).map(common::mock_account).collect(),
+            Default::default(),
+            Default::default(),
+        ),
         // Mock `BLOCK_SIZE` transactions sending some tokens to itself.
         // Skipping `Address::ZERO` as the beneficiary account.
         (1..=BLOCK_SIZE)

crates/pevm/tests/common/mod.rs

@@ -3,14 +3,17 @@
 use std::{
     fs::{self, File},
     io::BufReader,
+    sync::Arc,
 };
 
 use alloy_consensus::{Signed, TxLegacy};
 use alloy_primitives::{Address, Bytes, PrimitiveSignature, TxKind, B256, U256};
 use alloy_rpc_types_eth::{Block, BlockTransactions, Header};
 use flate2::bufread::GzDecoder;
 use hashbrown::HashMap;
-use pevm::{chain::PevmChain, BlockHashes, BuildSuffixHasher, EvmAccount, InMemoryStorage};
+use pevm::{
+    chain::PevmChain, BlockHashes, BuildSuffixHasher, ChainState, EvmAccount, InMemoryStorage,
+};
 
 /// runner module
 pub mod runner;
@@ -26,18 +29,20 @@ pub const RAW_TRANSFER_GAS_LIMIT: u64 = 21_000;
 
 // TODO: Put somewhere better?
 /// Iterates over blocks stored on disk and processes each block using the provided handler.
-pub fn for_each_block_from_disk(mut handler: impl FnMut(Block, InMemoryStorage<'_>)) {
+pub fn for_each_block_from_disk(mut handler: impl FnMut(Block, InMemoryStorage)) {
     let data_dir = std::path::PathBuf::from("../../data");
 
     // TODO: Deduplicate logic with [bin/fetch.rs] when there is more usage
     let bytecodes = bincode::deserialize_from(GzDecoder::new(BufReader::new(
         File::open(data_dir.join("bytecodes.bincode.gz")).unwrap(),
     )))
+    .map(Arc::new)
     .unwrap();
 
     let block_hashes = bincode::deserialize_from::<_, BlockHashes>(BufReader::new(
         File::open(data_dir.join("block_hashes.bincode")).unwrap(),
     ))
+    .map(Arc::new)
     .unwrap();
 
     for block_path in fs::read_dir(data_dir.join("blocks")).unwrap() {
@@ -57,7 +62,7 @@ pub fn for_each_block_from_disk(mut handler: impl FnMut(Block, InMemoryStorage<'
 
         handler(
             block,
-            InMemoryStorage::new(accounts, Some(&bytecodes), block_hashes.clone()),
+            InMemoryStorage::new(accounts, Arc::clone(&bytecodes), Arc::clone(&block_hashes)),
         );
     }
 }
@@ -67,7 +72,7 @@ pub fn test_independent_raw_transfers<C>(chain: &C, block_size: usize)
 where
     C: PevmChain + Send + Sync + PartialEq,
 {
-    let accounts: Vec<(Address, EvmAccount)> = (0..block_size).map(mock_account).collect();
+    let accounts = (0..block_size).map(mock_account).collect::<ChainState>();
     let block: Block<C::Transaction> = Block {
         header: Header {
             inner: alloy_consensus::Header {
@@ -104,6 +109,6 @@ where
         ),
         ..Block::<C::Transaction>::default()
     };
-    let storage = InMemoryStorage::new(accounts, None, []);
+    let storage = InMemoryStorage::new(accounts, Default::default(), Default::default());
     test_execute_alloy(&storage, chain, block, false);
 }

crates/pevm/tests/erc20/main.rs

@@ -12,13 +12,17 @@ use common::test_execute_revm;
 use erc20::generate_cluster;
 use pevm::{Bytecodes, ChainState, EvmAccount, InMemoryStorage};
 use revm::primitives::{Address, TxEnv};
+use std::sync::Arc;
 
 #[test]
 fn erc20_independent() {
     const N: usize = 37123;
     let (mut state, bytecodes, txs) = generate_cluster(N, 1, 1);
     state.insert(Address::ZERO, EvmAccount::default()); // Beneficiary
-    test_execute_revm(InMemoryStorage::new(state, Some(&bytecodes), []), txs);
+    test_execute_revm(
+        InMemoryStorage::new(state, Arc::new(bytecodes), Default::default()),
+        txs,
+    );
 }
 
 #[test]
@@ -43,7 +47,7 @@ fn erc20_clusters() {
         final_txs.extend(txs);
     }
     common::test_execute_revm(
-        InMemoryStorage::new(final_state, Some(&final_bytecodes), []),
+        InMemoryStorage::new(final_state, Arc::new(final_bytecodes), Default::default()),
         final_txs,
     )
 }

crates/pevm/tests/ethereum/main.rs

@@ -30,6 +30,7 @@ use revme::cmd::statetest::{
 };
 use std::path::Path;
 use std::str::FromStr;
+use std::sync::Arc;
 use std::{fs, num::NonZeroUsize};
 use walkdir::{DirEntry, WalkDir};
 
@@ -148,7 +149,7 @@ fn run_test_unit(path: &Path, unit: TestUnit) {
             match (
                 test.expect_exception.as_deref(),
                 Pevm::default().execute_revm_parallel(
-                    &InMemoryStorage::new(chain_state.clone(), Some(&bytecodes), []),
+                    &InMemoryStorage::new(chain_state.clone(), Arc::new(bytecodes), Default::default()),
                     &PevmEthereum::mainnet(),
                     spec_name.to_spec_id(),
                     build_block_env(&unit.env),

crates/pevm/tests/mixed.rs

@@ -3,6 +3,7 @@
 use pevm::{Bytecodes, ChainState, EvmAccount, InMemoryStorage};
 use rand::random;
 use revm::primitives::{env::TxEnv, Address, TransactTo, U256};
+use std::sync::Arc;
 
 pub mod common;
 pub mod erc20;
@@ -58,7 +59,7 @@ fn mixed_block() {
         }
     }
     common::test_execute_revm(
-        InMemoryStorage::new(final_state, Some(&final_bytecodes), []),
+        InMemoryStorage::new(final_state, Arc::new(final_bytecodes), Default::default()),
         // TODO: Shuffle transactions to scatter dependencies around the block.
         // Note that we'll need to guarantee that the nonces are increasing.
         final_txs,

crates/pevm/tests/raw_transfers.rs

@@ -11,7 +11,11 @@ fn raw_transfers_independent() {
     let block_size = 100_000; // number of transactions
     common::test_execute_revm(
         // Mock the beneficiary account (`Address:ZERO`) and the next `block_size` user accounts.
-        InMemoryStorage::new((0..=block_size).map(common::mock_account), None, []),
+        InMemoryStorage::new(
+            (0..=block_size).map(common::mock_account).collect(),
+            Default::default(),
+            Default::default(),
+        ),
         // Mock `block_size` transactions sending some tokens to itself.
         // Skipping `Address::ZERO` as the beneficiary account.
         (1..=block_size)
@@ -41,7 +45,11 @@ fn raw_transfers_same_sender_multiple_txs() {
 
     common::test_execute_revm(
         // Mock the beneficiary account (`Address:ZERO`) and the next `block_size` user accounts.
-        InMemoryStorage::new((0..=block_size).map(common::mock_account), None, []),
+        InMemoryStorage::new(
+            (0..=block_size).map(common::mock_account).collect(),
+            Default::default(),
+            Default::default(),
+        ),
         (1..=block_size)
             .map(|i| {
                 // Insert a "parallel" transaction every ~256 transactions

crates/pevm/tests/small_blocks.rs

@@ -15,7 +15,11 @@ fn empty_revm_block() {
 #[test]
 fn one_tx_revm_block() {
     common::test_execute_revm(
-        InMemoryStorage::new([common::mock_account(0)], None, []),
+        InMemoryStorage::new(
+            [common::mock_account(0)].into_iter().collect(),
+            Default::default(),
+            Default::default(),
+        ),
         vec![TxEnv {
             caller: Address::ZERO,
             transact_to: TransactTo::Call(Address::ZERO),

crates/pevm/tests/uniswap/main.rs

@@ -14,6 +14,7 @@ pub mod uniswap;
 use crate::uniswap::generate_cluster;
 use pevm::{Bytecodes, ChainState, EvmAccount, InMemoryStorage};
 use revm::primitives::{Address, TxEnv};
+use std::sync::Arc;
 
 #[test]
 fn uniswap_clusters() {
@@ -33,7 +34,7 @@ fn uniswap_clusters() {
         final_txs.extend(txs);
     }
     common::test_execute_revm(
-        InMemoryStorage::new(final_state, Some(&final_bytecodes), []),
+        InMemoryStorage::new(final_state, Arc::new(final_bytecodes), Default::default()),
         final_txs,
     )
 }