fix zk_env

matter-labs · Dec 16, 2024 · f0d0f9d · f0d0f9d
1 parent 589c8ba
commit f0d0f9d
Show file tree

Hide file tree

Showing 9 changed files with 58 additions and 31 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/cheatcodes/src/config.rs b/crates/cheatcodes/src/config.rs
@@ -63,7 +63,6 @@ pub struct CheatsConfig {
     pub assertions_revert: bool,
     /// Optional seed for the RNG algorithm.
     pub seed: Option<U256>,
-
 }
 
 impl CheatsConfig {
@@ -108,7 +107,6 @@ impl CheatsConfig {
             strategy,
             assertions_revert: config.assertions_revert,
             seed: config.fuzz.seed,
-            zk_env,
         }
     }
 
@@ -240,7 +238,6 @@ impl Default for CheatsConfig {
             strategy: Arc::new(Mutex::new(EvmCheatcodeInspectorStrategy::default())),
             assertions_revert: true,
             seed: None,
-            zk_env: Default::default(),
         }
     }
 }

diff --git a/crates/evm/evm/Cargo.toml b/crates/evm/evm/Cargo.toml
@@ -29,6 +29,7 @@ foundry-zksync-inspectors.workspace = true
 alloy-dyn-abi = { workspace = true, features = [ "arbitrary", "eip712" ] }
 alloy-json-abi.workspace = true
 alloy-serde.workspace = true
