fix: restore bucket capacity overrides

Author: WonderLawrence

bin/mega-evme/src/common/env.rs

@@ -5,17 +5,19 @@ use std::str::FromStr;
 use alloy_primitives::{Address, B256, U256};
 use clap::{Args, Parser};
 use mega_evm::{
-    alloy_evm::Database,
     revm::{
         context::{block::BlockEnv, cfg::CfgEnv},
+        database::{DatabaseRef, WrapDatabaseRef},
         primitives::eip4844,
     },
-    BucketId, MegaContext, MegaSpecId, TestExternalEnvs,
+    BucketId, MegaContext, MegaSpecId, TestDatabaseWrapper, TestExternalEnvs,
 };
 use tracing::{debug, trace};
 
 use super::{EvmeError, Result};
 
+type ConfiguredRefDatabase<DB> = TestDatabaseWrapper<WrapDatabaseRef<DB>>;
+
 /// Chain configuration arguments (spec and chain ID)
 #[derive(Args, Debug, Clone)]
 #[command(next_help_heading = "Chain Options")]
@@ -127,19 +129,20 @@ pub struct ExtEnvArgs {
 }
 
 impl ExtEnvArgs {
-    /// Creates [`TestExternalEnvs`] and returns bucket capacity configuration.
-    pub fn create_external_envs(&self) -> Result<(TestExternalEnvs, Vec<(BucketId, u64)>)> {
-        let external_envs = TestExternalEnvs::new();
+    /// Parses CLI-provided bucket capacity overrides.
+    pub fn bucket_capacities(&self) -> Result<Vec<(BucketId, u64)>> {
+        self.bucket_capacity
+            .iter()
+            .map(|bucket_capacity| parse_bucket_capacity(bucket_capacity))
+            .collect()
+    }
 
-        // Parse bucket capacities
-        let mut bucket_capacities = Vec::new();
-        for bucket_capacity_str in &self.bucket_capacity {
-            let (bucket_id, capacity) = parse_bucket_capacity(bucket_capacity_str)?;
-            bucket_capacities.push((bucket_id, capacity));
-        }
-        debug!(external_envs = ?external_envs, bucket_capacities = ?bucket_capacities, "Evm TestExternalEnvs created");
+    /// Creates [`TestExternalEnvs`].
+    pub fn create_external_envs(&self) -> Result<TestExternalEnvs> {
+        let external_envs = TestExternalEnvs::new();
+        debug!(external_envs = ?external_envs, "Evm TestExternalEnvs created");
 
-        Ok((external_envs, bucket_capacities))
+        Ok(external_envs)
     }
 }
 
@@ -175,19 +178,26 @@ impl EnvArgs {
         self.block.create_block_env()
     }
 
-    /// Creates [`TestExternalEnvs`] and returns bucket capacity configuration.
-    pub fn create_external_envs(&self) -> Result<(TestExternalEnvs, Vec<(BucketId, u64)>)> {
+    /// Creates [`TestExternalEnvs`].
+    pub fn create_external_envs(&self) -> Result<TestExternalEnvs> {
         self.ext.create_external_envs()
     }
 
     /// Creates a [`MegaContext`] with all environment configurations.
-    pub fn create_evm_context<DB: Database>(
+    pub fn create_evm_context<DB>(
         &self,
         db: DB,
-    ) -> Result<MegaContext<DB, TestExternalEnvs>> {
+    ) -> Result<MegaContext<ConfiguredRefDatabase<DB>, TestExternalEnvs>>
+    where
+        DB: DatabaseRef,
+        <DB as DatabaseRef>::Error: Send + Sync + 'static,
+    {
         let cfg = self.create_cfg_env()?;
         let block = self.create_block_env()?;
-        let (external_envs, _bucket_capacities) = self.create_external_envs()?;
+        let external_envs = self.create_external_envs()?;
+        let bucket_capacities = self.ext.bucket_capacities()?;
+        let db = TestDatabaseWrapper::new(WrapDatabaseRef::from(db))
+            .with_bucket_capacities(bucket_capacities);
 
         Ok(MegaContext::new(db, cfg.spec)
             .with_cfg(cfg)
@@ -218,3 +228,25 @@ pub fn parse_bucket_capacity(s: &str) -> Result<(u32, u64)> {
     trace!(string = %s, bucket_id = %bucket_id, capacity = %capacity, "Parsed bucket capacity");
     Ok((bucket_id, capacity))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use mega_evm::{revm::Database as _, test_utils::MemoryDatabase, MIN_BUCKET_SIZE};
+
+    #[test]
+    fn bucket_capacities_apply_to_wrapped_database() {
+        let address = Address::repeat_byte(0x11);
+        let bucket_id = TestDatabaseWrapper::<MemoryDatabase>::bucket_id_for_account(address);
+        let capacity = MIN_BUCKET_SIZE as u64 * 4;
+        let args = ExtEnvArgs { bucket_capacity: vec![format!("{bucket_id}:{capacity}")] };
+
+        let wrapped = TestDatabaseWrapper::new(MemoryDatabase::default())
+            .with_bucket_capacities(args.bucket_capacities().unwrap());
+
+        assert_eq!(
+            wrapped.salt_bucket_capacity(address, None).unwrap(),
+            (bucket_id as usize, capacity)
+        );
+    }
+}

bin/mega-evme/src/common/trace.rs

@@ -8,6 +8,7 @@ use alloy_rpc_types_trace::geth::{
 };
 use clap::{Parser, ValueEnum};
 use mega_evm::{
+    alloy_evm::Database as AlloyDatabase,
     revm::{
         context::{
             result::{ExecutionResult, ResultAndState},
@@ -22,7 +23,7 @@ use mega_evm::{
 use revm_inspectors::tracing::{TracingInspector, TracingInspectorConfig};
 use tracing::{debug, info, trace};
 
-use super::{EvmeError, EvmeState};
+use super::EvmeError;
 
 /// Tracer type for execution analysis
 #[derive(Debug, Clone, Copy, ValueEnum, Default)]
@@ -213,14 +214,13 @@ impl TraceArgs {
     }
 
     /// Execute transaction with optional tracing
-    pub fn execute_transaction<N, P>(
+    pub fn execute_transaction<DB>(
         &self,
-        evm_context: MegaContext<&mut EvmeState<N, P>, TestExternalEnvs>,
+        evm_context: MegaContext<DB, TestExternalEnvs>,
         tx: MegaTransaction,
     ) -> Result<(ExecutionResult<MegaHaltReason>, EvmState, Option<String>), EvmeError>
     where
-        N: alloy_network::Network,
-        P: alloy_provider::Provider<N> + std::fmt::Debug,
+        DB: AlloyDatabase + DatabaseRef,
     {
         if self.is_tracing_enabled() {
             info!(tracer = ?self.tracer, "Evm executing with tracing");

bin/mega-evme/src/replay/cmd.rs

@@ -11,11 +11,11 @@ use mega_evm::{
     revm::{
         context::{result::ExecutionResult, BlockEnv, ContextTr},
         context_interface::block::BlobExcessGasAndPrice,
-        database::{states::bundle_state::BundleRetention, StateBuilder},
+        database::{states::bundle_state::BundleRetention, StateBuilder, WrapDatabaseRef},
         DatabaseRef,
     },
     BlockLimits, EvmTxRuntimeLimits, MegaBlockExecutionCtx, MegaBlockExecutorFactory,
-    MegaEvmFactory, MegaHardforks, MegaSpecId,
+    MegaEvmFactory, MegaHardforks, MegaSpecId, TestDatabaseWrapper,
 };
 use tracing::{debug, info, trace, warn};
 
@@ -226,7 +226,7 @@ impl Cmd {
         let transaction_index = preceding_transactions.len() as u64;
         debug!(transaction_index, "Setting up block executor");
 
-        let (external_env_factory, _bucket_capacities) = self.ext_args.create_external_envs()?;
+        let external_env_factory = self.ext_args.create_external_envs()?;
         let evm_factory = MegaEvmFactory::new().with_external_env_factory(external_env_factory);
         let block_executor_factory = MegaBlockExecutorFactory::new(
             &hardforks,
@@ -263,6 +263,9 @@ impl Cmd {
         let mut inspector = self.trace_args.create_inspector();
 
         // Create state and block executor with inspector
+        let bucket_capacities = self.ext_args.bucket_capacities()?;
+        let database = TestDatabaseWrapper::new(WrapDatabaseRef::from(&*database))
+            .with_bucket_capacities(bucket_capacities);
         let mut state = StateBuilder::new().with_database(database).with_bundle_update().build();
         let mut block_executor = block_executor_factory.create_executor_with_inspector(
             &mut state,

bin/mega-evme/src/run/cmd.rs

@@ -95,7 +95,7 @@ impl Cmd {
         }
 
         // Create EVM context and execute transaction
-        let evm_context = self.env_args.create_evm_context(&mut state)?;
+        let evm_context = self.env_args.create_evm_context(&state)?;
         let start = Instant::now();
         let (exec_result, evm_state, trace_data) =
             self.trace_args.execute_transaction(evm_context, tx)?;

bin/mega-evme/src/tx/cmd.rs

@@ -92,7 +92,7 @@ impl Cmd {
 
         // Step 3: Execute transaction
         info!("Executing transaction");
-        let evm_context = self.env_args.create_evm_context(&mut state)?;
+        let evm_context = self.env_args.create_evm_context(&state)?;
         let start = Instant::now();
         let (exec_result, evm_state, trace_data) =
             self.trace_args.execute_transaction(evm_context, tx.clone())?;

crates/mega-evm/src/external/test_utils.rs

@@ -141,12 +141,14 @@ const PLAIN_STORAGE_KEY_LEN: usize = PLAIN_ACCOUNT_KEY_LEN + SLOT_KEY_LEN;
 ///
 /// This wrapper delegates all database operations to the inner database,
 /// but adds configurable bucket capacity tracking for testing dynamic gas costs.
-#[derive(Debug, Clone)]
+#[derive(derive_more::Debug, Clone)]
 pub struct TestDatabaseWrapper<DB> {
     /// The inner database implementation.
+    #[debug(skip)]
     pub inner: DB,
     /// Bucket capacities. Maps bucket IDs to their capacity values.
     /// Buckets not in this map default to [`MIN_BUCKET_SIZE`](crate::MIN_BUCKET_SIZE).
+    #[debug(skip)]
     bucket_capacity: Rc<RefCell<HashMap<BucketId, u64>>>,
 }
 
@@ -170,7 +172,19 @@ impl<DB> TestDatabaseWrapper<DB> {
     ///
     /// `self` for method chaining
     pub fn with_bucket_capacity(self, bucket_id: BucketId, capacity: u64) -> Self {
-        self.bucket_capacity.borrow_mut().insert(bucket_id, capacity);
+        self.with_bucket_capacities([(bucket_id, capacity)])
+    }
+
+    /// Configures multiple bucket capacities at once.
+    ///
+    /// # Returns
+    ///
+    /// `self` for method chaining
+    pub fn with_bucket_capacities(
+        self,
+        bucket_capacities: impl IntoIterator<Item = (BucketId, u64)>,
+    ) -> Self {
+        self.bucket_capacity.borrow_mut().extend(bucket_capacities);
         self
     }
 
@@ -234,3 +248,50 @@ impl<DB: revm::Database> revm::Database for TestDatabaseWrapper<DB> {
         Ok((bucket_id as usize, capacity))
     }
 }
+
+impl<DB: revm::DatabaseRef> revm::DatabaseRef for TestDatabaseWrapper<DB> {
+    type Error = DB::Error;
+
+    fn basic_ref(&self, address: Address) -> Result<Option<revm::state::AccountInfo>, Self::Error> {
+        self.inner.basic_ref(address)
+    }
+
+    fn code_by_hash_ref(&self, code_hash: B256) -> Result<revm::state::Bytecode, Self::Error> {
+        self.inner.code_by_hash_ref(code_hash)
+    }
+
+    fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
+        self.inner.storage_ref(address, index)
+    }
+
+    fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
+        self.inner.block_hash_ref(number)
+    }
+
+    fn salt_bucket_capacity_ref(
+        &self,
+        address: Address,
+        index: Option<U256>,
+    ) -> Result<(usize, u64), Self::Error> {
+        let bucket_id = if let Some(key) = index {
+            Self::bucket_id_for_slot(address, key)
+        } else {
+            Self::bucket_id_for_account(address)
+        };
+
+        let capacity = self
+            .bucket_capacity
+            .borrow()
+            .get(&bucket_id)
+            .copied()
+            .unwrap_or(salt::constant::MIN_BUCKET_SIZE as u64);
+
+        Ok((bucket_id as usize, capacity))
+    }
+}
+
+impl<DB: revm::DatabaseCommit> revm::DatabaseCommit for TestDatabaseWrapper<DB> {
+    fn commit(&mut self, changes: revm::primitives::HashMap<Address, revm::state::Account>) {
+        self.inner.commit(changes);
+    }
+}