feat: add the ability to use specific gas params in era vm environment and use them on script estimations (#773)

---------

Co-authored-by: Juan Rigada <[email protected]>
Co-authored-by: Dustin Brickwood <[email protected]>

Author: 3 people

Cargo.lock

@@ -9,6 +9,7 @@ use foundry_config::{
 };
 use foundry_evm_core::opts::EvmOpts;
 use foundry_zksync_compiler::DualCompiledContracts;
+use foundry_zksync_core::vm::ZkEnv;
 use semver::Version;
 use std::{
     path::{Path, PathBuf},
@@ -62,10 +63,13 @@ pub struct CheatsConfig {
     pub assertions_revert: bool,
     /// Optional seed for the RNG algorithm.
     pub seed: Option<U256>,
+    /// Era Vm environment
+    pub zk_env: Option<ZkEnv>,
 }
 
 impl CheatsConfig {
     /// Extracts the necessary settings from the Config
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         config: &Config,
         evm_opts: EvmOpts,
@@ -74,6 +78,7 @@ impl CheatsConfig {
         running_version: Option<Version>,
         dual_compiled_contracts: DualCompiledContracts,
         use_zk: bool,
+        zk_env: Option<ZkEnv>,
     ) -> Self {
         let mut allowed_paths = vec![config.root.0.clone()];
         allowed_paths.extend(config.libs.clone());
@@ -107,6 +112,7 @@ impl CheatsConfig {
             use_zk,
             assertions_revert: config.assertions_revert,
             seed: config.fuzz.seed,
+            zk_env,
         }
     }
 
@@ -239,6 +245,7 @@ impl Default for CheatsConfig {
             use_zk: false,
             assertions_revert: true,
             seed: None,
+            zk_env: Default::default(),
         }
     }
 }
@@ -257,6 +264,7 @@ mod tests {
             None,
             Default::default(),
             false,
+            None,
         )
     }
 

crates/cheatcodes/src/config.rs

@@ -46,8 +46,9 @@ use foundry_wallets::multi_wallet::MultiWallet;
 use foundry_zksync_compiler::{DualCompiledContract, DualCompiledContracts};
 use foundry_zksync_core::{
     convert::{ConvertAddress, ConvertH160, ConvertH256, ConvertRU256, ConvertU256},
-    get_account_code_key, get_balance_key, get_nonce_key, Call, ZkPaymasterData,
-    ZkTransactionMetadata, DEFAULT_CREATE2_DEPLOYER_ZKSYNC,
+    get_account_code_key, get_balance_key, get_nonce_key,
+    vm::ZkEnv,
+    Call, ZkPaymasterData, ZkTransactionMetadata, DEFAULT_CREATE2_DEPLOYER_ZKSYNC,
 };
 use foundry_zksync_inspectors::TraceCollector;
 use itertools::Itertools;
@@ -650,6 +651,8 @@ pub struct Cheatcodes {
 
     /// Nonce update persistence behavior in zkEVM for the tx caller.
     pub zk_persist_nonce_update: ZkPersistNonceUpdate,
+
+    pub zk_env: ZkEnv,
 }
 
 // This is not derived because calling this in `fn new` with `..Default::default()` creates a second
