fix: Add missing injection of factory deps (#753)

* Add missing injection and clear it afterwards * Update crates/cheatcodes/src/inspector.rs Co-authored-by: Federico Rodríguez <[email protected]> * Update crates/cheatcodes/src/inspector.rs Co-authored-by: Federico Rodríguez <[email protected]> * Update crates/cheatcodes/src/inspector.rs Co-authored-by: Federico Rodríguez <[email protected]> * clarify clear message * Cargo fmt --------- Co-authored-by: Federico Rodríguez <[email protected]>
matter-labs · Dec 4, 2024 · 2eb72ab · 2eb72ab
1 parent 15b4758
commit 2eb72ab
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 8 deletions.
diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs
@@ -1265,7 +1265,27 @@ impl Cheatcodes {
             constructor_args.to_vec(),
         );
 
-        let factory_deps = self.dual_compiled_contracts.fetch_all_factory_deps(contract);
+        let mut factory_deps = self.dual_compiled_contracts.fetch_all_factory_deps(contract);
+        let injected_factory_deps = self
+            .zk_use_factory_deps
+            .iter()
+            .flat_map(|contract| {
+                let artifact_code = crate::fs::get_artifact_code(self, contract, false)
+                    .inspect(|_| info!(contract, "pushing factory dep"))
+                    .unwrap_or_else(|_| {
+                        panic!(
+                            "failed to get bytecode for injected factory deps contract {contract}"
+                        )
+                    })
+                    .to_vec();
+                let res = self.dual_compiled_contracts.find_bytecode(&artifact_code).unwrap();
+                self.dual_compiled_contracts.fetch_all_factory_deps(res.contract())
+            })
+            .collect_vec();
+        factory_deps.extend(injected_factory_deps);
+
+        // NOTE(zk): Clear injected factory deps so that they are not sent on further transactions
+        self.zk_use_factory_deps.clear();
         tracing::debug!(contract = contract.name, "using dual compiled contract");
 
         let zk_persist_nonce_update = self.zk_persist_nonce_update.check();
@@ -1740,13 +1760,15 @@ where {
                         ecx_inner.journaled_state.state().get_mut(&broadcast.new_origin).unwrap();
 
                     let 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)
+                        let injected_factory_deps = self.zk_use_factory_deps.iter().flat_map(|contract| {
+                            let artifact_code = 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()
+                                .to_vec();
+                            let res = self.dual_compiled_contracts.find_bytecode(&artifact_code).unwrap();
+                            self.dual_compiled_contracts.fetch_all_factory_deps(res.contract())
                         }).collect_vec();
                         factory_deps.extend(injected_factory_deps.clone());
 
@@ -1899,6 +1921,11 @@ where {
 
         info!("running call in zkEVM {:#?}", call);
         let zk_persist_nonce_update = self.zk_persist_nonce_update.check();
+
+        // NOTE(zk): Clear injected factory deps here even though it's actually used in broadcast.
+        // To be consistent with where we clear factory deps in try_create_in_zk.
+        self.zk_use_factory_deps.clear();
+
         let ccx = foundry_zksync_core::vm::CheatcodeTracerContext {
             mocked_calls: self.mocked_calls.clone(),
             expected_calls: Some(&mut self.expected_calls),

diff --git a/crates/forge/tests/fixtures/zk/DeployCounterWithBytecodeHash.s.sol b/crates/forge/tests/fixtures/zk/DeployCounterWithBytecodeHash.s.sol
@@ -17,7 +17,7 @@ contract DeployCounterWithBytecodeHash is Script {
         Factory factory = new Factory(counterBytecodeHash);
         (bool _success,) = address(vm).call(abi.encodeWithSignature("zkUseFactoryDep(string)", "Counter"));
         require(_success, "Cheatcode failed");
-        address counter = factory.deployAccount(salt);
+        address counter = factory.deployContract(salt, abi.encode());
         require(counter != address(0), "Counter deployment failed");
         vm.stopBroadcast();
     }

diff --git a/crates/forge/tests/fixtures/zk/Factory.sol b/crates/forge/tests/fixtures/zk/Factory.sol
@@ -4,25 +4,47 @@ 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 {
+contract AAFactory {
     bytes32 public aaBytecodeHash;
 
     constructor(bytes32 _aaBytecodeHash) {
         aaBytecodeHash = _aaBytecodeHash;
     }
 
-    function deployAccount(bytes32 salt) external returns (address accountAddress) {
+    function deployAccount(bytes32 salt, bytes memory constructorArgs) 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)
+                (salt, aaBytecodeHash, constructorArgs, IContractDeployer.AccountAbstractionVersion.Version1)
             )
         );
         require(success, "Deployment failed");
 
         (accountAddress) = abi.decode(returnData, (address));
     }
 }
+
+contract Factory {
+    bytes32 public bytecodeHash;
+
+    constructor(bytes32 _bytecodeHash) {
+        bytecodeHash = _bytecodeHash;
+    }
+
+
+    function deployContract(bytes32 salt, bytes memory constructorArgs) external returns (address contractAddress) {
+        (bool success, bytes memory returnData) = SystemContractsCaller.systemCallWithReturndata(
+            uint32(gasleft()),
+            address(DEPLOYER_SYSTEM_CONTRACT),
+            uint128(0),
+            abi.encodeCall(DEPLOYER_SYSTEM_CONTRACT.create2, (salt, bytecodeHash, constructorArgs))
+        );
+
+        require(success, "Deployment failed");
+
+        (contractAddress) = abi.decode(returnData, (address));
+    }
+}