refactor: move salt bucket capacity lookup into revm::Database trait

Cargo.toml

@@ -13,7 +13,7 @@ exclude = []
 # https://doc.rust-lang.org/edition-guide/rust-2021/default-cargo-resolver.html
 resolver = "2"
 [workspace.package]
-version = "1.4.1"
+version = "1.4.2"
 edition = "2021"
 rust-version = "1.86"
 license = "MIT OR Apache-2.0"
@@ -177,6 +177,10 @@ tracing-subscriber = { version = "0.3", default-features = false }
 triehash = { version = "0.8", default-features = false }
 walkdir = { version = "2.5", default-features = false }
 
+[patch.crates-io]
+revm-database-interface = { git = "https://github.com/megaeth-labs/revm-database-interface.git", branch = "main" }
+revm-database = { git = "https://github.com/megaeth-labs/revm-database.git", branch = "main" }
+
 # Speed up compilation time for dev builds by reducing emitted debug info.
 # NOTE: Debuggers may provide less useful information with this setting.
 # Uncomment this section if you're using a debugger.

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,
     },
-    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,15 +129,17 @@ pub struct ExtEnvArgs {
 }
 
 impl ExtEnvArgs {
+    /// 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()
+    }
+
     /// Creates [`TestExternalEnvs`].
     pub fn create_external_envs(&self) -> Result<TestExternalEnvs> {
-        let mut external_envs = TestExternalEnvs::new();
-
-        // Parse and configure bucket capacities
-        for bucket_capacity_str in &self.bucket_capacity {
-            let (bucket_id, capacity) = parse_bucket_capacity(bucket_capacity_str)?;
-            external_envs = external_envs.with_bucket_capacity(bucket_id, capacity);
-        }
+        let external_envs = TestExternalEnvs::new();
         debug!(external_envs = ?external_envs, "Evm TestExternalEnvs created");
 
         Ok(external_envs)
@@ -180,13 +184,20 @@ impl EnvArgs {
     }
 
     /// 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 = 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)
@@ -217,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};
 
@@ -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/constants.rs

@@ -2,6 +2,18 @@
 //!
 //! It groups the constants for different EVM specs as sub-modules.
 
+// SALT bucket constants
+
+/// Number of bits to represent the minimum bucket size (8 bits = 256 slots).
+pub const MIN_BUCKET_SIZE_BITS: usize = 8;
+
+/// Minimum capacity of a SALT bucket in number of slots (256).
+///
+/// Buckets hold accounts or storage slots and can grow beyond this size. The gas cost
+/// multiplier is calculated as `capacity / MIN_BUCKET_SIZE`, so a bucket at minimum
+/// capacity has a 1x multiplier.
+pub const MIN_BUCKET_SIZE: usize = 1 << MIN_BUCKET_SIZE_BITS;
+
 /// Constants for the `EQUIVALENCE` spec.
 pub mod equivalence {
     use revm::interpreter::gas;

crates/mega-evm/src/evm/context.rs

@@ -54,7 +54,7 @@ pub struct MegaContext<DB: Database, ExtEnvs: ExternalEnvTypes> {
     pub additional_limit: Rc<RefCell<AdditionalLimit>>,
 
     /// Calculator for dynamic gas costs during transaction execution.
-    pub dynamic_storage_gas_cost: Rc<RefCell<DynamicGasCost<ExtEnvs::SaltEnv>>>,
+    pub dynamic_storage_gas_cost: Rc<RefCell<DynamicGasCost>>,
 
     /// The oracle environment.
     pub oracle_env: Rc<RefCell<ExtEnvs::OracleEnv>>,
@@ -109,7 +109,6 @@ impl<DB: Database> MegaContext<DB, EmptyExternalEnv> {
             additional_limit: Rc::new(RefCell::new(AdditionalLimit::new(spec, tx_limits))),
             dynamic_storage_gas_cost: Rc::new(RefCell::new(DynamicGasCost::new(
                 spec,
-                EmptyExternalEnv,
                 inner.block.number.to::<u64>().saturating_sub(1),
             ))),
             oracle_env: Rc::new(RefCell::new(EmptyExternalEnv)),
@@ -169,7 +168,6 @@ impl<DB: Database, ExtEnvTypes: ExternalEnvTypes> MegaContext<DB, ExtEnvTypes> {
             additional_limit: Rc::new(RefCell::new(AdditionalLimit::new(spec, tx_limits))),
             dynamic_storage_gas_cost: Rc::new(RefCell::new(DynamicGasCost::new(
                 spec,
-                external_envs.salt_env,
                 inner.block.number.to::<u64>() - 1,
             ))),
             oracle_env: Rc::new(RefCell::new(external_envs.oracle_env)),
@@ -274,10 +272,9 @@ impl<DB: Database, ExtEnvTypes: ExternalEnvTypes> MegaContext<DB, ExtEnvTypes> {
 
     /// Sets the external environments for the EVM.
     ///
-    /// This method updates the external environments used for gas cost calculations,
-    /// including the salt environment and oracle environment. When setting new
-    /// external environments, the dynamic gas cost calculator and oracle environment
-    /// are reinitialized with the new configurations.
+    /// This method updates the external environments used for oracle environment.
+    /// When setting new external environments, the oracle environment is reinitialized
+    /// with the new configuration.
     ///
     /// # Arguments
     ///
@@ -290,18 +287,12 @@ impl<DB: Database, ExtEnvTypes: ExternalEnvTypes> MegaContext<DB, ExtEnvTypes> {
         self,
         external_envs: ExternalEnvs<NewExtEnvTypes>,
     ) -> MegaContext<DB, NewExtEnvTypes> {
-        let parent_block_number = self.inner.block.number.to::<u64>().saturating_sub(1);
-        let spec = self.spec;
         MegaContext {
             inner: self.inner,
-            spec,
+            spec: self.spec,
             disable_beneficiary: self.disable_beneficiary,
             additional_limit: self.additional_limit,
-            dynamic_storage_gas_cost: Rc::new(RefCell::new(DynamicGasCost::new(
-                spec,
-                external_envs.salt_env,
-                parent_block_number,
-            ))),
+            dynamic_storage_gas_cost: self.dynamic_storage_gas_cost,
             oracle_env: Rc::new(RefCell::new(external_envs.oracle_env)),
             volatile_data_tracker: self.volatile_data_tracker,
             disable_sandbox: self.disable_sandbox,

crates/mega-evm/src/evm/factory.rs

@@ -141,10 +141,9 @@ impl<ExtEnvFactory: ExternalEnvFactory + Clone> alloy_evm::EvmFactory
         evm_env: EvmEnv<Self::Spec>,
     ) -> Self::Evm<DB, revm::inspector::NoOpInspector> {
         let spec_id = *evm_env.spec_id();
-        let block_number = evm_env.block_env.number.to();
         let runtime_limits = EvmTxRuntimeLimits::from_spec(spec_id);
         let ctx = MegaContext::new(db, spec_id)
-            .with_external_envs(self.external_env_factory.external_envs(block_number))
+            .with_external_envs(self.external_env_factory.external_envs())
             .with_tx(MegaTransaction::default())
             .with_block(evm_env.block_env)
             .with_cfg(evm_env.cfg_env)