From 2888a3ce143db5fc81f4ef67ca25dd758637af27 Mon Sep 17 00:00:00 2001
From: Jrigada <jrigada@frba.utn.edu.ar>
Date: Thu, 19 Sep 2024 16:29:15 -0300
Subject: [PATCH] Add check in zk create to check whether we are deploying the
 test contract or not

---
 crates/cheatcodes/src/inspector.rs | 17 +++++++++++++++++
 crates/script/src/runner.rs        |  2 +-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs
index d3acedfa5..86da6b20f 100644
--- a/crates/cheatcodes/src/inspector.rs
+++ b/crates/cheatcodes/src/inspector.rs
@@ -939,6 +939,23 @@ impl Cheatcodes {
             return None
         }
 
+        if let Some(CreateScheme::Create) = input.scheme() {
+            let caller = input.caller();
+            let nonce = ecx
+                .inner
+                .journaled_state
+                .load_account(input.caller(), &mut ecx.inner.db)
+                .unwrap()
+                .0
+                .info
+                .nonce;
+            let address = caller.create(nonce);
+            if let Some(true) = ecx.db.get_test_contract_address().map(|addr| address == addr) {
+                info!("running create in EVM, instead of zkEVM (Test Contract) {:#?}", address);
+                return None
+            }
+        }
+
         if input.init_code().0 == DEFAULT_CREATE2_DEPLOYER_CODE {
             info!("running create in EVM, instead of zkEVM (DEFAULT_CREATE2_DEPLOYER_CODE)");
             return None
diff --git a/crates/script/src/runner.rs b/crates/script/src/runner.rs
index 300f57809..136973e4b 100644
--- a/crates/script/src/runner.rs
+++ b/crates/script/src/runner.rs
@@ -130,6 +130,7 @@ impl ScriptRunner {
 
         let address = CALLER.create(self.executor.get_nonce(CALLER)?);
 
+        self.executor.backend_mut().set_test_contract(address);
         // Set the contracts initial balance before deployment, so it is available during the
         // construction
         self.executor.set_balance(address, self.evm_opts.initial_balance)?;
@@ -147,7 +148,6 @@ impl ScriptRunner {
 
         // Optionally call the `setUp` function
         let (success, gas_used, labeled_addresses, transactions) = if !setup {
-            self.executor.backend_mut().set_test_contract(address);
             (true, 0, Default::default(), Some(library_transactions))
         } else {
             match self.executor.setup(Some(self.evm_opts.sender), address, None) {