fix: erroneous nonce updates

crates/cheatcodes/spec/src/vm.rs

@@ -908,6 +908,14 @@ interface Vm {
     #[cheatcode(group = Testing, safety = Safe)]
     function zkRegisterContract(string calldata name, bytes32 evmBytecodeHash, bytes calldata evmDeployedBytecode, bytes calldata evmBytecode, bytes32 zkBytecodeHash, bytes calldata zkDeployedBytecode) external pure;
 
+    /// Gets the transaction nonce of a zksync account.
+    #[cheatcode(group = Evm, safety = Safe)]
+    function zkGetTransactionNonce(address account) external view returns (uint64 nonce);
+
+    /// Gets the deployment nonce of a zksync account.
+    #[cheatcode(group = Evm, safety = Safe)]
+    function zkGetDeploymentNonce(address account) external view returns (uint64 nonce);
+
     /// If the condition is false, discard this run's fuzz inputs and generate new ones.
     #[cheatcode(group = Testing, safety = Safe)]
     function assume(bool condition) external pure;

crates/cheatcodes/src/evm.rs

@@ -143,6 +143,20 @@ impl Cheatcode for getNonce_1Call {
     }
 }
 
+impl Cheatcode for zkGetTransactionNonceCall {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { account } = self;
+        get_nonce(ccx, account)
+    }
+}
+
+impl Cheatcode for zkGetDeploymentNonceCall {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { account } = self;
+        get_nonce(ccx, account)
+    }
+}
+
 impl Cheatcode for loadCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { target, slot } = *self;

crates/forge/tests/fixtures/zk/ScriptSetup.s.sol

@@ -4,26 +4,42 @@ pragma solidity ^0.8.13;
 import {Script} from "forge-std/Script.sol";
 import {Greeter} from "../src/Greeter.sol";
 
