fix: merge journaled state during forks

.github/workflows/test.yml

@@ -124,7 +124,7 @@ jobs:
           env:
             RUST_BACKTRACE: full
             TEST_MAINNET_URL: http://localhost:8011
-          run: cargo test zk
+          run: git config --global user.name "test-runner" && cargo test zk
 
     check-ci-install:
       name: CI install

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

@@ -29,7 +29,7 @@ use revm::{
     Database, DatabaseCommit, JournaledState,
 };
 use std::{
-    collections::{BTreeMap, HashSet},
+    collections::{hash_map::Entry, BTreeMap, HashSet},
     time::Instant,
 };
 
@@ -1945,6 +1945,13 @@ pub(crate) fn merge_account_data<ExtDB: DatabaseRef>(
             merge_zk_account_data(addr, active, &mut target_fork.db);
         }
         merge_journaled_state_data(addr, active_journaled_state, &mut target_fork.journaled_state);
+        if merge_zk_db {
+            merge_zk_journaled_state_data(
+                addr,
+                active_journaled_state,
+                &mut target_fork.journaled_state,
+            );
+        }
     }
 
     // need to mock empty journal entries in case the current checkpoint is higher than the existing
@@ -1991,7 +1998,6 @@ fn merge_db_account_data<ExtDB: DatabaseRef>(
     }
 
     // port account storage over
-    use std::collections::hash_map::Entry;
     match fork_db.accounts.entry(addr) {
         Entry::Vacant(vacant) => {
             trace!("target account not present - inserting from active");
@@ -2016,28 +2022,79 @@ fn merge_zk_account_data<ExtDB: DatabaseRef>(
     fork_db: &mut ForkDB,
 ) {
     let mut merge_system_contract_entry = |system_contract: Address, slot: U256| {
-        let mut acc = if let Some(acc) = active.accounts.get(&system_contract).cloned() {
-            acc
-        } else {
-            // Account does not exist
-            return;
-        };
+        let Some(acc) = active.accounts.get(&system_contract) else { return };
+
+        // port contract cache over
+        if let Some(code) = active.contracts.get(&acc.info.code_hash) {
+            trace!("merging contract cache");
+            fork_db.contracts.insert(acc.info.code_hash, code.clone());
+        }
 
-        let mut storage = Map::<U256, U256>::default();
+        // prepare only the specified slot in account storage
+        let mut new_acc = acc.clone();
+        new_acc.storage = Default::default();
         if let Some(value) = acc.storage.get(&slot) {
-            storage.insert(slot, *value);
+            new_acc.storage.insert(slot, *value);
         }
 
-        if let Some(fork_account) = fork_db.accounts.get_mut(&system_contract) {
-            // This will merge the fork's tracked storage with active storage and update values
-            fork_account.storage.extend(storage);
-            // swap them so we can insert the account as whole in the next step
-            std::mem::swap(&mut fork_account.storage, &mut acc.storage);
-        } else {
-            std::mem::swap(&mut storage, &mut acc.storage)
+        // port account storage over
+        match fork_db.accounts.entry(system_contract) {
+            Entry::Vacant(vacant) => {
+                trace!("target account not present - inserting from active");
+                // if the fork_db doesn't have the target account
+                // insert the entire thing
+                vacant.insert(new_acc);
+            }
+            Entry::Occupied(mut occupied) => {
+                trace!("target account present - merging storage slots");
+                // if the fork_db does have the system,
+                // extend the existing storage (overriding)
+                let fork_account = occupied.get_mut();
+                fork_account.storage.extend(&new_acc.storage);
+            }
         }
+    };
+
+    merge_system_contract_entry(
+        L2_BASE_TOKEN_ADDRESS.to_address(),
+        foundry_zksync_core::get_balance_key(addr),
+    );
+    merge_system_contract_entry(
+        ACCOUNT_CODE_STORAGE_ADDRESS.to_address(),
+        foundry_zksync_core::get_account_code_key(addr),
+    );
+    merge_system_contract_entry(
+        NONCE_HOLDER_ADDRESS.to_address(),
+        foundry_zksync_core::get_nonce_key(addr),
+    );
+}
+
+/// Clones the account data from the `active_journaled_state` into the `fork_journaled_state` for
+/// zksync storage.
+fn merge_zk_journaled_state_data(
+    addr: Address,
+    active_journaled_state: &JournaledState,
+    fork_journaled_state: &mut JournaledState,
+) {
+    let mut merge_system_contract_entry = |system_contract: Address, slot: U256| {
+        if let Some(acc) = active_journaled_state.state.get(&system_contract) {
+            // prepare only the specified slot in account storage
+            let mut new_acc = acc.clone();
+            new_acc.storage = Default::default();
+            if let Some(value) = acc.storage.get(&slot).cloned() {
+                new_acc.storage.insert(slot, value);
+            }
 
-        fork_db.accounts.insert(system_contract, acc);
+            match fork_journaled_state.state.entry(system_contract) {
+                Entry::Vacant(vacant) => {
+                    vacant.insert(new_acc);
+                }
+                Entry::Occupied(mut occupied) => {
+                    let fork_account = occupied.get_mut();
+                    fork_account.storage.extend(new_acc.storage);
+                }
+            }
+        }
     };
 
     merge_system_contract_entry(

crates/forge/tests/it/zk/repros.rs

@@ -43,3 +43,6 @@ test_repro!(565; |cfg| {
     cfg.runner.test_options.invariant.fail_on_revert = true;
     cfg.runner.test_options.invariant.runs = 2;
 });
+
+// https://github.com/matter-labs/foundry-zksync/issues/687
+test_repro!(687);

testdata/zk/repros/Issue687.sol

@@ -0,0 +1,46 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity ^0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+import {Globals} from "../Globals.sol";
+
+contract Counter {
+    uint256 number;
+
+    function get() public returns (uint256) {
+        return number;
+    }
+
+    function inc() public {
+        number += 1;
+    }
+}
+
+// https://github.com/matter-labs/foundry-zksync/issues/687
+contract Issue687 is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    uint256 constant ERA_FORK_BLOCK = 19579636;
+
+    uint256 forkEra;
+
+    function setUp() public {
+        forkEra = vm.createSelectFork(Globals.ZKSYNC_MAINNET_URL, ERA_FORK_BLOCK);
+    }
+
+    function testZkEnsureContractMigratedWhenForkedIfPersistent() external {
+        Counter counter = new Counter();
+        counter.inc();
+        assertEq(1, counter.get());
+        vm.makePersistent(address(counter));
+        assertTrue(vm.isPersistent(address(counter)));
+
+        vm.createSelectFork(Globals.ZKSYNC_MAINNET_URL, ERA_FORK_BLOCK - 10);
+
+        assertTrue(vm.isPersistent(address(counter)));
+        assertEq(1, counter.get());
+        counter.inc();
+        assertEq(2, counter.get());
+    }
+}