fix: fetch all factory deps for a given contract (#316)

Co-authored-by: Jrigada <jrigada@est.frba.utn.edu.ar>
Co-authored-by: Herman Obst Demaestri <70286869+HermanObst@users.noreply.github.com>
Co-authored-by: Nisheeth Barthwal <nbaztec@gmail.com>

crates/cheatcodes/src/config.rs

@@ -8,7 +8,7 @@ use foundry_config::{
     ResolvedRpcEndpoints,
 };
 use foundry_evm_core::opts::EvmOpts;
-use foundry_zksync_compiler::DualCompiledContract;
+use foundry_zksync_compiler::DualCompiledContracts;
 use std::{
     collections::HashMap,
     path::{Path, PathBuf},
@@ -42,7 +42,7 @@ pub struct CheatsConfig {
     /// Script wallets
     pub script_wallets: Option<ScriptWallets>,
     /// ZKSolc -> Solc Contract codes
-    pub dual_compiled_contracts: Vec<DualCompiledContract>,
+    pub dual_compiled_contracts: DualCompiledContracts,
     /// Use ZK-VM on startup
     pub use_zk: bool,
 }
@@ -53,7 +53,7 @@ impl CheatsConfig {
         config: &Config,
         evm_opts: EvmOpts,
         script_wallets: Option<ScriptWallets>,
-        dual_compiled_contracts: Vec<DualCompiledContract>,
+        dual_compiled_contracts: DualCompiledContracts,
         use_zk: bool,
     ) -> Self {
         let mut allowed_paths = vec![config.__root.0.clone()];

crates/cheatcodes/src/inspector.rs

@@ -28,10 +28,10 @@ use foundry_evm_core::{
         HARDHAT_CONSOLE_ADDRESS,
     },
 };
-use foundry_zksync_compiler::{DualCompiledContract, FindContract};
+use foundry_zksync_compiler::DualCompiledContracts;
 use foundry_zksync_core::{
     convert::{ConvertAddress, ConvertH160, ConvertH256, ConvertRU256, ConvertU256},
-    ZkTransactionMetadata,
+    hash_bytecode, ZkTransactionMetadata,
 };
 use itertools::Itertools;
 use revm::{
@@ -58,7 +58,7 @@ use zksync_types::{
     get_code_key, get_nonce_key,
     utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
     ACCOUNT_CODE_STORAGE_ADDRESS, CONTRACT_DEPLOYER_ADDRESS, CURRENT_VIRTUAL_BLOCK_INFO_POSITION,
-    KNOWN_CODES_STORAGE_ADDRESS, L2_ETH_TOKEN_ADDRESS, NONCE_HOLDER_ADDRESS,
+    H256, KNOWN_CODES_STORAGE_ADDRESS, L2_ETH_TOKEN_ADDRESS, NONCE_HOLDER_ADDRESS,
     SYSTEM_CONTEXT_ADDRESS,
 };
 
@@ -231,7 +231,7 @@ pub struct Cheatcodes {
     pub use_zk_vm: bool,
 
     /// Dual compiled contracts
-    pub dual_compiled_contracts: Vec<DualCompiledContract>,
+    pub dual_compiled_contracts: DualCompiledContracts,
 
     /// Logs printed during ZK-VM execution.
     /// EVM logs have the value `None` so they can be interpolated later, since
@@ -240,6 +240,13 @@ pub struct Cheatcodes {
 
     /// Starts the cheatcode inspector in ZK mode
     pub startup_zk: bool,
+
+    /// The list of factory_deps seen so far during a test or script execution.
+    /// Ideally these would be persisted in the storage, but since modifying [revm::JournaledState]
+    /// would be a significant refactor, we maintain the factory_dep part in the [Cheatcodes].
+    /// This can be done as each test runs with its own [Cheatcodes] instance, thereby
+    /// providing the necessary level of isolation.
+    pub persisted_factory_deps: HashMap<H256, Vec<u8>>,
 }
 
 impl Cheatcodes {
@@ -404,10 +411,8 @@ impl Cheatcodes {
                 .map(|(value, _)| value)
                 .ok()
                 .and_then(|zk_bytecode_hash| {
-                    let zk_bytecode_hash = zk_bytecode_hash.to_h256();
                     self.dual_compiled_contracts
-                        .iter()
-                        .find(|c| c.zk_bytecode_hash == zk_bytecode_hash)
+                        .find_by_zk_bytecode_hash(zk_bytecode_hash.to_h256())
                         .map(|contract| {
                             (
                                 contract.evm_bytecode_hash,
@@ -1189,31 +1194,15 @@ impl<DB: DatabaseExt + Send> Inspector<DB> for Cheatcodes {
             }
 
             info!("running call in zk vm {:#?}", call);
-
-            let code_hash = data
-                .journaled_state
-                .load_account(call.contract, data.db)
-                .map(|(account, _)| account.info.code_hash)
-                .unwrap_or_default();
-            let contract = if code_hash != KECCAK_EMPTY {
-                self.dual_compiled_contracts
-                    .find_zk_bytecode_hash(zksync_types::H256::from(code_hash.0))
-            } else {
-                None
-            };
-            if let Some(contract) = contract {
-                tracing::debug!(contract = contract.name, "using dual compiled contract");
-            } else {
-                error!("no zk contract was found for {code_hash:?}");
-            }
+            let persisted_factory_deps = self.persisted_factory_deps.clone();
 
             let ccx = foundry_zksync_core::vm::CheatcodeTracerContext {
                 mocked_calls: self.mocked_calls.clone(),
                 expected_calls: Some(&mut self.expected_calls),
+                persisted_factory_deps,
             };
             if let Ok(result) = foundry_zksync_core::vm::call::<_, DatabaseError>(
                 call,
-                contract,
                 data.env,
                 data.db,
                 &mut data.journaled_state,
@@ -1569,10 +1558,13 @@ impl<DB: DatabaseExt + Send> Inspector<DB> for Cheatcodes {
                         ) as u64;
                         let contract = self
                             .dual_compiled_contracts
-                            .find_evm_bytecode(&call.init_code.0)
+                            .find_by_evm_bytecode(&call.init_code.0)
                             .unwrap_or_else(|| {
                                 panic!("failed finding contract for {:?}", call.init_code)
                             });
+                        let factory_deps =
+                            self.dual_compiled_contracts.fetch_all_factory_deps(contract);
+
                         let constructor_input =
                             call.init_code[contract.evm_bytecode.len()..].to_vec();
                         let create_input = foundry_zksync_core::encode_create_params(
@@ -1581,7 +1573,6 @@ impl<DB: DatabaseExt + Send> Inspector<DB> for Cheatcodes {
                             constructor_input,
                         );
                         bytecode = Bytes::from(create_input);
-                        let factory_deps = vec![contract.zk_deployed_bytecode.clone()];
 
                         Some(ZkTransactionMetadata { factory_deps })
                     } else {
@@ -1674,17 +1665,27 @@ impl<DB: DatabaseExt + Send> Inspector<DB> for Cheatcodes {
 
             let zk_contract = self
                 .dual_compiled_contracts
-                .find_evm_bytecode(&call.init_code.0)
+                .find_by_evm_bytecode(&call.init_code.0)
                 .unwrap_or_else(|| panic!("failed finding contract for {:?}", call.init_code));
 
+            let factory_deps = self.dual_compiled_contracts.fetch_all_factory_deps(zk_contract);
+
+            // get the current persisted factory deps to pass to zk create
+            let persisted_factory_deps = self.persisted_factory_deps.clone();
+            // and extend it for future calls
+            self.persisted_factory_deps
+                .extend(factory_deps.iter().map(|dep| (hash_bytecode(dep), dep.clone())));
+
             tracing::debug!(contract = zk_contract.name, "using dual compiled contract");
             let ccx = foundry_zksync_core::vm::CheatcodeTracerContext {
                 mocked_calls: self.mocked_calls.clone(),
                 expected_calls: Some(&mut self.expected_calls),
+                persisted_factory_deps,
             };
             if let Ok(result) = foundry_zksync_core::vm::create::<_, DatabaseError>(
                 call,
                 zk_contract,
+                factory_deps,
                 data.env,
                 data.db,
                 &mut data.journaled_state,

crates/cheatcodes/src/test.rs

@@ -38,6 +38,8 @@ impl Cheatcode for zkRegisterContractCall {
             name: name.clone(),
             zk_bytecode_hash: zkBytecodeHash.0.into(),
             zk_deployed_bytecode: zkDeployedBytecode.clone(),
+            //TODO: add argument to cheatcode
+            zk_factory_deps: vec![],
             evm_bytecode_hash: *evmBytecodeHash,
             evm_deployed_bytecode: evmDeployedBytecode.clone(),
             evm_bytecode: evmBytecode.clone(),

crates/forge/bin/cmd/create.rs

@@ -30,9 +30,7 @@ use foundry_compilers::{
     utils::canonicalized,
 };
 use foundry_wallets::WalletSigner;
-use foundry_zksync_compiler::{
-    new_dual_compiled_contracts, DualCompiledContract, FindContract, ZkSolc,
-};
+use foundry_zksync_compiler::{DualCompiledContract, DualCompiledContracts, ZkSolc};
 use serde_json::json;
 use std::{borrow::Borrow, marker::PhantomData, path::PathBuf, sync::Arc};
 
@@ -123,7 +121,7 @@ impl CreateArgs {
             Ok(compiled) => compiled,
             Err(e) => return Err(eyre::eyre!("Failed to compile with zksolc: {}", e)),
         };
-        let dual_compiled_contracts = new_dual_compiled_contracts(&output, &zk_output);
+        let dual_compiled_contracts = DualCompiledContracts::new(&output, &zk_output);
 
         if let Some(ref mut path) = self.contract.path {
             // paths are absolute in the project's output
@@ -136,11 +134,12 @@ impl CreateArgs {
             let contract = bin
                 .object
                 .as_bytes()
-                .and_then(|bytes| dual_compiled_contracts.find_evm_bytecode(&bytes.0))
+                .and_then(|bytes| dual_compiled_contracts.find_by_evm_bytecode(&bytes.0))
                 .ok_or(eyre::eyre!(
                     "Could not find zksolc contract for contract {}",
                     self.contract.name
                 ))?;
+
             let zk_bin = CompactBytecode {
                 object: BytecodeObject::Bytecode(Bytes::from(
                     contract.zk_deployed_bytecode.clone(),
@@ -149,8 +148,9 @@ impl CreateArgs {
                 source_map: Default::default(),
             };
 
-            let mut factory_deps = Vec::with_capacity(self.factory_deps.len());
+            let mut factory_deps = dual_compiled_contracts.fetch_all_factory_deps(contract);
 
+            // for manual specified factory deps
             for mut contract in std::mem::take(&mut self.factory_deps) {
                 if let Some(path) = contract.path.as_mut() {
                     *path = canonicalized(project.root().join(&path)).to_string_lossy().to_string();
@@ -163,16 +163,27 @@ impl CreateArgs {
                 let zk = bin
                     .object
                     .as_bytes()
-                    .and_then(|bytes| dual_compiled_contracts.find_evm_bytecode(&bytes.0))
+                    .and_then(|bytes| dual_compiled_contracts.find_by_evm_bytecode(&bytes.0))
                     .ok_or(eyre::eyre!(
                         "Could not find zksolc contract for contract {}",
                         contract.name
                     ))?;
 
-                factory_deps.push(zk.zk_deployed_bytecode.clone());
+                // if the dep isn't already present,
+                // fetch all deps and add them to the final list
+                if !factory_deps.contains(&zk.zk_deployed_bytecode) {
+                    let additional_factory_deps =
+                        dual_compiled_contracts.fetch_all_factory_deps(zk);
+                    factory_deps.extend(additional_factory_deps);
+                    factory_deps.dedup();
+                }
             }
 
-            (abi, zk_bin, Some((contract, factory_deps)))
+            (
+                abi,
+                zk_bin,
+                Some((contract, factory_deps.into_iter().map(|bc| bc.to_vec()).collect())),
+            )
         } else {
             (abi, bin, None)
         };
@@ -297,16 +308,7 @@ impl CreateArgs {
         let factory = ContractFactory::new(abi.clone(), bin.clone(), provider.clone());
 
         let is_args_empty = args.is_empty();
-        let (zk_contract, factory_deps) = match zk_data {
-            Some((zk_contract, mut factory_deps)) => {
-                //add this contract to the list of factory deps
-                factory_deps.push(zk_contract.zk_deployed_bytecode.clone());
-                (Some(zk_contract), factory_deps)
-            }
-            None => (None, vec![]),
-        };
-
-        let deployer = if let Some(contract) = zk_contract {
+        let deployer = if let Some((contract, factory_deps)) = &zk_data {
             factory.deploy_tokens_zk(args.clone(), contract).context("failed to deploy contract")
                 .map(|deployer| deployer.set_zk_factory_deps(factory_deps.clone()))
         } else {
@@ -328,9 +330,9 @@ impl CreateArgs {
             deployer.tx.set_value(value.to_ethers());
         }
 
-        match zk_contract {
+        match zk_data {
             None => provider.fill_transaction(&mut deployer.tx, None).await?,
-            Some(contract) => {
+            Some((contract, factory_deps)) => {
                 let chain_id = provider.get_chainid().await?.as_u64();
                 deployer.tx.set_chain_id(chain_id);
 
@@ -358,9 +360,12 @@ impl CreateArgs {
                     .tx
                     .set_to(NameOrAddress::from(foundry_zksync_core::CONTRACT_DEPLOYER_ADDRESS));
 
-                let estimated_gas =
-                    foundry_zksync_core::estimate_gas(&deployer.tx, factory_deps, &provider)
-                        .await?;
+                let estimated_gas = foundry_zksync_core::estimate_gas(
+                    &deployer.tx,
+                    factory_deps.clone(),
+                    &provider,
+                )
+                .await?;
                 deployer.tx.set_gas(estimated_gas.limit.to_ethers());
                 deployer.tx.set_gas_price(estimated_gas.price.to_ethers());
             }

crates/forge/bin/cmd/script/broadcast.rs

@@ -29,7 +29,7 @@ use foundry_common::{
 use foundry_compilers::{artifacts::Libraries, ArtifactId};
 use foundry_config::Config;
 use foundry_wallets::WalletSigner;
-use foundry_zksync_compiler::DualCompiledContract;
+use foundry_zksync_compiler::DualCompiledContracts;
 use futures::StreamExt;
 use std::{
     cmp::min,
@@ -317,7 +317,7 @@ impl ScriptArgs {
         mut script_config: ScriptConfig,
         verify: VerifyBundle,
         signers: &HashMap<Address, WalletSigner>,
-        dual_compiled_contracts: Option<Vec<DualCompiledContract>>,
+        dual_compiled_contracts: Option<DualCompiledContracts>,
     ) -> Result<()> {
         if let Some(txs) = result.transactions.take() {
             script_config.collect_rpcs(&txs);
@@ -419,7 +419,7 @@ impl ScriptArgs {
         script_config: &mut ScriptConfig,
         decoder: &CallTraceDecoder,
         known_contracts: &ContractsByArtifact,
-        dual_compiled_contracts: Option<Vec<DualCompiledContract>>,
+        dual_compiled_contracts: Option<DualCompiledContracts>,
     ) -> Result<Vec<ScriptSequence>> {
         if !txs.is_empty() {
             let gas_filled_txs = self
@@ -458,7 +458,7 @@ impl ScriptArgs {
         script_config: &ScriptConfig,
         decoder: &CallTraceDecoder,
         known_contracts: &ContractsByArtifact,
-        dual_compiled_contracts: Option<Vec<DualCompiledContract>>,
+        dual_compiled_contracts: Option<DualCompiledContracts>,
     ) -> Result<VecDeque<TransactionWithMetadata>> {
         let gas_filled_txs = if self.skip_simulation {
             shell::println("\nSKIPPING ON CHAIN SIMULATION.")?;

crates/forge/bin/cmd/script/build.rs

@@ -11,7 +11,7 @@ use foundry_compilers::{
     info::ContractInfo,
     ArtifactId, Project, ProjectCompileOutput,
 };
-use foundry_zksync_compiler::{new_dual_compiled_contracts, DualCompiledContract, ZkSolc};
+use foundry_zksync_compiler::{DualCompiledContracts, ZkSolc};
 use std::str::FromStr;
 
 impl ScriptArgs {
@@ -46,7 +46,7 @@ impl ScriptArgs {
             Ok(compiled) => compiled,
             Err(e) => return Err(eyre::eyre!("Failed to compile with zksolc: {}", e)),
         };
-        let dual_compiled_contracts = new_dual_compiled_contracts(&output, &zk_output);
+        let dual_compiled_contracts = DualCompiledContracts::new(&output, &zk_output);
 
         let sources = ContractSources::from_project_output(&output, root)?;
         let contracts = output.into_artifacts().collect();
@@ -228,5 +228,5 @@ pub struct BuildOutput {
     pub libraries: Libraries,
     pub predeploy_libraries: Vec<Bytes>,
     pub sources: ContractSources,
-    pub dual_compiled_contracts: Option<Vec<DualCompiledContract>>,
+    pub dual_compiled_contracts: Option<DualCompiledContracts>,
 }

crates/forge/bin/cmd/script/cmd.rs

@@ -19,7 +19,7 @@ use foundry_compilers::{
 use foundry_debugger::Debugger;
 use foundry_evm::inspectors::cheatcodes::{BroadcastableTransaction, ScriptWallets};
 use foundry_wallets::WalletSigner;
-use foundry_zksync_compiler::DualCompiledContract;
+use foundry_zksync_compiler::DualCompiledContracts;
 use std::{collections::HashMap, sync::Arc};
 
 /// Helper alias type for the collection of data changed due to the new sender.
@@ -159,7 +159,7 @@ impl ScriptArgs {
         predeploy_libraries: Vec<Bytes>,
         result: &mut ScriptResult,
         script_wallets: ScriptWallets,
-        dual_compiled_contracts: Option<Vec<DualCompiledContract>>,
+        dual_compiled_contracts: Option<DualCompiledContracts>,
     ) -> Result<Option<NewSenderChanges>> {
         if let Some(new_sender) = self.maybe_new_sender(
             &script_config.evm_opts,
@@ -332,7 +332,7 @@ impl ScriptArgs {
         first_run_result: &mut ScriptResult,
         linker: Linker,
         script_wallets: ScriptWallets,
-        dual_compiled_contracts: Option<Vec<DualCompiledContract>>,
+        dual_compiled_contracts: Option<DualCompiledContracts>,
     ) -> Result<(Libraries, ArtifactContracts<ContractBytecodeSome>)> {
         // if we had a new sender that requires relinking, we need to
         // get the nonce mainnet for accurate addresses for predeploy libs