+interface VmExt {
+    function zkGetTransactionNonce(
+        address account
+    ) external view returns (uint64 nonce);
+    function zkGetDeploymentNonce(
+        address account
+    ) external view returns (uint64 nonce);
+}
+
 contract ScriptSetupNonce is Script {
+    VmExt internal constant vmExt = VmExt(VM_ADDRESS);
+
     function setUp() public {
-        uint256 initial_nonce = checkNonce(address(tx.origin));
+        uint256 initialNonceTx = checkTxNonce(address(tx.origin));
+        uint256 initialNonceDeploy = checkDeployNonce(address(tx.origin));
         // Perform transactions and deploy contracts in setup to increment nonce and verify broadcast nonce matches onchain
         new Greeter();
         new Greeter();
         new Greeter();
         new Greeter();
-        assert(checkNonce(address(tx.origin)) == initial_nonce);
+        assert(checkTxNonce(address(tx.origin)) == initialNonceTx);
+        assert(checkDeployNonce(address(tx.origin)) == initialNonceDeploy);
     }
 
     function run() public {
         // Get initial nonce
-        uint256 initial_nonce = checkNonce(address(tx.origin));
-        assert(initial_nonce == vm.getNonce(address(tx.origin)));
+        uint256 initialNonceTx = checkTxNonce(tx.origin);
+        uint256 initialNonceDeploy = checkDeployNonce(tx.origin);
+        assert(initialNonceTx == vmExt.zkGetTransactionNonce(tx.origin));
+        assert(initialNonceDeploy == vmExt.zkGetDeploymentNonce(tx.origin));
 
         // Create and interact with non-broadcasted contract to verify nonce is not incremented
         Greeter notBroadcastGreeter = new Greeter();
         notBroadcastGreeter.greeting("john");
-        assert(checkNonce(address(tx.origin)) == initial_nonce);
+        assert(checkTxNonce(tx.origin) == initialNonceTx);
+        assert(checkDeployNonce(tx.origin) == initialNonceDeploy);
 
         // Start broadcasting transactions
         vm.startBroadcast();
@@ -33,37 +49,71 @@ contract ScriptSetupNonce is Script {
 
         // Deploy checker and verify nonce
         NonceChecker checker = new NonceChecker();
+
+        vm.stopBroadcast();
+
         // We expect the nonce to be incremented by 1 because the check is done in an external
         // call
-        checker.assertNonce(vm.getNonce(address(tx.origin)) + 1);
-        vm.stopBroadcast();
+        checker.assertTxNonce(
+            vmExt.zkGetTransactionNonce(address(tx.origin)) + 1
+        );
+        checker.assertDeployNonce(
+            vmExt.zkGetDeploymentNonce(address(tx.origin))
+        );
     }
 
-    function checkNonce(address addr) public returns (uint256) {
+    function checkTxNonce(address addr) public returns (uint256) {
         // We prank here to avoid accidentally "polluting" the nonce of `addr` during the call
         // for example when `addr` is `tx.origin`
         vm.prank(address(this), address(this));
-        return NonceLib.getNonce(addr);
+        return NonceLib.getTxNonce(addr);
+    }
+
+    function checkDeployNonce(address addr) public returns (uint256) {
+        // We prank here to avoid accidentally "polluting" the nonce of `addr` during the call
+        // for example when `addr` is `tx.origin`
+        vm.prank(address(this), address(this));
+        return NonceLib.getDeployNonce(addr);
     }
 }
 
 contract NonceChecker {
-    function checkNonce() public returns (uint256) {
-        return NonceLib.getNonce(address(tx.origin));
+    function checkTxNonce() public returns (uint256) {
+        return NonceLib.getTxNonce(address(tx.origin));
+    }
+
+    function checkDeployNonce() public returns (uint256) {
+        return NonceLib.getDeployNonce(address(tx.origin));
+    }
+
+    function assertTxNonce(uint256 expected) public {
+        uint256 real_nonce = checkTxNonce();
+        require(real_nonce == expected, "tx nonce mismatch");
     }
 
-    function assertNonce(uint256 expected) public {
-        uint256 real_nonce = checkNonce();
-        require(real_nonce == expected, "Nonce mismatch");
+    function assertDeployNonce(uint256 expected) public {
+        uint256 real_nonce = checkDeployNonce();
+        require(real_nonce == expected, "deploy nonce mismatch");
     }
 }
 
 library NonceLib {
     address constant NONCE_HOLDER = address(0x8003);
 
     /// Retrieve tx nonce for `addr` from the NONCE_HOLDER system contract
-    function getNonce(address addr) internal returns (uint256) {
-        (bool success, bytes memory data) = NONCE_HOLDER.call(abi.encodeWithSignature("getMinNonce(address)", addr));
+    function getTxNonce(address addr) internal returns (uint256) {
+        (bool success, bytes memory data) = NONCE_HOLDER.call(
+            abi.encodeWithSignature("getMinNonce(address)", addr)
+        );
+        require(success, "Failed to get nonce");
+        return abi.decode(data, (uint256));
+    }
+
+    /// Retrieve tx nonce for `addr` from the NONCE_HOLDER system contract
+    function getDeployNonce(address addr) internal returns (uint256) {
+        (bool success, bytes memory data) = NONCE_HOLDER.call(
+            abi.encodeWithSignature("getDeploymentNonce(address)", addr)
+        );
         require(success, "Failed to get nonce");
         return abi.decode(data, (uint256));
     }

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

@@ -20,4 +20,5 @@ mod paymaster;
 mod proxy;
 mod repros;
 mod script;
+mod sender;
 mod traces;

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

@@ -13,7 +13,7 @@ forgetest_async!(setup_block_on_script_test, |prj, cmd| {
         "./script/ScriptSetup.s.sol",
         "ScriptSetupNonce",
         None,
-        4,
+        3,
         Some(&["-vvvvv"]),
     )
     .await;

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

@@ -0,0 +1,25 @@
+//! Forge tests for basic zkysnc functionality.
+
+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_correctly_updates_sender_nonce() {
+    let runner = TEST_DATA_DEFAULT.runner_zksync();
+    let filter = Filter::new(
+        "testZkTxSenderNoncesAreConsistent|testZkTxSenderNoncesAreConsistentInBroadcast",
+        "ZkTxSenderTest",
+        ".*",
+    );
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::SHANGHAI).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_zk_correctly_updates_sender_balance() {
+    let runner = TEST_DATA_DEFAULT.runner_zksync();
+    let filter = Filter::new("testZkTxSenderBalancesAreConsistent", "ZkTxSenderTest", ".*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::SHANGHAI).run().await;
+}

crates/strategy/zksync/src/cheatcode/context.rs

@@ -10,7 +10,7 @@ use foundry_zksync_compilers::dual_compiled_contracts::{
 use foundry_zksync_core::{vm::ZkEnv, ZkPaymasterData, H256};
 use revm::primitives::Bytecode;
 
-use super::types::{ZkPersistNonceUpdate, ZkStartupMigration};
+use super::types::ZkStartupMigration;
 
 /// Context for [ZksyncCheatcodeInspectorStrategyRunner].
 #[derive(Debug, Default, Clone)]
@@ -48,9 +48,6 @@ pub struct ZksyncCheatcodeInspectorStrategyContext {
     /// providing the necessary level of isolation.
     pub persisted_factory_deps: HashMap<H256, Vec<u8>>,
 
-    /// Nonce update persistence behavior in zkEVM for the tx caller.
-    pub zk_persist_nonce_update: ZkPersistNonceUpdate,
-
     /// Stores the factory deps that were detected as part of CREATE2 deployer call.
     /// Must be cleared every call.
     pub set_deployer_call_input_factory_deps: Vec<Vec<u8>>,
@@ -108,7 +105,6 @@ impl ZksyncCheatcodeInspectorStrategyContext {
             zk_startup_migration: ZkStartupMigration::Defer,
             zk_use_factory_deps: Default::default(),
             persisted_factory_deps: Default::default(),
-            zk_persist_nonce_update: Default::default(),
             set_deployer_call_input_factory_deps: Default::default(),
             zk_env,
         }

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

@@ -8,8 +8,9 @@ use foundry_cheatcodes::{
         createFork_0Call, createFork_1Call, createFork_2Call, createSelectFork_0Call,
         createSelectFork_1Call, createSelectFork_2Call, dealCall, etchCall, getCodeCall,
         getNonce_0Call, mockCallRevert_0Call, mockCall_0Call, resetNonceCall, rollCall,
-        selectForkCall, setNonceCall, setNonceUnsafeCall, warpCall, zkRegisterContractCall,
-        zkUseFactoryDepCall, zkUsePaymasterCall, zkVmCall, zkVmSkipCall,
+        selectForkCall, setNonceCall, setNonceUnsafeCall, warpCall, zkGetDeploymentNonceCall,
+        zkGetTransactionNonceCall, zkRegisterContractCall, zkUseFactoryDepCall, zkUsePaymasterCall,
+        zkVmCall, zkVmSkipCall,
     },
 };
 use foundry_evm::backend::LocalForkId;
@@ -100,6 +101,26 @@ impl ZksyncCheatcodeInspectorStrategyRunner {
                 let nonce = foundry_zksync_core::cheatcodes::get_nonce(account, ccx.ecx);
                 Ok(nonce.abi_encode())
             }
+            t if using_zk_vm && is::<zkGetTransactionNonceCall>(t) => {
+                let &zkGetTransactionNonceCall { account } =
+                    cheatcode.as_any().downcast_ref().unwrap();
+
+                info!(?account, "cheatcode zkGetTransactionNonce");
+
+                let (tx_nonce, _) =
+                    foundry_zksync_core::cheatcodes::get_full_nonce(account, ccx.ecx);
+                Ok(tx_nonce.abi_encode())
+            }
+            t if using_zk_vm && is::<zkGetDeploymentNonceCall>(t) => {
+                let &zkGetDeploymentNonceCall { account } =
+                    cheatcode.as_any().downcast_ref().unwrap();
+
+                info!(?account, "cheatcode zkGetDeploymentNonce");
+
+                let (_, deploy_nonce) =
+                    foundry_zksync_core::cheatcodes::get_full_nonce(account, ccx.ecx);
+                Ok(deploy_nonce.abi_encode())
+            }
             t if using_zk_vm && is::<mockCall_0Call>(t) => {
                 let mockCall_0Call { callee, data, returnData } =
                     cheatcode.as_any().downcast_ref().unwrap();