From 758eb86dd56c1fa13f6c8dc90a1733e5b3587ec7 Mon Sep 17 00:00:00 2001
From: Nisheeth Barthwal <nbaztec@gmail.com>
Date: Thu, 12 Sep 2024 21:56:28 +0200
Subject: [PATCH] sticky addresses deployed in skip

---
 crates/cheatcodes/src/inspector.rs | 24 +++++++++++++++++++++++-
 testdata/zk/Cheatcodes.t.sol       | 14 +++++++++++---
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs
index b3d76d72b..a390c954f 100644
--- a/crates/cheatcodes/src/inspector.rs
+++ b/crates/cheatcodes/src/inspector.rs
@@ -51,7 +51,7 @@ use revm::{
 use rustc_hash::FxHashMap;
 use serde_json::Value;
 use std::{
-    collections::{BTreeMap, HashMap, VecDeque},
+    collections::{BTreeMap, HashMap, HashSet, VecDeque},
     fs::File,
     io::BufReader,
     ops::Range,
@@ -351,6 +351,14 @@ pub struct Cheatcodes {
     /// When in zkEVM context, execute the next CALL or CREATE in the EVM instead.
     pub skip_zk_vm: bool,
 
+    /// Any contracts that were deployed in `skip_zk_vm` step.
+    /// This makes it easier to dispatch calls to any of these addresses in zkEVM context, directly
+    /// to EVM. Alternatively, we'd need to add `vm.zkVmSkip()` to these calls manually.
+    pub skip_zk_vm_addresses: HashSet<Address>,
+
+    /// Records the next create address for `skip_zk_vm_addresses`.
+    pub record_next_create_address: bool,
+
     /// Dual compiled contracts
     pub dual_compiled_contracts: DualCompiledContracts,
 
@@ -448,6 +456,8 @@ impl Cheatcodes {
             breakpoints: Default::default(),
             use_zk_vm: Default::default(),
             skip_zk_vm: Default::default(),
+            skip_zk_vm_addresses: Default::default(),
+            record_next_create_address: Default::default(),
             persisted_factory_deps: Default::default(),
         }
     }
@@ -900,6 +910,7 @@ impl Cheatcodes {
             if self.skip_zk_vm {
                 // reset skip_zk_vm
                 self.skip_zk_vm = false;
+                self.record_next_create_address = true;
                 info!("running create in EVM, instead of zkEVM");
             } else {
                 info!("running create in zkEVM");
@@ -1123,6 +1134,14 @@ impl Cheatcodes {
                 }
             }
         }
+
+        if self.record_next_create_address {
+            self.record_next_create_address = false;
+            if let Some(address) = outcome.address {
+                self.skip_zk_vm_addresses.insert(address);
+            }
+        }
+
         outcome
     }
 
@@ -1412,6 +1431,9 @@ impl Cheatcodes {
         }
 
         if self.use_zk_vm {
+            // also skip if the target was created during a zkEVM skip
+            self.skip_zk_vm =
+                self.skip_zk_vm || self.skip_zk_vm_addresses.contains(&call.target_address);
             if self.skip_zk_vm {
                 // reset skip_zk_vm
                 self.skip_zk_vm = false;
diff --git a/testdata/zk/Cheatcodes.t.sol b/testdata/zk/Cheatcodes.t.sol
index 509d4760a..84fcdd955 100644
--- a/testdata/zk/Cheatcodes.t.sol
+++ b/testdata/zk/Cheatcodes.t.sol
@@ -360,7 +360,7 @@ contract ZkCheatcodesInZkVmTest is DSTest {
 
 contract Calculator {
     event Added(uint8 indexed sum);
-    
+
     function add(uint8 a, uint8 b) public returns (uint8) {
         uint8 sum = a + b;
         emit Added(sum);
@@ -396,12 +396,20 @@ contract ZkCheatcodeZkVmSkipTest is DSTest {
         vm.makePersistent(address(helper));
     }
 
-    function testFail_ContractDeploysInZkEvmWithoutSkip() external {
+    function testFail_ContractUsesCheatcodesInEvmWithSkip() external {
         helper.exec();
     }
 
-    function testContractDeploysInEvmWithSkip() external {
+    function testContractUsesCheatcodesInEvmWithSkip() external {
         vm.zkVmSkip();
         helper.exec();
     }
+
+    function testContractDeploysInEvmWithSkipAndCallsAutoSkip() external {
+        vm.zkVmSkip();
+        EvmTargetContract helper2 = new EvmTargetContract();
+
+        // this should auto execute in EVM
+        helper2.exec();
+    }
 }