+alloy-zksync.workspace = true
 alloy-primitives = { workspace = true, features = [
     "serde",
     "getrandom",

diff --git a/crates/evm/evm/src/executors/builder.rs b/crates/evm/evm/src/executors/builder.rs
@@ -2,7 +2,6 @@ use std::sync::{Arc, Mutex};
 
 use crate::{executors::Executor, inspectors::InspectorStackBuilder};
 use foundry_evm_core::backend::Backend;
-use foundry_zksync_core::vm::ZkEnv;
 use revm::primitives::{Env, EnvWithHandlerCfg, SpecId};
 
 use super::strategy::ExecutorStrategyExt;
@@ -25,9 +24,6 @@ pub struct ExecutorBuilder {
     spec_id: SpecId,
 
     legacy_assertions: bool,
-
-    use_zk: bool,
-    zk_env: ZkEnv,
 }
 
 impl Default for ExecutorBuilder {
@@ -38,8 +34,6 @@ impl Default for ExecutorBuilder {
             gas_limit: None,
             spec_id: SpecId::LATEST,
             legacy_assertions: false,
-            use_zk: false,
-            zk_env: Default::default(),
         }
     }
 }
@@ -82,20 +76,6 @@ impl ExecutorBuilder {
         self
     }
 
-    /// Sets the EVM spec to use
-    #[inline]
-    pub fn use_zk_vm(mut self, enable: bool) -> Self {
-        self.use_zk = enable;
-        self
-    }
-
-    /// Sets zk_env
-    #[inline]
-    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(

diff --git a/crates/evm/evm/src/executors/strategy.rs b/crates/evm/evm/src/executors/strategy.rs
@@ -5,6 +5,7 @@ use std::{
 
 use alloy_primitives::{Address, U256};
 use alloy_serde::OtherFields;
+use alloy_zksync::types::BlockDetails;
 use eyre::Context;
 use foundry_cheatcodes::strategy::{CheatcodeInspectorStrategyExt, EvmCheatcodeInspectorStrategy};
 use foundry_evm_core::{
@@ -72,6 +73,8 @@ pub trait ExecutorStrategyExt: ExecutorStrategy {
         _dual_compiled_contracts: DualCompiledContracts,
     ) {
     }
+
+    fn zksync_set_env(&mut self, _block_details: BlockDetails) {}
 }
 
 #[derive(Debug, Default, Clone)]

diff --git a/crates/script/src/lib.rs b/crates/script/src/lib.rs
@@ -58,7 +58,6 @@ use foundry_evm::{
 };
 use foundry_wallets::MultiWalletOpts;
 use foundry_zksync_compiler::DualCompiledContracts;
-use foundry_zksync_core::vm::ZkEnv;
 use serde::Serialize;
 use std::path::PathBuf;
 
@@ -640,6 +639,20 @@ impl ScriptConfig {
                 .expect("failed acquiring strategy")
                 .zksync_set_dual_compiled_contracts(dual_compiled_contracts);
 
+            if let Some(fork_url) = &self.evm_opts.fork_url {
+                let provider =
+                    zksync_provider().with_recommended_fillers().on_http(fork_url.parse()?);
+                // TODO(zk): switch to getFeeParams call when it is implemented for anvil-zksync
+                let maybe_details =
+                    provider.get_block_details(env.block.number.try_into()?).await?;
+                if let Some(block_details) = maybe_details {
+                    strategy
+                        .lock()
+                        .expect("failed acquiring strategy")
+                        .zksync_set_env(block_details);
+                }
+            }
+
             builder = builder.inspectors(|stack| {
                 stack
                     .cheatcodes(

diff --git a/crates/strategy/zksync/Cargo.toml b/crates/strategy/zksync/Cargo.toml
@@ -15,6 +15,7 @@ workspace = true
 [dependencies]
 alloy-sol-types.workspace = true
 alloy-json-abi.workspace = true
+alloy-zksync.workspace = true
 foundry-common.workspace = true
 foundry-config.workspace = true
 foundry-evm.workspace = true

diff --git a/crates/strategy/zksync/src/cheatcode.rs b/crates/strategy/zksync/src/cheatcode.rs
@@ -32,10 +32,11 @@ use foundry_evm_core::{
 use foundry_zksync_compiler::{ContractType, DualCompiledContract, DualCompiledContracts};
 use foundry_zksync_core::{
     convert::{ConvertAddress, ConvertH160, ConvertH256, ConvertRU256, ConvertU256},
-    get_account_code_key, get_balance_key, get_nonce_key, PaymasterParams, ZkPaymasterData,
-    ZkTransactionMetadata, ACCOUNT_CODE_STORAGE_ADDRESS, CONTRACT_DEPLOYER_ADDRESS,
-    DEFAULT_CREATE2_DEPLOYER_ZKSYNC, H256, KNOWN_CODES_STORAGE_ADDRESS, L2_BASE_TOKEN_ADDRESS,
-    NONCE_HOLDER_ADDRESS,
+    get_account_code_key, get_balance_key, get_nonce_key,
+    vm::ZkEnv,
+    PaymasterParams, ZkPaymasterData, ZkTransactionMetadata, ACCOUNT_CODE_STORAGE_ADDRESS,
+    CONTRACT_DEPLOYER_ADDRESS, DEFAULT_CREATE2_DEPLOYER_ZKSYNC, H256, KNOWN_CODES_STORAGE_ADDRESS,
+    L2_BASE_TOKEN_ADDRESS, NONCE_HOLDER_ADDRESS,
 };
 use itertools::Itertools;
 use revm::{
@@ -128,11 +129,11 @@ pub struct ZksyncCheatcodeInspectorStrategy {
     pub set_deployer_call_input_factory_deps: Vec<Vec<u8>>,
 
     /// Era Vm environment
-    pub zk_env: Option<ZkEnv>,
+    pub zk_env: ZkEnv,
 }
 
 impl ZksyncCheatcodeInspectorStrategy {
-    pub fn new(dual_compiled_contracts: DualCompiledContracts) -> Self {
+    pub fn new(dual_compiled_contracts: DualCompiledContracts, zk_env: ZkEnv) -> Self {
         // We add the empty bytecode manually so it is correctly translated in zk mode.
         // This is used in many places in foundry, e.g. in cheatcode contract's account code.
         let empty_bytes = Bytes::from_static(&[0]);
@@ -183,6 +184,7 @@ impl ZksyncCheatcodeInspectorStrategy {
             persisted_factory_deps: Default::default(),
             zk_persist_nonce_update: Default::default(),
             set_deployer_call_input_factory_deps: Default::default(),
+            zk_env,
         }
     }
 }
@@ -877,6 +879,7 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy {
             persisted_factory_deps: Some(&mut self.persisted_factory_deps),
             paymaster_data: self.paymaster_params.take(),
             persist_nonce_update: state.broadcast.is_some() || zk_persist_nonce_update,
+            zk_env: self.zk_env.clone(),
         };
 
         let zk_create = foundry_zksync_core::vm::ZkCreateInputs {
@@ -1058,6 +1061,7 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy {
             persisted_factory_deps: Some(&mut self.persisted_factory_deps),
             paymaster_data: self.paymaster_params.take(),
             persist_nonce_update: state.broadcast.is_some() || zk_persist_nonce_update,
+            zk_env: self.zk_env.clone(),
         };
 
         let mut gas = Gas::new(call.gas_limit);

diff --git a/crates/strategy/zksync/src/executor.rs b/crates/strategy/zksync/src/executor.rs
@@ -2,6 +2,7 @@ use std::sync::{Arc, Mutex};
 
 use alloy_primitives::{Address, U256};
 use alloy_rpc_types::serde_helpers::OtherFields;
+use alloy_zksync::types::BlockDetails;
 use foundry_cheatcodes::strategy::CheatcodeInspectorStrategyExt;
 
 use foundry_evm::{
@@ -13,7 +14,7 @@ use foundry_evm::{
     InspectorExt,
 };
 use foundry_zksync_compiler::DualCompiledContracts;
-use foundry_zksync_core::ZkTransactionMetadata;
+use foundry_zksync_core::{vm::ZkEnv, ZkTransactionMetadata};
 use revm::{
     primitives::{EnvWithHandlerCfg, HashMap, ResultAndState},
     Database,
@@ -31,6 +32,7 @@ pub struct ZksyncExecutorStrategy {
     inspect_context: Option<ZkTransactionMetadata>,
     persisted_factory_deps: HashMap<H256, Vec<u8>>,
     dual_compiled_contracts: DualCompiledContracts,
+    zk_env: ZkEnv,
 }
 
 impl ExecutorStrategy for ZksyncExecutorStrategy {
@@ -87,6 +89,7 @@ impl ExecutorStrategy for ZksyncExecutorStrategy {
     fn new_cheatcode_inspector_strategy(&self) -> Arc<Mutex<dyn CheatcodeInspectorStrategyExt>> {
         Arc::new(Mutex::new(ZksyncCheatcodeInspectorStrategy::new(
             self.dual_compiled_contracts.clone(),
+            self.zk_env.clone(),
         )))
     }
 
@@ -103,6 +106,7 @@ impl ExecutorStrategy for ZksyncExecutorStrategy {
                 Some(zk_tx.factory_deps.clone()),
                 zk_tx.paymaster_data.clone(),
                 env,
+                &self.zk_env,
                 db,
             ),
         }
@@ -128,6 +132,7 @@ impl ExecutorStrategy for ZksyncExecutorStrategy {
                     Some(zk_tx.factory_deps),
                     zk_tx.paymaster_data,
                     env,
+                    &self.zk_env,
                     db,
                 )
             }
@@ -146,6 +151,27 @@ impl ExecutorStrategyExt for ZksyncExecutorStrategy {
     ) {
         self.dual_compiled_contracts = dual_compiled_contracts;
     }
+
+    fn zksync_set_env(&mut self, block_details: BlockDetails) {
+        self.zk_env = ZkEnv {
+            l1_gas_price: block_details
+                .l1_gas_price
+                .try_into()
+                .expect("failed to convert l1_gas_price to u64"),
+            fair_l2_gas_price: block_details
+                .l2_fair_gas_price
+                .try_into()
+                .expect("failed to convert fair_l2_gas_price to u64"),
+            fair_pubdata_price: block_details
+                .fair_pubdata_price
+                // TODO(zk): None as a value might mean L1Pegged model
+                // we need to find out if it will ever be relevant to
+                // us
+                .unwrap_or_default()
+                .try_into()
+                .expect("failed to convert fair_pubdata_price to u64"),
+        };
+    }
 }
 
 /// Retrieve metadata for zksync tx