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
+    #[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 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,9 @@ pub struct Cheatcodes {
     /// This is set to `false`, once the startup migration is completed.
     pub startup_zk: bool,
 
+    /// Factory deps stored through `zkUseFactoryDep`
+    pub contracts_as_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 +606,7 @@ impl Cheatcodes {
             record_next_create_address: Default::default(),
             persisted_factory_deps: Default::default(),
             paymaster_params: None,
+            contracts_as_factory_deps: Vec::new(),
         }
     }
 
@@ -982,9 +986,18 @@ impl Cheatcodes {
                             paymaster: paymaster_data.address.to_h160(),
                             paymaster_input: paymaster_data.input.to_vec(),
                         });
-                    if let Some(factory_deps) = zk_tx {
-                        let mut batched =
-                            foundry_zksync_core::vm::batch_factory_dependencies(factory_deps);
+                    if let Some(mut factory_deps) = zk_tx {
+                        for factory_deps_contract in &self.contracts_as_factory_deps {
+                            let bytecode =
+                                crate::fs::get_artifact_code(self, factory_deps_contract, false)
+                                    .unwrap_or_else(|_| panic!("Failed to get bytecode for factory deps contract {factory_deps_contract}"));
+                            info!("Pushing {:?} factory deps bytecode", factory_deps_contract);
+                            factory_deps.push(bytecode.to_vec());
+                        }
+                        self.contracts_as_factory_deps.clear();
+                        let mut batched = foundry_zksync_core::vm::batch_factory_dependencies(
+                            factory_deps.clone(),
+                        );
                         debug!(batches = batched.len(), "splitting factory deps for broadcast");
                         // the last batch is the final one that does the deployment
                         zk_tx = batched.pop();
@@ -1580,7 +1593,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 {
+                        info!("factory deps contract {:?}", self.contracts_as_factory_deps);
+                        for factory_deps_contract in &self.contracts_as_factory_deps {
+                            let bytecode =
+                                crate::fs::get_artifact_code(self, factory_deps_contract, false)
+                                    .unwrap_or_else(|_| panic!("Failed to get bytecode for factory deps contract {factory_deps_contract}"));
+                            info!("Pushing {:?} factory deps bytecode", factory_deps_contract);
+                            factory_deps.push(bytecode.to_vec());
+                        }
+                        self.contracts_as_factory_deps.clear();
+
                         let paymaster_params =
                             self.paymaster_params.clone().map(|paymaster_data| PaymasterParams {
                                 paymaster: paymaster_data.address.to_h160(),
@@ -1601,6 +1624,13 @@ impl Cheatcodes {
                     } else {
                         None
                     };
+                    warn!("zk tx {:?}", zk_tx);
+                    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(),

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.contracts_as_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/AAFactory.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 AAFactory {
+    bytes32 public aaBytecodeHash;
+
+    constructor(bytes32 _aaBytecodeHash) {
+        aaBytecodeHash = _aaBytecodeHash;
+    }
+
+    function deployAccount(bytes32 salt, address owner1, address owner2) 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(owner1, owner2), IContractDeployer.AccountAbstractionVersion.Version1)
+            )
+        );
+        require(success, "Deployment failed");
+
+        (accountAddress) = abi.decode(returnData, (address));
+    }
+}

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/DeployMultisig.s.sol

@@ -0,0 +1,88 @@
+// 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/AAFactory.sol";
+import "../src/TwoUserMultisig.sol";
+
+contract DeployMultisig is Script {
+    function run() external {
+        // Owners for the multisig account
+        // Can be random
+        address owner1 = makeAddr("OWNER_1");
+        address owner2 = makeAddr("OWNER_2");
+
+        // Read artifact file and get the bytecode hash
+        string memory artifact = vm.readFile("zkout/TwoUserMultisig.sol/TwoUserMultisig.json");
+        bytes32 multisigBytecodeHash = vm.parseJsonBytes32(artifact, ".hash");
+        console.log("Bytecode hash: ");
+        console.logBytes32(multisigBytecodeHash);
+        bytes32 salt = "CEACREACREA";
+
+        vm.startBroadcast();
+        AAFactory factory = new AAFactory(multisigBytecodeHash);
+        console.log("Factory deployed at: ", address(factory));
+        (bool _success,) = address(vm).call(abi.encodeWithSignature("zkUseFactoryDep(string)", "TwoUserMultisig"));
+        require(_success, "Deployment failed");
+        factory.deployAccount(salt, owner1, owner2);
+        string memory factoryArtifact = vm.readFile("zkout/AAFactory.sol/AAFactory.json");
+        bytes32 factoryBytecodeHash = vm.parseJsonBytes32(factoryArtifact, ".hash");
+        Create2Factory create2Factory = new Create2Factory();
+        address multisigAddress = create2Factory.create2(salt, factoryBytecodeHash, abi.encode(owner1, owner2));
+        console.log("Multisig deployed at: ", multisigAddress);
+
+        vm.stopBroadcast();
+    }
+}
+
+import {DEPLOYER_SYSTEM_CONTRACT} from "zksync-contracts/zksync-contracts/l2/system-contracts/Constants.sol";
+import {EfficientCall} from "zksync-contracts/zksync-contracts/l2/system-contracts/libraries/EfficientCall.sol";
+import {IContractDeployer} from "zksync-contracts/zksync-contracts/l2/system-contracts/interfaces/IContractDeployer.sol";
+
+/// @custom:security-contact [email protected]
+/// @author Matter Labs
+/// @notice The contract that can be used for deterministic contract deployment.
+contract Create2Factory {
+    /// @notice Function that calls the `create2` method of the `ContractDeployer` contract.
+    /// @dev This function accepts the same parameters as the `create2` function of the ContractDeployer system contract,
+    /// so that we could efficiently relay the calldata.
+    function create2(
+        bytes32, // _salt
+        bytes32, // _bytecodeHash
+        bytes calldata // _input
+    ) external payable returns (address) {
+        _relayCall();
+    }
+
+    /// @notice Function that calls the `create2Account` method of the `ContractDeployer` contract.
+    /// @dev This function accepts the same parameters as the `create2Account` function of the ContractDeployer system contract,
+    /// so that we could efficiently relay the calldata.
+    function create2Account(
+        bytes32, // _salt
+        bytes32, // _bytecodeHash
+        bytes calldata, // _input
+        IContractDeployer.AccountAbstractionVersion // _aaVersion
+    ) external payable returns (address) {
+        _relayCall();
+    }
+
+    /// @notice Function that efficiently relays the calldata to the contract deployer system contract. After that,
+    /// it also relays full result.
+    function _relayCall() internal {
+        bool success = EfficientCall.rawCall({
+            _gas: gasleft(),
+            _address: address(DEPLOYER_SYSTEM_CONTRACT),
+            _value: msg.value,
+            _data: msg.data,
+            _isSystem: true
+        });
+
+        assembly {
+            returndatacopy(0, 0, returndatasize())
+            if success { return(0, returndatasize()) }
+            revert(0, returndatasize())
+        }
+    }
+}

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);
     }
 }