Fix contract storage migration (#500)

Author: smiasojed

crates/evm/core/src/backend/cow.rs

@@ -291,6 +291,10 @@ impl DatabaseExt for CowBackend<'_> {
         self.backend.cached_accounts()
     }
 
+    fn cached_storage(&self, address: Address) -> Option<Map<U256, U256>> {
+        self.backend.cached_storage(address)
+    }
+
     fn set_blockhash(&mut self, block_number: U256, block_hash: B256) {
         self.backend.to_mut().set_blockhash(block_number, block_hash);
     }

crates/evm/core/src/backend/mod.rs

@@ -356,6 +356,11 @@ pub trait DatabaseExt: Database<Error = DatabaseError> + DatabaseCommit + Debug
     /// Returns all accounts in the memory database cache
     fn cached_accounts(&self) -> Vec<Address>;
 
+    /// Returns the cached storage for an account from the memory database cache.
+    /// This returns storage that was written during contract construction before
+    /// startup migration, which may not yet be in the journaled state.
+    fn cached_storage(&self, address: Address) -> Option<Map<U256, U256>>;
+
     /// Ensures that `account` is allowed to execute cheatcodes
     ///
     /// Returns an error if [`Self::has_cheatcode_access`] returns `false`
@@ -1535,6 +1540,14 @@ impl DatabaseExt for Backend {
             .collect()
     }
 
+    fn cached_storage(&self, address: Address) -> Option<Map<U256, U256>> {
+        self.mem_db
+            .cache
+            .accounts
+            .get(&address)
+            .map(|acc| acc.storage.iter().map(|(k, v)| (*k, *v)).collect())
+    }
+
     fn set_blockhash(&mut self, block_number: U256, block_hash: B256) {
         if let Some(db) = self.active_fork_db_mut() {
             db.cache.block_hashes.insert(block_number.saturating_to(), block_hash);

crates/forge/tests/it/revive/migration.rs

@@ -176,3 +176,14 @@ async fn test_initial_contract_deployment(#[case] runtime_mode: ReviveRuntimeMod
         Filter::new("testInitialContractMigration", "InitialMigrationTest", ".*/revive/.*");
     TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
 }
+
+#[rstest]
+#[case::pvm(ReviveRuntimeMode::Pvm)]
+#[case::evm(ReviveRuntimeMode::Evm)]
+#[tokio::test(flavor = "multi_thread")]
+async fn test_contract_with_constructor_deployment(#[case] runtime_mode: ReviveRuntimeMode) {
+    let runner = TEST_DATA_REVIVE.runner_revive(runtime_mode);
+    let filter =
+        Filter::new("testConstructorArgPreserved", "ContractArgsMigrationTest", ".*/revive/.*");
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}

crates/revive-strategy/src/cheatcodes/mod.rs

@@ -354,15 +354,15 @@ impl CheatcodeInspectorStrategyRunner for PvmCheatcodeInspectorStrategyRunner {
             }
             t if using_revive && is::<rollCall>(t) => {
                 let &rollCall { newHeight } = cheatcode.as_any().downcast_ref().unwrap();
-                let new_block_number: u64 = newHeight.try_into().expect("Block number exceeds u32");
+                let new_block_number: u64 = newHeight.saturating_to();
 
                 // blockhash should be the same on both revive and revm sides, so fetch it before
                 // changing the block number.
                 let prev_new_height_hash = ccx
                     .ecx
                     .journaled_state
                     .database
-                    .block_hash(new_block_number - 1)
+                    .block_hash(new_block_number.saturating_sub(1))
                     .expect("Should not fail");
                 let new_height_hash = ccx
                     .ecx
@@ -695,10 +695,11 @@ fn select_revive(
             }
 
             for address in accounts {
-                tracing::info!("Migrating account {:?} (is_test_contract: {})", address, test_contract_addr == Some(address));
+                tracing::info!("Migrating account {:?} (is_test_contract: {})", address, test_contract_addr == Some(address)); 
                 let acc = data.journaled_state.load_account(address).expect("failed to load account");
-                let amount = acc.data.info.balance;
-                let nonce = acc.data.info.nonce;
+
+                let amount = acc.info.balance;
+                let nonce = acc.info.nonce;
                 let account = H160::from_slice(address.as_slice());
                 let account_id =
                     AccountId32Mapper::<Runtime>::to_fallback_account_id(&account);
@@ -710,7 +711,7 @@ fn select_revive(
                     a.nonce = nonce.min(u32::MAX.into()).try_into().expect("shouldn't happen");
                 });
 
-                if let Some(bytecode) = acc.data.info.code.as_ref() {
+                if let Some(bytecode) = acc.info.code.as_ref() {
                     let account_h160 = H160::from_slice(address.as_slice());
 
                     // Skip if contract already exists in pallet-revive
@@ -825,26 +826,73 @@ fn select_revive(
                             }
                         }
                     }
-                    if  AccountInfo::<Runtime>::load_contract(&account_h160).is_some() {
-                           // Migrate complete account state (storage) for newly created/existing contract
-                           for (slot, storage_slot) in &acc.data.storage {
-                            let slot_bytes = slot.to_be_bytes::<32>();
-                            let value_bytes = storage_slot.present_value.to_be_bytes::<32>();
-
-                            if !storage_slot.present_value.is_zero() {
-                                let _ = Pallet::<Runtime>::set_storage(
-                                    account_h160,
-                                    slot_bytes,
-                                    Some(value_bytes.to_vec()),
-                                );
-                            }
-                        }
+                    if AccountInfo::<Runtime>::load_contract(&account_h160).is_some() {
+                        migrate_contract_storage(data, address, account_h160);
                     }
                 }
             }
         });
 }
 
