feat: ZkUseFactoryDep cheatcode

crates/cheatcodes/spec/src/vm.rs

@@ -697,6 +697,10 @@ interface Vm {
     #[cheatcode(group = Testing, safety = Safe)]
     function zkUsePaymaster(address paymaster_address, bytes calldata paymaster_input) external pure;
 
+    /// Marks the contract to be injected as a factory dependency in the next transaction
+    #[cheatcode(group = Testing, safety = Safe)]
+    function zkUseFactoryDep(string calldata name) external pure;
+
     /// Registers bytecodes for ZK-VM for transact/call and create instructions.
     #[cheatcode(group = Testing, safety = Safe)]
     function zkRegisterContract(string calldata name, bytes32 evmBytecodeHash, bytes calldata evmDeployedBytecode, bytes calldata evmBytecode, bytes32 zkBytecodeHash, bytes calldata zkDeployedBytecode) external pure;

crates/cheatcodes/src/fs.rs

@@ -328,7 +328,7 @@ impl Cheatcode for deployCode_1Call {
 /// - `path/to/contract.sol:0.8.23`
 /// - `ContractName`
 /// - `ContractName:0.8.23`
-fn get_artifact_code(state: &Cheatcodes, path: &str, deployed: bool) -> Result<Bytes> {
+pub(crate) fn get_artifact_code(state: &Cheatcodes, path: &str, deployed: bool) -> Result<Bytes> {
     let path = if path.ends_with(".json") {
         PathBuf::from(path)
     } else {

crates/cheatcodes/src/inspector.rs

@@ -506,6 +506,10 @@ pub struct Cheatcodes {
     /// This is set to `false`, once the startup migration is completed.
     pub startup_zk: bool,
 
+    /// Factory deps stored through `zkUseFactoryDep`. These factory deps are used in the next
+    /// CREATE or CALL, and cleared after.
+    pub zk_use_factory_deps: Vec<String>,
+
     /// 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].
@@ -603,6 +607,7 @@ impl Cheatcodes {
             record_next_create_address: Default::default(),
             persisted_factory_deps: Default::default(),
             paymaster_params: None,
+            zk_use_factory_deps: Default::default(),
         }
     }
 
@@ -982,7 +987,16 @@ impl Cheatcodes {
                             paymaster: paymaster_data.address.to_h160(),
                             paymaster_input: paymaster_data.input.to_vec(),
                         });
-                    if let Some(factory_deps) = zk_tx {
+                    if let Some(mut factory_deps) = zk_tx {
+                        let injected_factory_deps = self.zk_use_factory_deps.iter().map(|contract| {
+                            crate::fs::get_artifact_code(self, contract, false)
+                                .inspect(|_| info!(contract, "pushing factory dep"))
+                                .unwrap_or_else(|_| {
+                                    panic!("failed to get bytecode for factory deps contract {contract}")
+                                })
+                                .to_vec()
+                        }).collect_vec();
+                        factory_deps.extend(injected_factory_deps);
                         let mut batched =
                             foundry_zksync_core::vm::batch_factory_dependencies(factory_deps);
                         debug!(batches = batched.len(), "splitting factory deps for broadcast");
@@ -1580,7 +1594,17 @@ impl Cheatcodes {
                     let account =
                         ecx_inner.journaled_state.state().get_mut(&broadcast.new_origin).unwrap();
 
-                    let zk_tx = if self.use_zk_vm {
+                    let mut zk_tx = if self.use_zk_vm {
+                        let injected_factory_deps = self.zk_use_factory_deps.iter().map(|contract| {
+                            crate::fs::get_artifact_code(self, contract, false)
+                                .inspect(|_| info!(contract, "pushing factory dep"))
+                                .unwrap_or_else(|_| {
+                                    panic!("failed to get bytecode for factory deps contract {contract}")
+                                })
+                                .to_vec()
+                        }).collect_vec();
+                        factory_deps.extend(injected_factory_deps);
+
                         let paymaster_params =
                             self.paymaster_params.clone().map(|paymaster_data| PaymasterParams {
                                 paymaster: paymaster_data.address.to_h160(),
@@ -1602,6 +1626,13 @@ impl Cheatcodes {
                         None
                     };
 
+                    if let Some(ref mut zk_tx) = zk_tx {
+                        zk_tx.factory_deps = factory_deps.clone();
+                        Some(zk_tx)
+                    } else {
+                        None
+                    };
+
                     self.broadcastable_transactions.push_back(BroadcastableTransaction {
                         rpc: ecx_inner.db.active_fork_url(),
                         transaction: TransactionRequest {

crates/cheatcodes/src/test.rs

@@ -45,6 +45,15 @@ impl Cheatcode for zkUsePaymasterCall {
     }
 }
 
+impl Cheatcode for zkUseFactoryDepCall {
+    fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { name } = self;
+        info!("Adding factory dependency: {:?}", name);
+        ccx.state.zk_use_factory_deps.push(name.clone());
+        Ok(Default::default())
+    }
+}
+
 impl Cheatcode for zkRegisterContractCall {
     fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self {

crates/forge/tests/fixtures/zk/Create2.s.sol

@@ -9,7 +9,7 @@ contract Create2Script is Script {
     function run() external {
         (bool success,) = address(vm).call(abi.encodeWithSignature("zkVm(bool)", true));
         require(success, "zkVm() call failed");
-        
+
         vm.startBroadcast();
 
         // Deploy Greeter using create2 with a salt

crates/forge/tests/fixtures/zk/DeployCounterWithBytecodeHash.s.sol

@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.17;
+
+import "forge-std/Script.sol";
+import "zksync-contracts/zksync-contracts/l2/system-contracts/libraries/SystemContractsCaller.sol";
+import "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
+import "../src/Factory.sol";
+
+contract DeployCounterWithBytecodeHash is Script {
+    function run() external {
+        // Read artifact file and get the bytecode hash
+        string memory artifact = vm.readFile("zkout/Counter.sol/Counter.json");
+        bytes32 counterBytecodeHash = vm.parseJsonBytes32(artifact, ".hash");
+        bytes32 salt = "JUAN";
+
+        vm.startBroadcast();
+        Factory factory = new Factory(counterBytecodeHash);
+        (bool _success,) = address(vm).call(abi.encodeWithSignature("zkUseFactoryDep(string)", "Counter"));
+        require(_success, "Cheatcode failed");
+        address counter = factory.deployAccount(salt);
+        require(counter != address(0), "Counter deployment failed");
+        vm.stopBroadcast();
+    }
+}

crates/forge/tests/fixtures/zk/Factory.sol

@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.17;
+
+import "zksync-contracts/zksync-contracts/l2/system-contracts/Constants.sol";
+import "zksync-contracts/zksync-contracts/l2/system-contracts/libraries/SystemContractsCaller.sol";
+
+contract Factory {
+    bytes32 public aaBytecodeHash;
+
+    constructor(bytes32 _aaBytecodeHash) {
+        aaBytecodeHash = _aaBytecodeHash;
+    }
+
+    function deployAccount(bytes32 salt) external returns (address accountAddress) {
+        (bool success, bytes memory returnData) = SystemContractsCaller.systemCallWithReturndata(
+            uint32(gasleft()),
+            address(DEPLOYER_SYSTEM_CONTRACT),
+            uint128(0),
+            abi.encodeCall(
+                DEPLOYER_SYSTEM_CONTRACT.create2Account,
+                (salt, aaBytecodeHash, abi.encode(), IContractDeployer.AccountAbstractionVersion.Version1)
+            )
+        );
+        require(success, "Deployment failed");
+
+        (accountAddress) = abi.decode(returnData, (address));
+    }
+}

crates/forge/tests/fixtures/zk/Paymaster.t.sol

@@ -74,7 +74,7 @@ contract TestPaymasterFlow is Test {
         vm.deal(address(do_stuff), 1 ether);
         require(address(alice).balance == 0, "Balance is not 0 ether");
         vm.prank(alice, alice);
-        
+
         do_stuff.do_stuff(bob);
     }
 }

crates/forge/tests/it/zk/cheats.rs

@@ -1,9 +1,14 @@
 //! Forge tests for cheatcodes.
 
-use crate::{config::*, test_helpers::TEST_DATA_DEFAULT};
+use std::path::Path;
+
+use crate::{
+    config::*,
+    test_helpers::{run_zk_script_test, TEST_DATA_DEFAULT},
+};
 use forge::revm::primitives::SpecId;
-use foundry_config::fs_permissions::PathPermission;
-use foundry_test_utils::Filter;
+use foundry_config::{fs_permissions::PathPermission, Config, FsPermissions};
+use foundry_test_utils::{forgetest_async, util, Filter, TestProject};
 
 #[tokio::test(flavor = "multi_thread")]
 async fn test_zk_cheat_roll_works() {
@@ -140,3 +145,35 @@ async fn test_zk_zk_vm_skip_works() {
 
     TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
 }
+
+forgetest_async!(test_zk_use_factory_dep, |prj, cmd| {
+    setup_deploy_prj(&mut prj);
+
+    cmd.forge_fuse();
+    run_zk_script_test(
+        prj.root(),
+        &mut cmd,
+        "./script/DeployCounterWithBytecodeHash.s.sol",
+        "DeployCounterWithBytecodeHash",
+        Some("transmissions11/solmate@v7 OpenZeppelin/openzeppelin-contracts cyfrin/zksync-contracts"),
+        2,
+        Some(&["-vvvvv", "--via-ir", "--system-mode", "true"]),
+    );
+});
+
+fn setup_deploy_prj(prj: &mut TestProject) {
+    util::initialize(prj.root());
+    let permissions = FsPermissions::new(vec![
+        PathPermission::read(Path::new("zkout/Counter.sol/Counter.json")),
+        PathPermission::read(Path::new("zkout/Factory.sol/Factory.json")),
+    ]);
+    let config = Config { fs_permissions: permissions, ..Default::default() };
+    prj.write_config(config);
+    prj.add_script(
+        "DeployCounterWithBytecodeHash.s.sol",
+        include_str!("../../fixtures/zk/DeployCounterWithBytecodeHash.s.sol"),
+    )
+    .unwrap();
+    prj.add_source("Factory.sol", include_str!("../../fixtures/zk/Factory.sol")).unwrap();
+    prj.add_source("Counter", "contract Counter {}").unwrap();
+}