chore: explicit behavior for nonce update persistence (#729)

* explicit behavior for nonce update persistence

* prevent short-circuit bug

crates/cheatcodes/src/inspector.rs

@@ -410,6 +410,39 @@ impl ArbitraryStorage {
 /// List of transactions that can be broadcasted.
 pub type BroadcastableTransactions = VecDeque<BroadcastableTransaction>;
 
+/// Allows overriding nonce update behavior for the tx caller in the zkEVM.
+///
+/// Since each CREATE or CALL is executed as a separate transaction within zkEVM, we currently skip
+/// persisting nonce updates as it erroneously increments the tx nonce. However, under certain
+/// situations, e.g. deploying contracts, transacts, etc. the nonce updates must be persisted.
+#[derive(Default, Debug, Clone)]
+pub enum ZkPersistNonceUpdate {
+    /// Never update the nonce. This is currently the default behavior.
+    #[default]
+    Never,
+    /// Override the default behavior, and persist nonce update for tx caller for the next
+    /// zkEVM execution _only_.
+    PersistNext,
+}
+
+impl ZkPersistNonceUpdate {
+    /// Persist nonce update for the tx caller for next execution.
+    pub fn persist_next(&mut self) {
+        *self = Self::PersistNext;
+    }
+
+    /// Retrieve if a nonce update must be persisted, or not. Resets the state to default.
+    pub fn get(&mut self) -> bool {
+        let persist_nonce_update = match self {
+            ZkPersistNonceUpdate::Never => false,
+            ZkPersistNonceUpdate::PersistNext => true,
+        };
+        *self = Default::default();
+
+        persist_nonce_update
+    }
+}
+
 /// Setting for migrating the database to zkEVM storage when starting in ZKsync mode.
 /// The migration is performed on the DB via the inspector so must only be performed once.
 #[derive(Debug, Default, Clone)]
@@ -609,9 +642,8 @@ pub struct Cheatcodes {
     /// providing the necessary level of isolation.
     pub persisted_factory_deps: HashMap<H256, Vec<u8>>,
 
-    /// Whether to persist nonce updates in ZK-VM. This is an option to act as a sticky flag
-    /// for the nonce updates.
-    pub zk_should_update_nonce: Option<bool>,
+    /// Nonce update persistence behavior in zkEVM for the tx caller.
+    pub zk_persist_nonce_update: ZkPersistNonceUpdate,
 }
 
 // This is not derived because calling this in `fn new` with `..Default::default()` creates a second
@@ -709,7 +741,7 @@ impl Cheatcodes {
             persisted_factory_deps: Default::default(),
             paymaster_params: None,
             zk_use_factory_deps: Default::default(),
-            zk_should_update_nonce: None,
+            zk_persist_nonce_update: Default::default(),
         }
     }
 
@@ -1236,14 +1268,14 @@ impl Cheatcodes {
         let factory_deps = self.dual_compiled_contracts.fetch_all_factory_deps(contract);
         tracing::debug!(contract = contract.name, "using dual compiled contract");
 
+        let zk_persist_nonce_update = self.zk_persist_nonce_update.get();
         let ccx = foundry_zksync_core::vm::CheatcodeTracerContext {
             mocked_calls: self.mocked_calls.clone(),
             expected_calls: Some(&mut self.expected_calls),
             accesses: self.accesses.as_mut(),
             persisted_factory_deps: Some(&mut self.persisted_factory_deps),
             paymaster_data: self.paymaster_params.take(),
-            should_update_nonce: self.broadcast.is_some() ||
-                self.zk_should_update_nonce.take().unwrap_or_default(),
+            persist_nonce_update: self.broadcast.is_some() || zk_persist_nonce_update,
         };
 
         let zk_create = foundry_zksync_core::vm::ZkCreateInputs {
@@ -1873,8 +1905,7 @@ where {
             accesses: self.accesses.as_mut(),
             persisted_factory_deps: Some(&mut self.persisted_factory_deps),
             paymaster_data: self.paymaster_params.take(),
-            should_update_nonce: self.broadcast.is_some() ||
-                self.zk_should_update_nonce.take().unwrap_or_default(),
+            persist_nonce_update: self.broadcast.is_some() || self.zk_persist_nonce_update.get(),
         };
 
         let mut gas = Gas::new(call.gas_limit);

crates/evm/evm/src/executors/mod.rs

@@ -429,7 +429,7 @@ impl Executor {
     pub fn call_with_env(&self, mut env: EnvWithHandlerCfg) -> eyre::Result<RawCallResult> {
         let mut inspector = self.inspector().clone();
         if let Some(cheatcodes) = inspector.cheatcodes.as_mut() {
-            cheatcodes.zk_should_update_nonce = Some(true);
+            cheatcodes.zk_persist_nonce_update.persist_next();
         }
         let mut backend = CowBackend::new_borrowed(self.backend());
         let result = match &self.zk_tx {
@@ -454,7 +454,7 @@ impl Executor {
     pub fn transact_with_env(&mut self, mut env: EnvWithHandlerCfg) -> eyre::Result<RawCallResult> {
         let mut inspector = self.inspector.clone();
         if let Some(cheatcodes) = inspector.cheatcodes.as_mut() {
-            cheatcodes.zk_should_update_nonce = Some(true);
+            cheatcodes.zk_persist_nonce_update.persist_next();
         }
         let backend = &mut self.backend;
         let result_and_state = match self.zk_tx.take() {

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

@@ -191,7 +191,7 @@ where
     info!(?call_ctx, "executing transaction in zk vm");
 
     let initiator_address = tx.common_data.initiator_address;
-    let should_update_nonce = ccx.should_update_nonce;
+    let persist_nonce_update = ccx.persist_nonce_update;
 
     if tx.common_data.signature.is_empty() {
         // FIXME: This is a hack to make sure that the signature is not empty.
@@ -336,7 +336,7 @@ where
     let filtered = modified_storage.iter().filter(|(k, _)| {
         !(k.address() == &NONCE_HOLDER_ADDRESS &&
             get_nonce_key(&initiator_address) == **k &&
-            !should_update_nonce)
+            !persist_nonce_update)
     });
 
     for (k, v) in filtered {

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

@@ -78,7 +78,7 @@ where
 
     let mut ccx = CheatcodeTracerContext {
         persisted_factory_deps,
-        should_update_nonce: true,
+        persist_nonce_update: true,
         ..Default::default()
     };
 

crates/zksync/core/src/vm/tracers/cheatcode.rs

@@ -90,8 +90,8 @@ pub struct CheatcodeTracerContext<'a> {
     pub persisted_factory_deps: Option<&'a mut HashMap<H256, Vec<u8>>>,
     /// Paymaster data
     pub paymaster_data: Option<ZkPaymasterData>,
-    /// Whether to persist the nonce updates or not
-    pub should_update_nonce: bool,
+    /// Whether to persist nonce update for the tx caller, or not.
+    pub persist_nonce_update: bool,
 }
 
 /// Tracer result to return back to foundry.