+/// Migrates contract storage from REVM state to pallet-revive.
+///
+/// Merges storage from two sources:
+/// 1. Journaled state: most recent storage values from current execution
+/// 2. Database cache: storage written before startup migration, which is run as separate
+///    transaction, already committed to cache
+///
+/// The journaled state takes precedence - cache values are only used for slots
+/// not present in the journaled state.
+fn migrate_contract_storage(data: Ecx<'_, '_, '_>, address: Address, account_h160: H160) {
+    use std::collections::HashSet;
+
+    // Track which slots we've already migrated from the journaled state
+    let mut migrated_slots: HashSet<U256> = HashSet::new();
+
+    // First, migrate storage from the journaled state (most up-to-date values)
+    if let Some(account_state) = data.journaled_state.state.get(&address) {
+        for (slot, storage_slot) in &account_state.storage {
+            migrated_slots.insert(*slot);
+
+            let slot_bytes = slot.to_be_bytes::<32>();
+            let value = storage_slot.present_value;
+
+            if !value.is_zero() {
+                let _ = Pallet::<Runtime>::set_storage(
+                    account_h160,
+                    slot_bytes,
+                    Some(value.to_be_bytes::<32>().to_vec()),
+                );
+            } else {
+                // Handle case where storage was explicitly cleared
+                let _ = Pallet::<Runtime>::set_storage(account_h160, slot_bytes, None);
+            }
+        }
+    }
+
+    // Then, migrate storage from the database cache for slots NOT in journaled state
+    if let Some(cached_storage) = data.journaled_state.database.cached_storage(address) {
+        for (slot, value) in cached_storage {
+            // Skip slots already migrated from the journaled state
+            if migrated_slots.contains(&slot) {
+                continue;
+            }
+
+            let slot_bytes = slot.to_be_bytes::<32>();
+            if !value.is_zero() {
+                let _ = Pallet::<Runtime>::set_storage(
+                    account_h160,
+                    slot_bytes,
+                    Some(value.to_be_bytes::<32>().to_vec()),
+                );
+            } else {
+                // Handle case where storage was explicitly cleared
+                let _ = Pallet::<Runtime>::set_storage(account_h160, slot_bytes, None);
+            }
+        }
+    }
+}
+
 fn select_evm(ctx: &mut PvmCheatcodeInspectorStrategyContext, data: Ecx<'_, '_, '_>) {
     if !ctx.using_revive {
         tracing::info!("already using REVM");
@@ -991,7 +1039,7 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
                 (contract.resolc_bytecode.as_bytes().unwrap().to_vec(), constructor_args.to_vec())
             }
             crate::ReviveRuntimeMode::Evm => {
-                // EVM mode: use EVM bytecode directly
+                // EVM mode: use EVM bytecode directly (includes constructor args)
                 tracing::info!("running create in EVM mode with EVM bytecode");
                 (init_code.0.to_vec(), vec![])
             }

crates/revive-strategy/src/state.rs

@@ -115,7 +115,7 @@ impl TestEnv {
     ) {
         // Set block number in pallet-revive runtime.
         self.0.lock().unwrap().externalities.execute_with(|| {
-            let new_block_number: u64 = new_height.try_into().expect("Block number exceeds u64");
+            let new_block_number: u64 = new_height.saturating_to();
             let digest = System::digest();
             if System::block_hash(new_block_number) == H256::zero() {
                 // First initialize and finalize the parent block to set up correct hashes.

testdata/default/revive/EvmToReviveMigration.t.sol

@@ -384,3 +384,30 @@ contract InitialMigrationTest is DSTest {
         assertEq(storageContract.get(), 42);
     }
 }
+
+// Simple contract that stores a constructor argument
+contract SimpleAddrStorage {
+    address public storedAddr;
+
+    constructor(address _addr) {
+        storedAddr = _addr;
+    }
+}
+
+contract ContractArgsMigrationTest is DSTest {
+    Vm constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));
+    address myAddr = vm.addr(1900);
+
+    // Expected address for private key 1900
+    address constant EXPECTED = 0x6Af741AA4Ff39CF3De9A0Cb02A8Bab387E41abFB;
+    SimpleAddrStorage store = new SimpleAddrStorage(myAddr);
+
+    function testConstructorArgPreserved() public {
+        emit log_named_address("myAddr (from vm.addr)", myAddr);
+        emit log_named_address("store.storedAddr()", store.storedAddr());
+        emit log_named_address("EXPECTED", EXPECTED);
+
+        assertEq(myAddr, EXPECTED, "vm.addr(1900) should return correct address");
+        assertEq(store.storedAddr(), myAddr, "Constructor arg should be preserved");
+    }
+}