fix(zk): merge code storage & nonce with zk-zk

crates/cheatcodes/src/inspector.rs

@@ -33,8 +33,8 @@ use foundry_evm_core::{
 };
 use foundry_zksync_compiler::{DualCompiledContract, DualCompiledContracts};
 use foundry_zksync_core::{
-    convert::{ConvertAddress, ConvertH160, ConvertH256, ConvertRU256, ConvertU256},
-    ZkTransactionMetadata,
+    convert::{ConvertH160, ConvertH256, ConvertRU256, ConvertU256},
+    get_account_code_key, get_balance_key, get_nonce_key, ZkTransactionMetadata,
 };
 use itertools::Itertools;
 use revm::{
@@ -60,8 +60,7 @@ use std::{
 };
 use zksync_types::{
     block::{pack_block_info, unpack_block_info},
-    get_code_key, get_nonce_key,
-    utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
+    utils::{decompose_full_nonce, nonces_to_full_nonce},
     ACCOUNT_CODE_STORAGE_ADDRESS, CONTRACT_DEPLOYER_ADDRESS, CURRENT_VIRTUAL_BLOCK_INFO_POSITION,
     H256, KNOWN_CODES_STORAGE_ADDRESS, L2_BASE_TOKEN_ADDRESS, NONCE_HOLDER_ADDRESS,
     SYSTEM_CONTEXT_ADDRESS,
@@ -281,14 +280,22 @@ impl Cheatcodes {
             evm_deployed_bytecode: Bytecode::new_raw(empty_bytes.clone()).bytecode().to_vec(),
             evm_bytecode: Bytecode::new_raw(empty_bytes.clone()).bytecode().to_vec(),
         });
+
+        let cheatcodes_bytecode = {
+            let mut bytecode = CHEATCODE_ADDRESS.abi_encode_packed();
+            bytecode.append(&mut [0; 12].to_vec());
+            Bytes::from(bytecode)
+        };
         dual_compiled_contracts.push(DualCompiledContract {
             name: String::from("CheatcodeBytecode"),
-            zk_bytecode_hash,
-            zk_deployed_bytecode: zk_deployed_bytecode.clone(),
+            // we put a different bytecode hash here so when importing back to EVM
+            // we avoid collision with EmptyEVMBytecode for the cheatcodes
+            zk_bytecode_hash: foundry_zksync_core::hash_bytecode(CHEATCODE_CONTRACT_HASH.as_ref()),
+            zk_deployed_bytecode: cheatcodes_bytecode.to_vec(),
             zk_factory_deps: Default::default(),
             evm_bytecode_hash: CHEATCODE_CONTRACT_HASH,
-            evm_deployed_bytecode: Bytecode::new_raw(empty_bytes.clone()).bytecode().to_vec(),
-            evm_bytecode: Bytecode::new_raw(empty_bytes).bytecode().to_vec(),
+            evm_deployed_bytecode: cheatcodes_bytecode.to_vec(),
+            evm_bytecode: cheatcodes_bytecode.to_vec(),
         });
 
         let mut persisted_factory_deps = HashMap::new();
@@ -476,16 +483,15 @@ impl Cheatcodes {
         for address in data.db.persistent_accounts().into_iter().chain([data.env.tx.caller]) {
             info!(?address, "importing to evm state");
 
-            let zk_address = address.to_h160();
-            let balance_key = storage_key_for_eth_balance(&zk_address).key().to_ru256();
-            let nonce_key = get_nonce_key(&zk_address).key().to_ru256();
+            let balance_key = get_balance_key(address);
+            let nonce_key = get_nonce_key(address);
 
             let (balance, _) = data.sload(balance_account, balance_key).unwrap_or_default();
             let (full_nonce, _) = data.sload(nonce_account, nonce_key).unwrap_or_default();
             let (tx_nonce, _deployment_nonce) = decompose_full_nonce(full_nonce.to_u256());
             let nonce = tx_nonce.as_u64();
 
-            let account_code_key = get_code_key(&zk_address).key().to_ru256();
+            let account_code_key = get_account_code_key(address);
             let (code_hash, code) = data
                 .sload(account_code_account, account_code_key)
                 .map(|(value, _)| value)
@@ -551,21 +557,21 @@ impl Cheatcodes {
 
             let account = journaled_account(data, address).expect("failed to load account");
             let info = &account.info;
-            let zk_address = address.to_h160();
 
-            let balance_key = storage_key_for_eth_balance(&zk_address).key().to_ru256();
-            let nonce_key = get_nonce_key(&zk_address).key().to_ru256();
+            let balance_key = get_balance_key(address);
             l2_eth_storage.insert(balance_key, EvmStorageSlot::new(info.balance));
 
             // TODO we need to find a proper way to handle deploy nonces instead of replicating
             let full_nonce = nonces_to_full_nonce(info.nonce.into(), info.nonce.into());
+
+            let nonce_key = get_nonce_key(address);
             nonce_storage.insert(nonce_key, EvmStorageSlot::new(full_nonce.to_ru256()));
 
             if let Some(contract) = self.dual_compiled_contracts.iter().find(|contract| {
                 info.code_hash != KECCAK_EMPTY && info.code_hash == contract.evm_bytecode_hash
             }) {
                 account_code_storage.insert(
-                    zk_address.to_h256().to_ru256(),
+                    get_account_code_key(address),
                     EvmStorageSlot::new(contract.zk_bytecode_hash.to_ru256()),
                 );
                 known_codes_storage
@@ -584,7 +590,7 @@ impl Cheatcodes {
                     },
                 );
             } else {
-                tracing::debug!("no zk contract found for {:?}", info.code_hash)
+                tracing::debug!(code_hash = ?info.code_hash, ?address, "no zk contract found")
             }
         }
 

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

@@ -37,12 +37,12 @@ impl CachedForkType {
 
         let is_zk_url = foundry_common::provider::try_get_http_provider(fork_url)
             .map(|provider| {
-                let is_zk_url =
-                    futures::executor::block_on(provider.raw_request("zks_L1ChainId".into(), ()))
-                        .map(|_: String| true)
-                        .unwrap_or_default();
-
-                is_zk_url
+                tokio::task::block_in_place(move || {
+                    tokio::runtime::Handle::current()
+                        .block_on(provider.raw_request::<_, String>("zks_L1ChainId".into(), ()))
+                        .map(|_| true)
+                })
+                .unwrap_or_default()
             })
             .unwrap_or_default();
 

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

@@ -13,7 +13,9 @@ use alloy_rpc_types::{Block, BlockNumberOrTag, BlockTransactions, Transaction};
 use alloy_serde::WithOtherFields;
 use eyre::Context;
 use foundry_common::{is_known_system_sender, SYSTEM_TRANSACTION_TYPE};
-use foundry_zksync_core::{convert::ConvertH160, L2_BASE_TOKEN_ADDRESS};
+use foundry_zksync_core::{
+    convert::ConvertH160, ACCOUNT_CODE_STORAGE_ADDRESS, L2_BASE_TOKEN_ADDRESS, NONCE_HOLDER_ADDRESS,
+};
 use itertools::Itertools;
 use revm::{
     db::{CacheDB, DatabaseRef},
@@ -1944,8 +1946,6 @@ fn merge_db_account_data<ExtDB: DatabaseRef>(
     active: &CacheDB<ExtDB>,
     fork_db: &mut ForkDB,
 ) {
-    trace!(?addr, "merging database data");
-
     let mut acc = if let Some(acc) = active.accounts.get(&addr).cloned() {
         acc
     } else {
@@ -1973,33 +1973,43 @@ fn merge_zk_account_data<ExtDB: DatabaseRef>(
     active: &CacheDB<ExtDB>,
     fork_db: &mut ForkDB,
 ) {
-    trace!(?addr, "merging zk database data");
+    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;
+        };
 
-    // TODO: do the same for nonce and codes
-    let balance_addr = L2_BASE_TOKEN_ADDRESS.to_address();
-    let mut acc = if let Some(acc) = active.accounts.get(&balance_addr).cloned() {
-        acc
-    } else {
-        // Account does not exist
-        return;
-    };
+        let mut storage = Map::<U256, U256>::default();
+        if let Some(value) = acc.storage.get(&slot) {
+            storage.insert(slot, *value);
+        }
 
-    let mut balances = Map::<U256, U256>::default();
-    let slot = foundry_zksync_core::get_balance_key(addr);
-    if let Some(value) = acc.storage.get(&slot) {
-        balances.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)
+        }
 
-    if let Some(fork_account) = fork_db.accounts.get_mut(&balance_addr) {
-        // This will merge the fork's tracked storage with active storage and update values
-        fork_account.storage.extend(balances);
-        // 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 balances, &mut acc.storage)
-    }
+        fork_db.accounts.insert(system_contract, acc);
+    };
 
-    fork_db.accounts.insert(balance_addr, acc);
+    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),
+    );
 }
 
 /// Returns true of the address is a contract

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

@@ -0,0 +1,15 @@
+//! Forge tests for zkysnc issues, to avoid regressions.
+//!
+//! Issue list: https://github.com/matter-labs/foundry-zksync/issues
+
+use crate::{config::*, test_helpers::TEST_DATA_DEFAULT};
+use forge::revm::primitives::SpecId;
+use foundry_test_utils::Filter;
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_zk_issue497() {
+    let runner = TEST_DATA_DEFAULT.runner_zksync();
+    let filter = Filter::new(".*", "Issue497", ".*");
+
+    TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
+}

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

@@ -2,4 +2,5 @@
 mod basic;
 mod cheats;
 mod contracts;
+mod issues;
 mod logs;

crates/zksync/core/src/lib.rs

@@ -58,6 +58,16 @@ pub fn get_balance_key(address: Address) -> rU256 {
     storage_key_for_eth_balance(&address.to_h160()).key().to_ru256()
 }
 
+/// Returns the account code storage key for a provided account address.
+pub fn get_account_code_key(address: Address) -> rU256 {
+    zksync_types::get_code_key(&address.to_h160()).key().to_ru256()
+}
+
+/// Returns the account nonce key for a provided account address.
+pub fn get_nonce_key(address: Address) -> rU256 {
+    zksync_types::get_nonce_key(&address.to_h160()).key().to_ru256()
+}
+
 /// Represents additional data for ZK transactions.
 #[derive(Clone, Debug, Default)]
 pub struct ZkTransactionMetadata {

crates/zksync/core/src/vm/storage_view.rs

@@ -76,22 +76,22 @@ impl<S: ReadStorage + fmt::Debug> ReadStorage for StorageView<S> {
         if key.address() == &ACCOUNT_CODE_STORAGE_ADDRESS && key.key() == &self.caller.to_h256() {
             let value = StorageValue::zero();
             tracing::trace!(
-                "override read value {:?} {:?} ({:?}/{:?})",
-                key.hashed_key(),
-                value,
-                key.address(),
-                key.key()
+                hashed_key = ?key.hashed_key(),
+                ?value,
+                address = ?key.address(),
+                key = ?key.key(),
+                "override read value",
             );
 
             return value
         }
 
         tracing::trace!(
-            "read value {:?} {:?} ({:?}/{:?})",
-            key.hashed_key(),
-            value,
-            key.address(),
-            key.key()
+            hashed_key = ?key.hashed_key(),
+            ?value,
+            address = ?key.address(),
+            key = ?key.key(),
+            "read value",
         );
 
         value
@@ -127,12 +127,12 @@ impl<S: ReadStorage + fmt::Debug> WriteStorage for StorageView<S> {
         let original = self.get_value_no_log(&key);
 
         tracing::trace!(
-            "write value {:?} value: {:?} original value: {:?} ({:?}/{:?})",
-            key.hashed_key(),
-            value,
-            original,
-            key.address(),
-            key.key()
+            hashed_key = ?key.hashed_key(),
+            ?value,
+            ?original,
+            address = ?key.address(),
+            key = ?key.key(),
+            "write value",
         );
         self.modified_storage_keys.insert(key, value);
 

testdata/zk/Issues.t.sol

@@ -0,0 +1,25 @@
+// 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";
+
+// https://github.com/matter-labs/foundry-zksync/issues/497
+contract Issue497 is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    uint256 constant ERA_FORK_BLOCK = 19579636;
+
+    uint256 forkEra;
+
+    function setUp() public {
+        forkEra = vm.createFork(Globals.ZKSYNC_MAINNET_URL, ERA_FORK_BLOCK);
+    }
+
+    function testZkEnsureContractMigratedWhenForkZkSyncThenZkVmOff() external {
+        vm.selectFork(forkEra);
+        vm.zkVm(false);
+        assert(address(vm).codehash != 0);
+    }
+}

zk-tests/src/Basic.t.sol

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.13;
 
-import {Test} from "forge-std/Test.sol";
+import "forge-std/Test.sol";
 
 contract BlockEnv {
     uint256 public number;

zk-tests/src/Issues.t.sol

@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity ^0.8.18;
+
+import "forge-std/Test.sol";
+
+// https://github.com/matter-labs/foundry-zksync/issues/497
+contract Issue497 is Test {
+    uint256 constant ERA_FORK_BLOCK = 19579636;
+
+    uint256 forkEra;
+
+    function setUp() public {
+        forkEra = vm.createFork("mainnet", ERA_FORK_BLOCK);
+    }
+
+    function testZkEnsureContractMigratedWhenForkZkSyncThenZkVmOff() external {
+       vm.selectFork(forkEra);
+       (bool success, ) = address(vm).call(
+           abi.encodeWithSignature("zkVm(bool)", false)
+       );
+       assert(address(vm).codehash != 0);
+    }
+}