fix: use existing deployment nonce during storage migration (#895)

* use existing deployment nonce during storage migration

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

@@ -1,8 +1,15 @@
 //! Fork tests.
 
 use crate::{config::*, test_helpers::TEST_DATA_DEFAULT};
+use alloy_provider::Provider;
 use forge::revm::primitives::SpecId;
-use foundry_test_utils::Filter;
+use foundry_common::provider::try_get_zksync_http_provider;
+use foundry_test_utils::{
+    forgetest_async,
+    util::{self},
+    Filter, ZkSyncNode,
+};
+use foundry_zksync_core::state::{get_nonce_storage, new_full_nonce};
 
 #[tokio::test(flavor = "multi_thread")]
 async fn test_zk_setup_fork_failure() {
@@ -28,3 +35,70 @@ async fn test_zk_consistent_storage_migration_after_fork() {
 
     TestConfig::with_filter(runner, filter).spec_id(SpecId::SHANGHAI).run().await;
 }
+
+forgetest_async!(test_zk_consistent_nonce_migration_after_fork, |prj, cmd| {
+    let test_address = alloy_primitives::address!("076d6da60aAAC6c97A8a0fE8057f9564203Ee545");
+    let transaction_nonce = 2;
+    let deployment_nonce = 1;
+
+    let node = ZkSyncNode::start().await;
+    // set nonce
+    let (nonce_key_addr, nonce_key_slot) = get_nonce_storage(test_address);
+    let full_nonce = new_full_nonce(transaction_nonce, deployment_nonce);
+    let result = try_get_zksync_http_provider(node.url())
+        .unwrap()
+        .raw_request::<_, bool>(
+            "anvil_setStorageAt".into(),
+            (nonce_key_addr, nonce_key_slot, full_nonce),
+        )
+        .await
+        .unwrap();
+    assert!(result, "failed setting nonce on anvil-zksync");
+
+    // prepare script
+    util::initialize(prj.root());
+    prj.add_script(
+        "ZkForkNonceTest.s.sol",
+        format!(r#"
+import "forge-std/Script.sol";
+import "forge-std/Test.sol";
+
+interface VmExt {{
+    function zkGetTransactionNonce(
+        address account
+    ) external view returns (uint64 nonce);
+    function zkGetDeploymentNonce(
+        address account
+    ) external view returns (uint64 nonce);
+}}
+
+contract ZkForkNonceTest is Script {{
+    VmExt internal constant vmExt = VmExt(VM_ADDRESS);
+
+    address constant TEST_ADDRESS = {test_address};
+    uint128 constant TEST_ADDRESS_TRANSACTION_NONCE = {transaction_nonce};
+    uint128 constant TEST_ADDRESS_DEPLOYMENT_NONCE = {deployment_nonce};
+
+    function run() external {{
+        require(TEST_ADDRESS_TRANSACTION_NONCE == vmExt.zkGetTransactionNonce(TEST_ADDRESS), "failed matching transaction nonce");
+        require(TEST_ADDRESS_DEPLOYMENT_NONCE == vmExt.zkGetDeploymentNonce(TEST_ADDRESS), "failed matching deployment nonce");
+    }}
+}}
+"#).as_str(),
+    )
+    .unwrap();
+
+    cmd.arg("script").args([
+        "ZkForkNonceTest",
+        "--zk-startup",
+        "./script/ForkNonce.s.sol",
+        "--no-storage-caching", // prevents rpc caching
+        "--rpc-url",
+        node.url().as_str(),
+        // set address as sender to be migrated on startup, so storage is read immediately
+        "--sender",
+        test_address.to_string().as_str(),
+    ]);
+
+    cmd.assert_success();
+});

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

@@ -51,7 +51,7 @@ async fn test_zk_logs_work_in_create() {
                 Some(vec![
                     "print".into(),
                     "outer print".into(),
-                    "0xF9E9ba9Ed9B96AB918c74B21dD0f1D5f2ac38a30".into(),
+                    "0xB5c1DF089600415B21FB76bf89900Adb575947c8".into(),
                     "print".into(),
                     "0xff".into(),
                     "print".into(),

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

@@ -14,10 +14,10 @@ use itertools::Itertools;
 use serde::Deserialize;
 
 const ADDRESS_ZK_TRACE_TEST: Address = address!("7fa9385be102ac3eac297483dd6233d62b3e1496");
-const ADDRESS_ADDER: Address = address!("f9e9ba9ed9b96ab918c74b21dd0f1d5f2ac38a30");
-const ADDRESS_NUMBER: Address = address!("f232f12e115391c535fd519b00efadf042fc8be5");
-const ADDRESS_FIRST_INNER_NUMBER: Address = address!("ed570f3f91621894e001df0fb70bfbd123d3c8ad");
-const ADDRESS_SECOND_INNER_NUMBER: Address = address!("abceaeac3d3a2ac3dcffd7a60ca00a3fac9490ca");
+const ADDRESS_ADDER: Address = address!("b5c1df089600415b21fb76bf89900adb575947c8");
+const ADDRESS_NUMBER: Address = address!("d6a7a38ee698efae2f48f3a62dc7a71c3c0930a1");
+const ADDRESS_FIRST_INNER_NUMBER: Address = address!("89c74b24fb24dda42a8465ee0f9ede2c1308deeb");
+const ADDRESS_SECOND_INNER_NUMBER: Address = address!("9359008843d2c083a14e9c17cde01893938047fa");
 const ADDRESS_CONSOLE: Address = address!("000000000000000000636f6e736f6c652e6c6f67");
 const SELECTOR_TEST_CALL: Bytes = Bytes::from_static(hex!("0d3282c4").as_slice());
 const SELECTOR_TEST_CREATE: Bytes = Bytes::from_static(hex!("61bdc916").as_slice());

crates/strategy/zksync/src/cheatcode/runner/mod.rs

@@ -22,10 +22,11 @@ use foundry_evm::{
 use foundry_evm_core::backend::DatabaseExt;
 use foundry_zksync_core::{
     convert::{ConvertAddress, ConvertH160, ConvertH256, ConvertRU256, ConvertU256},
-    get_account_code_key, get_balance_key, get_nonce_key, PaymasterParams, ZkTransactionMetadata,
-    ACCOUNT_CODE_STORAGE_ADDRESS, CONTRACT_DEPLOYER_ADDRESS, DEFAULT_CREATE2_DEPLOYER_ZKSYNC,
-    KNOWN_CODES_STORAGE_ADDRESS, L2_BASE_TOKEN_ADDRESS, NONCE_HOLDER_ADDRESS,
-    ZKSYNC_TRANSACTION_OTHER_FIELDS_KEY,
+    get_account_code_key, get_balance_key, get_nonce_key,
+    state::parse_full_nonce,
+    PaymasterParams, ZkTransactionMetadata, ACCOUNT_CODE_STORAGE_ADDRESS,
+    CONTRACT_DEPLOYER_ADDRESS, DEFAULT_CREATE2_DEPLOYER_ZKSYNC, KNOWN_CODES_STORAGE_ADDRESS,
+    L2_BASE_TOKEN_ADDRESS, NONCE_HOLDER_ADDRESS, ZKSYNC_TRANSACTION_OTHER_FIELDS_KEY,
 };
 use itertools::Itertools;
 use revm::{
@@ -38,7 +39,7 @@ use revm::{
         SignedAuthorization, KECCAK_EMPTY,
     },
 };
-use tracing::{debug, error, info};
+use tracing::{debug, error, info, trace, warn};
 use zksync_types::{
     block::{pack_block_info, unpack_block_info},
     utils::{decompose_full_nonce, nonces_to_full_nonce},
@@ -883,7 +884,10 @@ impl ZksyncCheatcodeInspectorStrategyRunner {
 
             let balance = data.sload(balance_account, balance_key).unwrap_or_default().data;
             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 (tx_nonce, deployment_nonce) = decompose_full_nonce(full_nonce.to_u256());
+            if !deployment_nonce.is_zero() {
+                warn!(?address, ?deployment_nonce, "discarding ZKsync deployment nonce for EVM context, might cause inconsistencies");
+            }
             let nonce = tx_nonce.as_u64();
 
             let account_code_key = get_account_code_key(address);
@@ -909,7 +913,7 @@ impl ZksyncCheatcodeInspectorStrategyRunner {
             let _ = std::mem::replace(&mut account.info.nonce, nonce);
 
             if test_contract.map(|addr| addr == address).unwrap_or_default() {
-                tracing::trace!(?address, "ignoring code translation for test contract");
+                trace!(?address, "ignoring code translation for test contract");
             } else {
                 account.info.code_hash = code_hash;
                 account.info.code.clone_from(&code);
@@ -951,21 +955,39 @@ impl ZksyncCheatcodeInspectorStrategyRunner {
         for address in data.db.persistent_accounts().into_iter().chain([data.env.tx.caller]) {
             info!(?address, "importing to zk state");
 
+            // Reuse the deployment nonce from storage if present.
+            let deployment_nonce = {
+                let nonce_key = get_nonce_key(address);
+                let nonce_addr = NONCE_HOLDER_ADDRESS.to_address();
+                let account = journaled_account(data, nonce_addr).expect("failed to load account");
+                if let Some(value) = account.storage.get(&nonce_key) {
+                    let full_nonce = parse_full_nonce(value.original_value);
+                    debug!(
+                        ?address,
+                        deployment_nonce = full_nonce.deploy_nonce,
+                        "reuse existing deployment nonce"
+                    );
+                    full_nonce.deploy_nonce.into()
+                } else {
+                    zksync_types::U256::zero()
+                }
+            };
+
             let account = journaled_account(data, address).expect("failed to load account");
             let info = &account.info;
 
             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());
+            debug!(?address, ?deployment_nonce, transaction_nonce=?info.nonce, "attempting to fit EVM nonce to ZKsync nonces, might cause inconsistencies");
+            let full_nonce = nonces_to_full_nonce(info.nonce.into(), deployment_nonce);
 
             let nonce_key = get_nonce_key(address);
             nonce_storage.insert(nonce_key, EvmStorageSlot::new(full_nonce.to_ru256()));
 
             if test_contract.map(|test_address| address == test_address).unwrap_or_default() {
                 // avoid migrating test contract code
-                tracing::trace!(?address, "ignoring code translation for test contract");
+                trace!(?address, "ignoring code translation for test contract");
                 continue;
             }
 

crates/test-utils/src/lib.rs

@@ -31,7 +31,7 @@ pub use script::{ScriptOutcome, ScriptTester};
 
 // TODO: remove once anvil supports zksync node
 mod zksync;
-pub use zksync::ZkSyncNode;
+pub use zksync::{Fork, ZkSyncNode};
 
 // re-exports for convenience
 pub use foundry_compilers;

crates/test-utils/src/zksync.rs

@@ -2,8 +2,12 @@
 use std::{net::SocketAddr, str::FromStr};
 
 use anvil_zksync_api_server::NodeServerBuilder;
-use anvil_zksync_config::{types::SystemContractsOptions, TestNodeConfig};
+use anvil_zksync_config::{
+    types::{CacheConfig, SystemContractsOptions},
+    TestNodeConfig,
+};
 use anvil_zksync_core::{
+    fork::ForkDetails,
     node::{
         BlockProducer, BlockSealer, BlockSealerMode, ImpersonationManager, InMemoryNode,
         TimestampManager, TxPool,
@@ -111,6 +115,25 @@ const RICH_WALLETS: [(&str, &str, &str); 10] = [
     ),
 ];
 
+/// Represents fork config for [ZkSyncNode].
+#[derive(Debug, Default)]
+pub struct Fork {
+    url: String,
+    block: Option<u64>,
+}
+
+impl Fork {
+    /// Create a fork config with the provided url and the latest block.
+    pub fn new(url: String) -> Self {
+        Self { url, ..Default::default() }
+    }
+
+    /// Create a fork config with the provided url and block.
+    pub fn new_with_block(url: String, block: u64) -> Self {
+        Self { url, block: Some(block) }
+    }
+}
+
 /// In-memory anvil-zksync that is stopped when dropped.
 pub struct ZkSyncNode {
     port: u16,
@@ -124,18 +147,34 @@ impl ZkSyncNode {
         format!("http://127.0.0.1:{}", self.port)
     }
 
-    /// Start anvil-zksync in memory, binding a random available port
+    /// Start anvil-zksync in memory, binding a random available port.
     ///
     /// The server is automatically stopped when the instance is dropped.
     pub async fn start() -> Self {
+        Self::start_inner(None).await
+    }
+
+    /// Start anvil-zksync in memory, binding a random available port and with the provided fork url
+    /// and block.
+    ///
+    /// The server is automatically stopped when the instance is dropped.
+    pub async fn start_with_fork(fork: Fork) -> Self {
+        Self::start_inner(Some(fork)).await
+    }
+
+    async fn start_inner(fork: Option<Fork>) -> Self {
         let (_guard, guard_rx) = tokio::sync::oneshot::channel::<()>();
         let (port_tx, port) = tokio::sync::oneshot::channel();
 
+        let fork = fork.map(|fork| {
+            ForkDetails::from_url(fork.url, fork.block, CacheConfig::Memory)
+                .expect("failed building ForkDetails")
+        });
         std::thread::spawn(move || {
             // We need to spawn a thread since `run_inner` future is not `Send`.
             let runtime =
                 tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap();
-            runtime.block_on(Self::run_inner(port_tx, guard_rx));
+            runtime.block_on(Self::run_inner(port_tx, guard_rx, fork));
         });
 
         // wait for server to start
@@ -147,6 +186,7 @@ impl ZkSyncNode {
     async fn run_inner(
         port_tx: tokio::sync::oneshot::Sender<u16>,
         stop_guard: tokio::sync::oneshot::Receiver<()>,
+        fork: Option<ForkDetails>,
     ) {
         const MAX_TRANSACTIONS: usize = 100; // Not that important for testing purposes.
 
@@ -159,7 +199,7 @@ impl ZkSyncNode {
         let block_sealer = BlockSealer::new(sealing_mode);
 
         let node: InMemoryNode = InMemoryNode::new(
-            None,
+            fork,
             None,
             &config,
             time,