diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index 142f9cb0e..70a63d0e4 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -410,6 +410,39 @@ impl ArbitraryStorage { /// List of transactions that can be broadcasted. pub type BroadcastableTransactions = VecDeque; +/// 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>, - /// 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, + /// 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); diff --git a/crates/evm/evm/src/executors/mod.rs b/crates/evm/evm/src/executors/mod.rs index 095bc2e67..476b4c2a5 100644 --- a/crates/evm/evm/src/executors/mod.rs +++ b/crates/evm/evm/src/executors/mod.rs @@ -429,7 +429,7 @@ impl Executor { pub fn call_with_env(&self, mut env: EnvWithHandlerCfg) -> eyre::Result { 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 { 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() { diff --git a/crates/zksync/core/src/vm/inspect.rs b/crates/zksync/core/src/vm/inspect.rs index 36e49db0f..8b0b9170f 100644 --- a/crates/zksync/core/src/vm/inspect.rs +++ b/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 { diff --git a/crates/zksync/core/src/vm/runner.rs b/crates/zksync/core/src/vm/runner.rs index 00a0bd2db..daf50047a 100644 --- a/crates/zksync/core/src/vm/runner.rs +++ b/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() }; diff --git a/crates/zksync/core/src/vm/tracers/cheatcode.rs b/crates/zksync/core/src/vm/tracers/cheatcode.rs index 9769ebc08..93d116c11 100644 --- a/crates/zksync/core/src/vm/tracers/cheatcode.rs +++ b/crates/zksync/core/src/vm/tracers/cheatcode.rs @@ -90,8 +90,8 @@ pub struct CheatcodeTracerContext<'a> { pub persisted_factory_deps: Option<&'a mut HashMap>>, /// Paymaster data pub paymaster_data: Option, - /// 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.