@@ -703,6 +706,7 @@ impl Cheatcodes {
         persisted_factory_deps.insert(zk_bytecode_hash, zk_deployed_bytecode);
 
         let zk_startup_migration = config.use_zk.then_some(ZkStartupMigration::Defer);
+        let zk_env = config.zk_env.clone().unwrap_or_default();
 
         Self {
             fs_commit: true,
@@ -750,6 +754,7 @@ impl Cheatcodes {
             paymaster_params: None,
             zk_use_factory_deps: Default::default(),
             zk_persist_nonce_update: Default::default(),
+            zk_env,
         }
     }
 
@@ -1304,6 +1309,7 @@ impl Cheatcodes {
             persisted_factory_deps: Some(&mut self.persisted_factory_deps),
             paymaster_data: self.paymaster_params.take(),
             persist_nonce_update: self.broadcast.is_some() || zk_persist_nonce_update,
+            zk_env: self.zk_env.clone(),
         };
 
         let zk_create = foundry_zksync_core::vm::ZkCreateInputs {
@@ -1968,6 +1974,7 @@ where {
             persisted_factory_deps: Some(&mut self.persisted_factory_deps),
             paymaster_data: self.paymaster_params.take(),
             persist_nonce_update: self.broadcast.is_some() || zk_persist_nonce_update,
+            zk_env: self.zk_env.clone(),
         };
 
         let mut gas = Gas::new(call.gas_limit);

crates/cheatcodes/src/inspector.rs

@@ -338,6 +338,7 @@ impl SessionSource {
                         Some(self.solc.version.clone()),
                         Default::default(),
                         false,
+                        None,
                     )
                     .into(),
                 )

crates/chisel/src/executor.rs

@@ -13,7 +13,7 @@ use alloy_primitives::{map::HashMap, Address, B256, U256};
 use alloy_rpc_types::TransactionRequest;
 use eyre::WrapErr;
 use foundry_fork_db::DatabaseError;
-use foundry_zksync_core::PaymasterParams;
+use foundry_zksync_core::{vm::ZkEnv, PaymasterParams};
 use revm::{
     db::DatabaseRef,
     primitives::{
@@ -65,6 +65,7 @@ impl<'a> CowBackend<'a> {
     pub fn inspect_ref_zk(
         &mut self,
         env: &mut Env,
+        zk_env: &ZkEnv,
         persisted_factory_deps: &mut HashMap<foundry_zksync_core::H256, Vec<u8>>,
         factory_deps: Option<Vec<Vec<u8>>>,
         paymaster_data: Option<PaymasterParams>,
@@ -78,6 +79,7 @@ impl<'a> CowBackend<'a> {
             factory_deps,
             paymaster_data,
             env,
+            zk_env,
             self,
         )
     }

crates/evm/core/src/backend/cow.rs

@@ -16,7 +16,7 @@ use eyre::Context;
 use foundry_common::{is_known_system_sender, SYSTEM_TRANSACTION_TYPE};
 pub use foundry_fork_db::{cache::BlockchainDbMeta, BlockchainDb, SharedBackend};
 use foundry_zksync_core::{
-    convert::ConvertH160, PaymasterParams, ACCOUNT_CODE_STORAGE_ADDRESS,
+    convert::ConvertH160, vm::ZkEnv, PaymasterParams, ACCOUNT_CODE_STORAGE_ADDRESS,
     IMMUTABLE_SIMULATOR_STORAGE_ADDRESS, KNOWN_CODES_STORAGE_ADDRESS, L2_BASE_TOKEN_ADDRESS,
     NONCE_HOLDER_ADDRESS,
 };
@@ -840,6 +840,7 @@ impl Backend {
     pub fn inspect_ref_zk(
         &mut self,
         env: &mut EnvWithHandlerCfg,
+        zk_env: &ZkEnv,
         persisted_factory_deps: &mut HashMap<foundry_zksync_core::H256, Vec<u8>>,
         factory_deps: Option<Vec<Vec<u8>>>,
         paymaster_data: Option<PaymasterParams>,
@@ -851,6 +852,7 @@ impl Backend {
             factory_deps,
             paymaster_data,
             env,
+            zk_env,
             self,
         )
     }

crates/evm/core/src/backend/mod.rs

@@ -1,5 +1,6 @@
 use crate::{executors::Executor, inspectors::InspectorStackBuilder};
 use foundry_evm_core::backend::Backend;
+use foundry_zksync_core::vm::ZkEnv;
 use revm::primitives::{Env, EnvWithHandlerCfg, SpecId};
 
 /// The builder that allows to configure an evm [`Executor`] which a stack of optional
@@ -19,8 +20,10 @@ pub struct ExecutorBuilder {
     /// The spec ID.
     spec_id: SpecId,
 
-    use_zk: bool,
     legacy_assertions: bool,
+
+    use_zk: bool,
+    zk_env: ZkEnv,
 }
 
 impl Default for ExecutorBuilder {
@@ -30,8 +33,9 @@ impl Default for ExecutorBuilder {
             stack: InspectorStackBuilder::new(),
             gas_limit: None,
             spec_id: SpecId::LATEST,
-            use_zk: false,
             legacy_assertions: false,
+            use_zk: false,
+            zk_env: Default::default(),
         }
     }
 }
@@ -67,24 +71,31 @@ impl ExecutorBuilder {
         self
     }
 
+    /// Sets the `legacy_assertions` flag.
+    #[inline]
+    pub fn legacy_assertions(mut self, legacy_assertions: bool) -> Self {
+        self.legacy_assertions = legacy_assertions;
+        self
+    }
+
     /// Sets the EVM spec to use
     #[inline]
     pub fn use_zk_vm(mut self, enable: bool) -> Self {
         self.use_zk = enable;
         self
     }
 
-    /// Sets the `legacy_assertions` flag.
+    /// Sets zk_env
     #[inline]
-    pub fn legacy_assertions(mut self, legacy_assertions: bool) -> Self {
-        self.legacy_assertions = legacy_assertions;
+    pub fn zk_env(mut self, zk_env: ZkEnv) -> Self {
+        self.zk_env = zk_env;
         self
     }
 
     /// Builds the executor as configured.
     #[inline]
     pub fn build(self, env: Env, db: Backend) -> Executor {
-        let Self { mut stack, gas_limit, spec_id, legacy_assertions, use_zk } = self;
+        let Self { mut stack, gas_limit, spec_id, legacy_assertions, use_zk, zk_env } = self;
         if stack.block.is_none() {
             stack.block = Some(env.block.clone());
         }
@@ -95,6 +106,7 @@ impl ExecutorBuilder {
         let env = EnvWithHandlerCfg::new_with_spec_id(Box::new(env), spec_id);
         let mut exec = Executor::new(db, env, stack.build(), gas_limit, legacy_assertions);
         exec.use_zk = use_zk;
+        exec.zk_env = zk_env;
         exec
     }
 }

crates/evm/evm/src/executors/builder.rs

@@ -27,7 +27,7 @@ use foundry_evm_core::{
 };
 use foundry_evm_coverage::HitMaps;
 use foundry_evm_traces::{SparsedTraceArena, TraceMode};
-use foundry_zksync_core::ZkTransactionMetadata;
+use foundry_zksync_core::{vm::ZkEnv, ZkTransactionMetadata};
 use revm::{
     db::{DatabaseCommit, DatabaseRef},
     interpreter::{return_ok, InstructionResult},
@@ -97,6 +97,7 @@ pub struct Executor {
     zk_persisted_factory_deps: HashMap<foundry_zksync_core::H256, Vec<u8>>,
 
     pub use_zk: bool,
+    pub zk_env: ZkEnv,
 }
 
 impl Executor {
@@ -137,6 +138,7 @@ impl Executor {
             zk_tx: None,
             zk_persisted_factory_deps: Default::default(),
             use_zk: false,
+            zk_env: Default::default(),
         }
     }
 
@@ -453,6 +455,7 @@ impl Executor {
                 env.tx.gas_price = self.env.tx.gas_price;
                 backend.inspect_ref_zk(
                     &mut env,
+                    &self.zk_env,
                     &mut self.zk_persisted_factory_deps.clone(),
                     Some(zk_tx.factory_deps.clone()),
                     zk_tx.paymaster_data.clone(),
@@ -476,6 +479,7 @@ impl Executor {
                 env.tx.gas_price = self.env.tx.gas_price;
                 backend.inspect_ref_zk(
                     &mut env,
+                    &self.zk_env,
                     // this will persist the added factory deps,
                     // no need to commit them later
                     &mut self.zk_persisted_factory_deps,

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

@@ -256,6 +256,7 @@ impl MultiContractRunner {
             Some(artifact_id.version.clone()),
             self.dual_compiled_contracts.clone(),
             self.use_zk,
+            None,
         );
 
         let trace_mode = TraceMode::default()

crates/forge/src/multi_runner.rs

@@ -6,6 +6,7 @@ use foundry_test_utils::{
     util::{self, OutputExt},
     Filter, ZkSyncNode,
 };
+use foundry_zksync_core::utils::MAX_L2_GAS_LIMIT;
 
 use crate::{config::TestConfig, test_helpers::TEST_DATA_DEFAULT};
 
@@ -44,9 +45,9 @@ contract ZkLargeFactoryDependenciesScript is Script {
 
     let node = ZkSyncNode::start();
 
-    // foundry default gas-limit is not enough to pay for factory deps in our current
-    // default environment
-    let gas_limit = u32::MAX >> 1;
+    // foundry default gas-limit is not enough to pay for factory deps
+    // with Anvil-zksync's environment
+    let gas_limit = MAX_L2_GAS_LIMIT;
 
     cmd.arg("script").args([
         "--zk-startup",