use existing deployment nonce during storage migration

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

@@ -2,7 +2,11 @@
 
 use crate::{config::*, test_helpers::TEST_DATA_DEFAULT};
 use forge::revm::primitives::SpecId;
-use foundry_test_utils::Filter;
+use foundry_test_utils::{
+    forgetest_async,
+    util::{self, OutputExt},
+    Filter, Fork, ZkSyncNode,
+};
 
 #[tokio::test(flavor = "multi_thread")]
 async fn test_zk_setup_fork_failure() {
@@ -28,3 +32,59 @@ 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| {
+    util::initialize(prj.root());
+
+    // Has deployment nonce (1) and transaction nonce (2) on mainnet block #55159219
+    let test_address = "0x076d6da60aAAC6c97A8a0fE8057f9564203Ee545";
+
+    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 = 2;
+    uint128 constant TEST_ADDRESS_DEPLOYMENT_NONCE = 1;
+
+    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();
+
+    let node = ZkSyncNode::start_with_fork(Fork::new_with_block(
+        String::from("https://mainnet.era.zksync.io"),
+        55159219,
+    ))
+    .await;
+
+    cmd.arg("script").args([
+        "ZkForkNonceTest",
+        "--zk-startup",
+        "./script/ForkNonce.s.sol",
+        "--rpc-url",
+        node.url().as_str(),
+        "--sender",
+        test_address,
+    ]);
+
+    cmd.assert_success();
+});

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

@@ -38,7 +38,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 +883,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 +912,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 +954,36 @@ impl ZksyncCheatcodeInspectorStrategyRunner {
         for address in data.db.persistent_accounts().into_iter().chain([data.env.tx.caller]) {
             info!(?address, "importing to zk state");
 
+            // Re-use 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 = value.original_value.to_u256();
+                    let (_tx_nonce, deployment_nonce) = decompose_full_nonce(full_nonce);
+                    debug!(?address, ?deployment_nonce, "reuse existing deployment nonce");
+                    deployment_nonce
+                } 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,14 +115,35 @@ 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 {
+        Fork { url, ..Default::default() }
+    }
+
+    /// Create a fork config with the provided url and block.
+    pub fn new_with_block(url: String, block: u64) -> Self {
+        Fork { url, block: Some(block) }
+    }
+}
+
 /// In-memory anvil-zksync that is stopped when dropped.
 pub struct ZkSyncNode {
     port: u16,
     _guard: tokio::sync::oneshot::Sender<()>,
 }
 
 impl ZkSyncNode {
-    /// Returns the server url.
+    /// Start anvil-zksync in memory, binding a random available port
+    ///
+    /// The server is automatically stopped when the instance is dropped.
     #[inline]
     pub fn url(&self) -> String {
         format!("http://127.0.0.1:{}", self.port)
@@ -128,14 +153,30 @@ impl ZkSyncNode {
     ///
     /// 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 +188,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 +201,7 @@ impl ZkSyncNode {
         let block_sealer = BlockSealer::new(sealing_mode);
 
         let node: InMemoryNode = InMemoryNode::new(
-            None,
+            fork,
             None,
             &config,
             time,