feat: add inner zk traces, fix create not registering logs (#558)

* add inner zk traces, fix create not registering logs

* rename trace to trace_zksync

crates/cheatcodes/src/inspector.rs

@@ -34,7 +34,7 @@ use foundry_evm_core::{
 use foundry_zksync_compiler::{DualCompiledContract, DualCompiledContracts};
 use foundry_zksync_core::{
     convert::{ConvertH160, ConvertH256, ConvertRU256, ConvertU256},
-    get_account_code_key, get_balance_key, get_nonce_key, ZkTransactionMetadata,
+    get_account_code_key, get_balance_key, get_nonce_key, Call, ZkTransactionMetadata,
 };
 use itertools::Itertools;
 use revm::{
@@ -158,6 +158,15 @@ pub trait CheatcodesExecutor {
     fn console_log<DB: DatabaseExt>(&mut self, ccx: &mut CheatsCtxt<DB>, message: String) {
         self.get_inspector::<DB>(ccx.state).console_log(message);
     }
+
+    fn trace<DB: DatabaseExt>(
+        &mut self,
+        ccx_state: &mut Cheatcodes,
+        ecx: &mut EvmContext<DB>,
+        call_traces: Vec<Call>,
+    ) {
+        self.get_inspector::<DB>(ccx_state).trace_zksync(ecx, call_traces);
+    }
 }
 
 /// Basic implementation of [CheatcodesExecutor] that simply returns the [Cheatcodes] instance as an
@@ -729,6 +738,7 @@ impl Cheatcodes {
         &mut self,
         ecx: &mut EvmContext<DB>,
         mut input: Input,
+        executor: &mut impl CheatcodesExecutor,
     ) -> Option<CreateOutcome>
     where
         DB: DatabaseExt,
@@ -926,7 +936,6 @@ impl Cheatcodes {
                 }
 
                 // append console logs from zkEVM to the current executor's LogTracer
-                let executor = &mut TransparentCheatcodesExecutor;
                 result.logs.iter().filter_map(decode_console_log).for_each(|decoded_log| {
                     executor.console_log(
                         &mut CheatsCtxt {
@@ -940,6 +949,9 @@ impl Cheatcodes {
                     );
                 });
 
+                // append traces
+                executor.trace(self, ecx, result.call_traces);
+
                 // for each log in cloned logs call handle_expect_emit
                 if !self.expected_emits.is_empty() {
                     for log in result.logs {
@@ -1102,6 +1114,15 @@ impl Cheatcodes {
         outcome
     }
 
+    pub fn create_with_executor<DB: DatabaseExt>(
+        &mut self,
+        ecx: &mut EvmContext<DB>,
+        call: &mut CreateInputs,
+        executor: &mut impl CheatcodesExecutor,
+    ) -> Option<CreateOutcome> {
+        self.create_common(ecx, call, executor)
+    }
+
     pub fn call_with_executor<DB: DatabaseExt>(
         &mut self,
         ecx: &mut EvmContext<DB>,
@@ -1419,6 +1440,9 @@ impl Cheatcodes {
                         }));
                     }
 
+                    // append traces
+                    executor.trace(self, ecx, result.call_traces);
+
                     // for each log in cloned logs call handle_expect_emit
                     if !self.expected_emits.is_empty() {
                         for log in result.logs {
@@ -1856,7 +1880,7 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
         ecx: &mut EvmContext<DB>,
         call: &mut CreateInputs,
     ) -> Option<CreateOutcome> {
-        self.create_common(ecx, call)
+        self.create_common(ecx, call, &mut TransparentCheatcodesExecutor)
     }
 
     fn create_end(
@@ -1873,7 +1897,7 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
         ecx: &mut EvmContext<DB>,
         call: &mut EOFCreateInputs,
     ) -> Option<CreateOutcome> {
-        self.create_common(ecx, call)
+        self.create_common(ecx, call, &mut TransparentCheatcodesExecutor)
     }
 
     fn eofcreate_end(

crates/evm/core/src/lib.rs

@@ -6,6 +6,7 @@
 #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
 
 use auto_impl::auto_impl;
+use foundry_zksync_core::Call;
 use revm::{inspectors::NoOpInspector, interpreter::CreateInputs, Database, EvmContext, Inspector};
 use revm_inspectors::access_list::AccessListInspector;
 
@@ -47,6 +48,9 @@ pub trait InspectorExt<DB: Database>: Inspector<DB> {
 
     // Simulates `console.log` invocation.
     fn console_log(&mut self, _input: String) {}
+
+    // Appends provided zksync traces.
+    fn trace_zksync(&mut self, _context: &mut EvmContext<DB>, _call_traces: Vec<Call>) {}
 }
 
 impl<DB: Database> InspectorExt<DB> for NoOpInspector {}

crates/evm/evm/src/inspectors/mod.rs

@@ -15,3 +15,6 @@ pub use logs::LogCollector;
 
 mod stack;
 pub use stack::{InspectorData, InspectorStack, InspectorStackBuilder};
+
+mod trace;
+pub use trace::TraceCollector;

crates/evm/evm/src/inspectors/stack.rs

@@ -1,6 +1,5 @@
 use super::{
-    Cheatcodes, CheatsConfig, ChiselState, CoverageCollector, Fuzzer, LogCollector,
-    TracingInspector,
+    Cheatcodes, CheatsConfig, ChiselState, CoverageCollector, Fuzzer, LogCollector, TraceCollector,
 };
 use alloy_primitives::{Address, Bytes, Log, TxKind, U256};
 use foundry_cheatcodes::CheatcodesExecutor;
@@ -10,6 +9,7 @@ use foundry_evm_core::{
 };
 use foundry_evm_coverage::HitMaps;
 use foundry_evm_traces::{CallTraceArena, TraceMode};
+use foundry_zksync_core::Call;
 use revm::{
     inspectors::CustomPrintTracer,
     interpreter::{
@@ -282,7 +282,7 @@ pub struct InspectorStackInner {
     pub fuzzer: Option<Fuzzer>,
     pub log_collector: Option<LogCollector>,
     pub printer: Option<CustomPrintTracer>,
-    pub tracer: Option<TracingInspector>,
+    pub tracer: Option<TraceCollector>,
     pub enable_isolation: bool,
 
     /// Flag marking if we are in the inner EVM context.
@@ -751,12 +751,23 @@ impl<'a, DB: DatabaseExt> Inspector<DB> for InspectorStackRefMut<'a> {
 
         call_inspectors_adjust_depth!(
             #[ret]
-            [&mut self.tracer, &mut self.coverage, &mut self.cheatcodes],
+            [&mut self.tracer, &mut self.coverage],
             |inspector| inspector.create(ecx, create).map(Some),
             self,
             ecx
         );
 
+        ecx.journaled_state.depth += self.in_inner_context as usize;
+        if let Some(cheatcodes) = self.cheatcodes.as_deref_mut() {
+            if let Some(output) = cheatcodes.create_with_executor(ecx, create, self.inner) {
+                if output.result.result != InstructionResult::Continue {
+                    ecx.journaled_state.depth -= self.in_inner_context as usize;
+                    return Some(output);
+                }
+            }
+        }
+        ecx.journaled_state.depth -= self.in_inner_context as usize;
+
         if !matches!(create.scheme, CreateScheme::Create2 { .. }) &&
             self.enable_isolation &&
             !self.in_inner_context &&
@@ -915,6 +926,14 @@ impl<'a, DB: DatabaseExt> InspectorExt<DB> for InspectorStackRefMut<'a> {
             inspector, input
         ));
     }
+
+    fn trace_zksync(&mut self, ecx: &mut EvmContext<DB>, call_traces: Vec<Call>) {
+        call_inspectors!([&mut self.tracer], |inspector| InspectorExt::<DB>::trace_zksync(
+            inspector,
+            ecx,
+            call_traces
+        ));
+    }
 }
 
 impl<DB: DatabaseExt> Inspector<DB> for InspectorStack {