Gradual contract code migration to offchain storage (#1445)

crates/evm/src/call.rs

@@ -64,7 +64,7 @@ impl<C: sov_modules_api::Context> Evm<C> {
             .map(|info| info.nonce)
             .unwrap_or(0);
 
-        let db: EvmDb<'_, C> = self.get_db(working_set);
+        let db: EvmDb<'_, C> = self.get_db(working_set, cfg_env.handler_cfg.spec_id);
         let system_txs = create_system_transactions(system_events, system_nonce, cfg_env.chain_id);
 
         let mut citrea_handler_ext = CitreaExternal::new(l1_fee_rate);
@@ -154,7 +154,7 @@ impl<C: sov_modules_api::Context> Evm<C> {
             log_index_start = tx.receipt.log_index_start + tx.receipt.receipt.logs.len() as u64;
         }
 
-        let evm_db: EvmDb<'_, C> = self.get_db(working_set);
+        let evm_db: EvmDb<'_, C> = self.get_db(working_set, cfg_env.handler_cfg.spec_id);
 
         let results = executor::execute_multiple_tx(
             evm_db,

crates/evm/src/evm/db.rs

@@ -2,7 +2,8 @@
 use std::collections::HashMap;
 
 use reth_primitives::{keccak256, Address, B256};
-use revm::primitives::{AccountInfo as ReVmAccountInfo, Bytecode, U256};
+use revm::primitives::SpecId::CANCUN;
+use revm::primitives::{AccountInfo as ReVmAccountInfo, Bytecode, SpecId, U256};
 use revm::Database;
 use sov_modules_api::{StateMapAccessor, WorkingSet};
 use sov_state::codec::BcsCodec;
@@ -34,23 +35,29 @@ impl std::error::Error for DBError {
 
 pub(crate) struct EvmDb<'a, C: sov_modules_api::Context> {
     pub(crate) accounts: sov_modules_api::StateMap<Address, AccountInfo, BcsCodec>,
-    pub(crate) code: sov_modules_api::OffchainStateMap<B256, Bytecode, BcsCodec>,
+    pub(crate) code: sov_modules_api::StateMap<B256, Bytecode, BcsCodec>,
+    pub(crate) offchain_code: sov_modules_api::OffchainStateMap<B256, Bytecode, BcsCodec>,
     pub(crate) last_block_hashes: sov_modules_api::StateMap<U256, B256, BcsCodec>,
     pub(crate) working_set: &'a mut WorkingSet<C>,
+    pub(crate) current_spec: SpecId,
 }
 
 impl<'a, C: sov_modules_api::Context> EvmDb<'a, C> {
     pub(crate) fn new(
         accounts: sov_modules_api::StateMap<Address, AccountInfo, BcsCodec>,
-        code: sov_modules_api::OffchainStateMap<B256, Bytecode, BcsCodec>,
+        code: sov_modules_api::StateMap<B256, Bytecode, BcsCodec>,
+        offchain_code: sov_modules_api::OffchainStateMap<B256, Bytecode, BcsCodec>,
         last_block_hashes: sov_modules_api::StateMap<U256, B256, BcsCodec>,
         working_set: &'a mut WorkingSet<C>,
+        current_spec: SpecId,
     ) -> Self {
         Self {
             accounts,
             code,
+            offchain_code,
             last_block_hashes,
             working_set,
+            current_spec,
         }
     }
 
@@ -80,6 +87,17 @@ impl<'a, C: sov_modules_api::Context> EvmDb<'a, C> {
             );
         }
     }
+
+    pub(crate) fn check_against_code_hash(
+        &self,
+        code: &Bytecode,
+        code_hash: &B256,
+    ) -> Result<(), DBError> {
+        if *code_hash != keccak256(code.original_bytes()) {
+            return Err(DBError::CodeHashMismatch);
+        }
+        Ok(())
+    }
 }
 
 impl<'a, C: sov_modules_api::Context> Database for EvmDb<'a, C> {
@@ -92,12 +110,24 @@ impl<'a, C: sov_modules_api::Context> Database for EvmDb<'a, C> {
 
     fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
         // TODO move to new_raw_with_hash for better performance
-        let code = self
-            .code
-            .get(&code_hash, &mut self.working_set.offchain_state());
+
+        // If CANCUN or later forks are activated, try to fetch code from offchain storage
+        // first. This is to prevent slower lookups in `code`.
+        if self.current_spec.is_enabled_in(CANCUN) {
+            if let Some(code) = self
+                .offchain_code
+                .get(&code_hash, &mut self.working_set.offchain_state())
+            {
+                self.check_against_code_hash(&code, &code_hash)?;
+                return Ok(code);
+            }
+        }
+        let code = self.code.get(&code_hash, self.working_set);
         if let Some(code) = code {
-            if code_hash != keccak256(code.original_bytes()) {
-                return Err(DBError::CodeHashMismatch);
+            // Gradually migrate contract codes into the offchain code state map.
+            if self.current_spec.is_enabled_in(CANCUN) {
+                self.offchain_code
+                    .set(&code_hash, &code, &mut self.working_set.offchain_state());
             }
             Ok(code)
         } else {

crates/evm/src/evm/db_commit.rs

@@ -1,6 +1,7 @@
 use std::collections::BTreeMap;
 
 use alloy_primitives::{Address, U256};
+use revm::primitives::SpecId::CANCUN;
 use revm::primitives::{Account, AccountInfo, HashMap};
 use revm::DatabaseCommit;
 use sov_modules_api::{StateMapAccessor, StateVecAccessor};
@@ -51,13 +52,22 @@ impl<'a, C: sov_modules_api::Context> DatabaseCommit for EvmDb<'a, C> {
 
             if let Some(ref code) = account_info.code {
                 if !code.is_empty() {
-                    let offchain_state = &mut self.working_set.offchain_state();
                     let exists_in_db = self
                         .code
-                        .get(&account_info.code_hash, offchain_state)
+                        .get(&account_info.code_hash, self.working_set)
                         .is_some();
+
                     if !exists_in_db {
-                        self.code.set(&account_info.code_hash, code, offchain_state);
+                        if self.current_spec.is_enabled_in(CANCUN) {
+                            self.offchain_code.set(
+                                &account_info.code_hash,
+                                code,
+                                &mut self.working_set.offchain_state(),
+                            );
+                        } else {
+                            self.code
+                                .set(&account_info.code_hash, code, self.working_set);
+                        }
                     }
                 }
             }

crates/evm/src/evm/db_init.rs

@@ -1,6 +1,7 @@
 use reth_primitives::U256;
 #[cfg(test)]
 use revm::db::{CacheDB, EmptyDB};
+use revm::primitives::SpecId::CANCUN;
 use revm::primitives::{Address, Bytecode, B256};
 use sov_modules_api::StateMapAccessor;
 
@@ -20,8 +21,12 @@ impl<'a, C: sov_modules_api::Context> InitEvmDb for EvmDb<'a, C> {
     }
 
     fn insert_code(&mut self, code_hash: B256, code: Bytecode) {
-        self.code
-            .set(&code_hash, &code, &mut self.working_set.offchain_state())
+        if self.current_spec.is_enabled_in(CANCUN) {
+            self.offchain_code
+                .set(&code_hash, &code, &mut self.working_set.offchain_state())
+        } else {
+            self.code.set(&code_hash, &code, self.working_set)
+        }
     }
 
     fn insert_storage(&mut self, address: Address, index: U256, value: U256) {

crates/evm/src/evm/tests.rs

@@ -1,6 +1,7 @@
 use std::str::FromStr;
 
 use reth_primitives::{Address, TxKind};
+use revm::primitives::SpecId::SHANGHAI;
 use revm::primitives::{CfgEnvWithHandlerCfg, ExecutionResult, Output, SpecId, U256};
 use revm::{Database, DatabaseCommit};
 use sov_modules_api::WorkingSet;
@@ -27,7 +28,7 @@ fn simple_contract_execution_sov_state() {
         WorkingSet::new(new_orphan_storage(tmpdir.path()).unwrap());
 
     let evm = Evm::<C>::default();
-    let evm_db: EvmDb<'_, C> = evm.get_db(&mut working_set);
+    let evm_db: EvmDb<'_, C> = evm.get_db(&mut working_set, SHANGHAI);
 
     simple_contract_execution(evm_db);
 }

crates/evm/src/genesis.rs

@@ -162,7 +162,29 @@ impl<C: sov_modules_api::Context> Evm<C> {
         config: &<Self as sov_modules_api::Module>::Config,
         working_set: &mut WorkingSet<C>,
     ) -> Result<()> {
-        let mut evm_db = self.get_db(working_set);
+        let mut spec = config
+            .spec
+            .iter()
+            .map(|(k, v)| {
+                // https://github.com/Sovereign-Labs/sovereign-sdk/issues/912
+                if *v == SpecId::CANCUN {
+                    panic!("Cancun is not supported");
+                }
+
+                (*k, *v)
+            })
+            .collect::<Vec<_>>();
+
+        spec.sort_by(|a, b| a.0.cmp(&b.0));
+
+        if spec.is_empty() {
+            spec.push((0, SpecId::SHANGHAI));
+        } else if spec[0].0 != 0u64 {
+            panic!("EVM spec must start from block 0");
+        }
+
+        let (_, current_spec) = spec.last().expect("Spec should be set");
+        let mut evm_db = self.get_db(working_set, *current_spec);
 
         for acc in &config.data {
             let code = Bytecode::new_raw(acc.code.clone());
@@ -190,27 +212,6 @@ impl<C: sov_modules_api::Context> Evm<C> {
             }
         }
 
-        let mut spec = config
-            .spec
-            .iter()
-            .map(|(k, v)| {
-                // https://github.com/Sovereign-Labs/sovereign-sdk/issues/912
-                if *v == SpecId::CANCUN {
-                    panic!("Cancun is not supported");
-                }
-
-                (*k, *v)
-            })
-            .collect::<Vec<_>>();
-
-        spec.sort_by(|a, b| a.0.cmp(&b.0));
-
-        if spec.is_empty() {
-            spec.push((0, SpecId::SHANGHAI));
-        } else if spec[0].0 != 0u64 {
-            panic!("EVM spec must start from block 0");
-        }
-
         let chain_cfg = EvmChainConfig {
             chain_id: config.chain_id,
             limit_contract_code_size: config.limit_contract_code_size,

crates/evm/src/lib.rs

@@ -80,7 +80,14 @@ pub struct Evm<C: sov_modules_api::Context> {
 
     /// Mapping from code hash to code. Used for lazy-loading code into a contract account.
     #[state(rename = "c")]
-    pub(crate) code: sov_modules_api::OffchainStateMap<
+    pub(crate) code:
+        sov_modules_api::StateMap<reth_primitives::B256, revm::primitives::Bytecode, BcsCodec>,
+
+    /// Mapping from code hash to code. Used for lazy-loading code into a contract account.
+    /// This is the new offchain version which is not counted in the state diff.
+    /// Activated after FORK1
+    #[state(rename = "occ")]
+    pub(crate) offchain_code: sov_modules_api::OffchainStateMap<
         reth_primitives::B256,
         revm::primitives::Bytecode,
         BcsCodec,
@@ -189,12 +196,18 @@ impl<C: sov_modules_api::Context> sov_modules_api::Module for Evm<C> {
 }
 
 impl<C: sov_modules_api::Context> Evm<C> {
-    pub(crate) fn get_db<'a>(&self, working_set: &'a mut WorkingSet<C>) -> EvmDb<'a, C> {
+    pub(crate) fn get_db<'a>(
+        &self,
+        working_set: &'a mut WorkingSet<C>,
+        current_spec: SpecId,
+    ) -> EvmDb<'a, C> {
         EvmDb::new(
             self.accounts.clone(),
             self.code.clone(),
+            self.offchain_code.clone(),
             self.latest_block_hashes.clone(),
             working_set,
+            current_spec,
         )
     }