diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index 88d09bfba..616787055 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/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(&mut self, ccx: &mut CheatsCtxt, message: String) { self.get_inspector::(ccx.state).console_log(message); } + + fn trace( + &mut self, + ccx_state: &mut Cheatcodes, + ecx: &mut EvmContext, + call_traces: Vec, + ) { + self.get_inspector::(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, mut input: Input, + executor: &mut impl CheatcodesExecutor, ) -> Option 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( + &mut self, + ecx: &mut EvmContext, + call: &mut CreateInputs, + executor: &mut impl CheatcodesExecutor, + ) -> Option { + self.create_common(ecx, call, executor) + } + pub fn call_with_executor( &mut self, ecx: &mut EvmContext, @@ -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 Inspector for Cheatcodes { ecx: &mut EvmContext, call: &mut CreateInputs, ) -> Option { - self.create_common(ecx, call) + self.create_common(ecx, call, &mut TransparentCheatcodesExecutor) } fn create_end( @@ -1873,7 +1897,7 @@ impl Inspector for Cheatcodes { ecx: &mut EvmContext, call: &mut EOFCreateInputs, ) -> Option { - self.create_common(ecx, call) + self.create_common(ecx, call, &mut TransparentCheatcodesExecutor) } fn eofcreate_end( diff --git a/crates/evm/core/src/lib.rs b/crates/evm/core/src/lib.rs index ed10c5d75..0a399c73a 100644 --- a/crates/evm/core/src/lib.rs +++ b/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: Inspector { // Simulates `console.log` invocation. fn console_log(&mut self, _input: String) {} + + // Appends provided zksync traces. + fn trace_zksync(&mut self, _context: &mut EvmContext, _call_traces: Vec) {} } impl InspectorExt for NoOpInspector {} diff --git a/crates/evm/evm/src/inspectors/mod.rs b/crates/evm/evm/src/inspectors/mod.rs index 41008397a..e0581f2c4 100644 --- a/crates/evm/evm/src/inspectors/mod.rs +++ b/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; diff --git a/crates/evm/evm/src/inspectors/stack.rs b/crates/evm/evm/src/inspectors/stack.rs index b1f28215b..4365d8c61 100644 --- a/crates/evm/evm/src/inspectors/stack.rs +++ b/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, pub log_collector: Option, pub printer: Option, - pub tracer: Option, + pub tracer: Option, pub enable_isolation: bool, /// Flag marking if we are in the inner EVM context. @@ -751,12 +751,23 @@ impl<'a, DB: DatabaseExt> Inspector 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 for InspectorStackRefMut<'a> { inspector, input )); } + + fn trace_zksync(&mut self, ecx: &mut EvmContext, call_traces: Vec) { + call_inspectors!([&mut self.tracer], |inspector| InspectorExt::::trace_zksync( + inspector, + ecx, + call_traces + )); + } } impl Inspector for InspectorStack { diff --git a/crates/evm/evm/src/inspectors/trace.rs b/crates/evm/evm/src/inspectors/trace.rs new file mode 100644 index 000000000..929f4644d --- /dev/null +++ b/crates/evm/evm/src/inspectors/trace.rs @@ -0,0 +1,323 @@ +use alloy_primitives::{Address, Bytes, Log, U256}; +use foundry_evm_core::InspectorExt; +use foundry_evm_traces::{ + CallTraceArena, GethTraceBuilder, ParityTraceBuilder, TracingInspector, TracingInspectorConfig, +}; +use foundry_zksync_core::{ + convert::{ConvertH160, ConvertU256}, + Call, +}; +use revm::{ + interpreter::{ + CallInputs, CallOutcome, CreateInputs, CreateOutcome, EOFCreateInputs, Gas, + InstructionResult, Interpreter, InterpreterResult, + }, + Database, EvmContext, Inspector, +}; + +/// A Wrapper around [TracingInspector] to allow adding zkEVM traces. +#[derive(Clone, Debug, Default)] +pub struct TraceCollector { + inner: TracingInspector, +} + +impl TraceCollector { + /// Returns a new instance for the given config + pub fn new(config: TracingInspectorConfig) -> Self { + Self { inner: TracingInspector::new(config) } + } + + /// Resets the inspector to its initial state of [Self::new]. + /// This makes the inspector ready to be used again. + /// + /// Note that this method has no effect on the allocated capacity of the vector. + #[inline] + pub fn fuse(&mut self) { + self.inner.fuse() + } + + /// Resets the inspector to it's initial state of [Self::new]. + #[inline] + pub fn fused(self) -> Self { + Self { inner: self.inner.fused() } + } + + /// Returns the config of the inspector. + pub const fn config(&self) -> &TracingInspectorConfig { + self.inner.config() + } + + /// Returns a mutable reference to the config of the inspector. + pub fn config_mut(&mut self) -> &mut TracingInspectorConfig { + self.inner.config_mut() + } + + /// Updates the config of the inspector. + pub fn update_config( + &mut self, + f: impl FnOnce(TracingInspectorConfig) -> TracingInspectorConfig, + ) { + self.inner.update_config(f); + } + + /// Gets a reference to the recorded call traces. + pub const fn traces(&self) -> &CallTraceArena { + self.inner.traces() + } + + /// Gets a mutable reference to the recorded call traces. + pub fn traces_mut(&mut self) -> &mut CallTraceArena { + self.inner.traces_mut() + } + + /// Consumes the inspector and returns the recorded call traces. + pub fn into_traces(self) -> CallTraceArena { + self.inner.into_traces() + } + + /// Manually the gas used of the root trace. + /// + /// This is useful if the root trace's gasUsed should mirror the actual gas used by the + /// transaction. + /// + /// This allows setting it manually by consuming the execution result's gas for example. + #[inline] + pub fn set_transaction_gas_used(&mut self, gas_used: u64) { + self.inner.set_transaction_gas_used(gas_used) + } + + /// Convenience function for [ParityTraceBuilder::set_transaction_gas_used] that consumes the + /// type. + #[inline] + pub fn with_transaction_gas_used(self, gas_used: u64) -> Self { + Self { inner: self.inner.with_transaction_gas_used(gas_used) } + } + + /// Consumes the Inspector and returns a [ParityTraceBuilder]. + #[inline] + pub fn into_parity_builder(self) -> ParityTraceBuilder { + self.inner.into_parity_builder() + } + + /// Consumes the Inspector and returns a [GethTraceBuilder]. + #[inline] + pub fn into_geth_builder(self) -> GethTraceBuilder { + self.inner.into_geth_builder() + } +} + +impl Inspector for TraceCollector +where + DB: Database, +{ + #[inline] + fn step(&mut self, interp: &mut Interpreter, context: &mut EvmContext) { + self.inner.step(interp, context) + } + + #[inline] + fn step_end(&mut self, interp: &mut Interpreter, context: &mut EvmContext) { + self.inner.step_end(interp, context) + } + + fn log(&mut self, interp: &mut Interpreter, context: &mut EvmContext, log: &Log) { + self.inner.log(interp, context, log) + } + + fn call( + &mut self, + context: &mut EvmContext, + inputs: &mut CallInputs, + ) -> Option { + self.inner.call(context, inputs) + } + + fn call_end( + &mut self, + context: &mut EvmContext, + inputs: &CallInputs, + outcome: CallOutcome, + ) -> CallOutcome { + self.inner.call_end(context, inputs, outcome) + } + + fn create( + &mut self, + context: &mut EvmContext, + inputs: &mut CreateInputs, + ) -> Option { + self.inner.create(context, inputs) + } + + fn create_end( + &mut self, + context: &mut EvmContext, + inputs: &CreateInputs, + outcome: CreateOutcome, + ) -> CreateOutcome { + self.inner.create_end(context, inputs, outcome) + } + + fn eofcreate( + &mut self, + context: &mut EvmContext, + inputs: &mut EOFCreateInputs, + ) -> Option { + self.inner.eofcreate(context, inputs) + } + + fn eofcreate_end( + &mut self, + context: &mut EvmContext, + inputs: &EOFCreateInputs, + outcome: CreateOutcome, + ) -> CreateOutcome { + self.inner.eofcreate_end(context, inputs, outcome) + } + + fn selfdestruct(&mut self, contract: Address, target: Address, value: U256) { + >::selfdestruct(&mut self.inner, contract, target, value) + } +} + +impl InspectorExt for TraceCollector { + fn trace_zksync(&mut self, context: &mut EvmContext, call_traces: Vec) { + fn trace_call_recursive( + tracer: &mut TracingInspector, + context: &mut EvmContext, + call: Call, + suppressed_top_call: bool, + ) -> u64 { + let inputs = &mut CallInputs { + input: call.input.into(), + gas_limit: call.gas, + scheme: revm::interpreter::CallScheme::Call, + caller: call.from.to_address(), + value: revm::interpreter::CallValue::Transfer(call.value.to_ru256()), + target_address: call.to.to_address(), + bytecode_address: call.to.to_address(), + is_eof: false, + is_static: false, + return_memory_offset: Default::default(), + }; + let is_first_non_system_call = if !suppressed_top_call { + !foundry_zksync_core::is_system_address(inputs.caller) && + !foundry_zksync_core::is_system_address(inputs.target_address) + } else { + false + }; + + // We ignore traces from system addresses, the default account abstraction calls on + // caller address, and the original call (identified when neither `to` or + // `from` are system addresses) since it is already included in EVM trace. + let record_trace = !is_first_non_system_call && + !foundry_zksync_core::is_system_address(inputs.target_address) && + inputs.target_address != context.env.tx.caller; + + let (new_depth, overflow) = context.journaled_state.depth.overflowing_add(1); + if !overflow && record_trace { + context.journaled_state.depth = new_depth; + } + + let mut outcome = if let Some(reason) = &call.revert_reason { + CallOutcome { + result: InterpreterResult { + result: InstructionResult::Revert, + output: reason.as_bytes().to_owned().into(), + gas: Gas::new_spent(call.gas_used), + }, + memory_offset: Default::default(), + } + } else { + CallOutcome { + result: InterpreterResult { + result: InstructionResult::Return, + output: call.output.clone().into(), + gas: Gas::new_spent(call.gas_used), + }, + memory_offset: Default::default(), + } + }; + + let is_create = matches!(call.r#type, foundry_zksync_core::CallType::Create); + let mut create_inputs = if is_create { + Some(CreateInputs { + caller: inputs.caller, + scheme: revm::primitives::CreateScheme::Create, + value: inputs.value.get(), + init_code: inputs.input.clone(), + gas_limit: inputs.gas_limit, + }) + } else { + None + }; + + // start span + if record_trace { + if let Some(inputs) = &mut create_inputs { + tracer.create(context, inputs); + } else { + tracer.call(context, inputs); + } + } + + // recurse into inner calls + // record extra gas from ignored traces, to add it at end + let mut extra_gas = if record_trace { 0u64 } else { call.gas_used }; + for inner_call in call.calls { + let inner_extra_gas = trace_call_recursive( + tracer, + context, + inner_call, + suppressed_top_call || is_first_non_system_call, + ); + extra_gas = extra_gas.saturating_add(inner_extra_gas); + } + + // finish span + if record_trace { + if let Some(inputs) = &mut create_inputs { + let outcome = if let Some(reason) = call.revert_reason { + CreateOutcome { + result: InterpreterResult { + result: InstructionResult::Revert, + output: reason.as_bytes().to_owned().into(), + gas: Gas::new_spent(call.gas_used + extra_gas), + }, + address: None, + } + } else { + // zkEVM traces do not have the create address in output + // it's always an empty slice, so we fill it from the `call.to` + // which contains the correct address + CreateOutcome { + result: InterpreterResult { + result: InstructionResult::Return, + output: Bytes::from(call.to.to_h256().0), + gas: Gas::new_spent(call.gas_used + extra_gas), + }, + address: Some(call.to.to_address()), + } + }; + + tracer.create_end(context, inputs, outcome); + } else { + if extra_gas != 0 { + outcome.result.gas = Gas::new_spent(outcome.result.gas.spent() + extra_gas); + } + tracer.call_end(context, inputs, outcome); + } + } + + if !overflow && record_trace { + context.journaled_state.depth = context.journaled_state.depth.saturating_sub(1); + } + + extra_gas + } + + for call in call_traces { + trace_call_recursive(&mut self.inner, context, call, false); + } + } +} diff --git a/crates/forge/tests/it/zk/logs.rs b/crates/forge/tests/it/zk/logs.rs index 8c0c2a159..ce0036457 100644 --- a/crates/forge/tests/it/zk/logs.rs +++ b/crates/forge/tests/it/zk/logs.rs @@ -7,9 +7,9 @@ use forge::revm::primitives::SpecId; use foundry_test_utils::Filter; #[tokio::test(flavor = "multi_thread")] -async fn test_zk_logs_work() { +async fn test_zk_logs_work_in_call() { let runner = TEST_DATA_DEFAULT.runner_zksync(); - let filter = Filter::new(".*", "ZkConsoleTest", ".*"); + let filter = Filter::new("testZkConsoleOutputDuringCall", "ZkConsoleTest", ".*"); let results = TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).test(); assert_multiple( @@ -17,7 +17,7 @@ async fn test_zk_logs_work() { BTreeMap::from([( "zk/Console.t.sol:ZkConsoleTest", vec![( - "testZkConsoleOutput()", + "testZkConsoleOutputDuringCall()", true, None, Some(vec![ @@ -33,3 +33,31 @@ async fn test_zk_logs_work() { )]), ); } + +#[tokio::test(flavor = "multi_thread")] +async fn test_zk_logs_work_in_create() { + let runner = TEST_DATA_DEFAULT.runner_zksync(); + let filter = Filter::new("testZkConsoleOutputDuringCreate", "ZkConsoleTest", ".*"); + + let results = TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).test(); + assert_multiple( + &results, + BTreeMap::from([( + "zk/Console.t.sol:ZkConsoleTest", + vec![( + "testZkConsoleOutputDuringCreate()", + true, + None, + Some(vec![ + "print".into(), + "outer print".into(), + "0xF9E9ba9Ed9B96AB918c74B21dD0f1D5f2ac38a30".into(), + "print".into(), + "0xff".into(), + "print".into(), + ]), + None, + )], + )]), + ); +} diff --git a/crates/forge/tests/it/zk/mod.rs b/crates/forge/tests/it/zk/mod.rs index 4149c7dfd..084d6b085 100644 --- a/crates/forge/tests/it/zk/mod.rs +++ b/crates/forge/tests/it/zk/mod.rs @@ -13,3 +13,4 @@ mod nft; mod ownership; mod proxy; mod repros; +mod traces; diff --git a/crates/forge/tests/it/zk/traces.rs b/crates/forge/tests/it/zk/traces.rs new file mode 100644 index 000000000..a72c03813 --- /dev/null +++ b/crates/forge/tests/it/zk/traces.rs @@ -0,0 +1,366 @@ +//! Forge tests for zksync logs. + +use crate::{config::*, test_helpers::TEST_DATA_DEFAULT}; +use alloy_primitives::{address, hex, Address, Bytes}; +use forge::{ + revm::primitives::SpecId, + traces::{CallKind, CallTraceArena, CallTraceNode, TraceKind}, +}; +use foundry_test_utils::Filter; +use foundry_zksync_core::convert::ConvertAddress; +use itertools::Itertools; + +const ADDRESS_ZK_TRACE_TEST: Address = address!("7fa9385be102ac3eac297483dd6233d62b3e1496"); +const ADDRESS_ADDER: Address = address!("f9e9ba9ed9b96ab918c74b21dd0f1d5f2ac38a30"); +const ADDRESS_NUMBER: Address = address!("f232f12e115391c535fd519b00efadf042fc8be5"); +const ADDRESS_FIRST_INNER_NUMBER: Address = address!("ed570f3f91621894e001df0fb70bfbd123d3c8ad"); +const ADDRESS_SECOND_INNER_NUMBER: Address = address!("abceaeac3d3a2ac3dcffd7a60ca00a3fac9490ca"); +const ADDRESS_CONSOLE: Address = address!("000000000000000000636f6e736f6c652e6c6f67"); +const SELECTOR_TEST_CALL: Bytes = Bytes::from_static(hex!("0d3282c4").as_slice()); +const SELECTOR_TEST_CREATE: Bytes = Bytes::from_static(hex!("61bdc916").as_slice()); +const SELECTOR_ADD: Bytes = Bytes::from_static(hex!("4f2be91f").as_slice()); +const SELECTOR_FIVE: Bytes = Bytes::from_static(hex!("af11c34c").as_slice()); +const SELECTOR_INNER_FIVE: Bytes = Bytes::from_static(hex!("3a0a858d").as_slice()); +const VALUE_FIVE: Bytes = Bytes::from_static( + hex!("0000000000000000000000000000000000000000000000000000000000000005").as_slice(), +); +const VALUE_TEN: Bytes = Bytes::from_static( + hex!("000000000000000000000000000000000000000000000000000000000000000a").as_slice(), +); +const VALUE_LOG_UINT_TEN: Bytes = Bytes::from_static( + hex!("f5b1bba9000000000000000000000000000000000000000000000000000000000000000a").as_slice(), +); // selector: log(uint) + +#[tokio::test(flavor = "multi_thread")] +async fn test_zk_traces_work_during_call() { + let runner = TEST_DATA_DEFAULT.runner_zksync(); + let filter = Filter::new("testZkTraceOutputDuringCall", "ZkTraceTest", ".*"); + + let results = TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).test(); + let traces = + &results["zk/Trace.t.sol:ZkTraceTest"].test_results["testZkTraceOutputDuringCall()"].traces; + + assert_execution_trace( + vec![TraceAssertion { + kind: Some(CallKind::Call), + address: Some(ADDRESS_ZK_TRACE_TEST), + data: Some(SELECTOR_TEST_CALL), + children: vec![ + TraceAssertion { + kind: Some(CallKind::Create), + address: Some(ADDRESS_ADDER), + output: Some(Bytes::from(ADDRESS_ADDER.to_h256().0)), + ..Default::default() + }, + TraceAssertion { + kind: Some(CallKind::Call), + address: Some(ADDRESS_ADDER), + data: Some(SELECTOR_ADD), + output: Some(VALUE_TEN), + children: vec![ + TraceAssertion { + kind: Some(CallKind::Create), + address: Some(ADDRESS_NUMBER), + output: Some(Bytes::from(ADDRESS_NUMBER.to_h256().0)), + ..Default::default() + }, + TraceAssertion { + kind: Some(CallKind::Call), + address: Some(ADDRESS_NUMBER), + data: Some(SELECTOR_FIVE), + output: Some(VALUE_FIVE), + children: vec![ + TraceAssertion { + kind: Some(CallKind::Create), + address: Some(ADDRESS_FIRST_INNER_NUMBER), + output: Some(Bytes::from( + ADDRESS_FIRST_INNER_NUMBER.to_h256().0, + )), + ..Default::default() + }, + TraceAssertion { + kind: Some(CallKind::Call), + address: Some(ADDRESS_FIRST_INNER_NUMBER), + data: Some(SELECTOR_INNER_FIVE), + output: Some(VALUE_FIVE), + ..Default::default() + }, + ], + }, + TraceAssertion { + kind: Some(CallKind::Call), + address: Some(ADDRESS_NUMBER), + data: Some(SELECTOR_FIVE), + output: Some(VALUE_FIVE), + children: vec![ + TraceAssertion { + kind: Some(CallKind::Create), + address: Some(ADDRESS_SECOND_INNER_NUMBER), + output: Some(Bytes::from( + ADDRESS_SECOND_INNER_NUMBER.to_h256().0, + )), + ..Default::default() + }, + TraceAssertion { + kind: Some(CallKind::Call), + address: Some(ADDRESS_SECOND_INNER_NUMBER), + data: Some(SELECTOR_INNER_FIVE), + output: Some(VALUE_FIVE), + ..Default::default() + }, + ], + }, + ], + }, + TraceAssertion { + kind: Some(CallKind::StaticCall), + address: Some(ADDRESS_CONSOLE), + data: Some(VALUE_LOG_UINT_TEN), + ..Default::default() + }, + ], + ..Default::default() + }], + traces, + ) +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_zk_traces_work_during_create() { + let runner = TEST_DATA_DEFAULT.runner_zksync(); + let filter = Filter::new("testZkTraceOutputDuringCreate", "ZkTraceTest", ".*"); + + let results = TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).test(); + let traces = &results["zk/Trace.t.sol:ZkTraceTest"].test_results + ["testZkTraceOutputDuringCreate()"] + .traces; + + assert_execution_trace( + vec![TraceAssertion { + kind: Some(CallKind::Call), + address: Some(ADDRESS_ZK_TRACE_TEST), + data: Some(SELECTOR_TEST_CREATE), + children: vec![TraceAssertion { + kind: Some(CallKind::Create), + address: Some(ADDRESS_ADDER), + output: Some(Bytes::from(ADDRESS_ADDER.to_h256().0)), + children: vec![ + TraceAssertion { + kind: Some(CallKind::Create), + address: Some(ADDRESS_NUMBER), + output: Some(Bytes::from(ADDRESS_NUMBER.to_h256().0)), + ..Default::default() + }, + TraceAssertion { + kind: Some(CallKind::Call), + address: Some(ADDRESS_NUMBER), + data: Some(SELECTOR_FIVE), + output: Some(VALUE_FIVE), + children: vec![ + TraceAssertion { + kind: Some(CallKind::Create), + address: Some(ADDRESS_FIRST_INNER_NUMBER), + output: Some(Bytes::from(ADDRESS_FIRST_INNER_NUMBER.to_h256().0)), + ..Default::default() + }, + TraceAssertion { + kind: Some(CallKind::Call), + address: Some(ADDRESS_FIRST_INNER_NUMBER), + data: Some(SELECTOR_INNER_FIVE), + output: Some(VALUE_FIVE), + ..Default::default() + }, + ], + }, + TraceAssertion { + kind: Some(CallKind::Call), + address: Some(ADDRESS_NUMBER), + data: Some(SELECTOR_FIVE), + output: Some(VALUE_FIVE), + children: vec![ + TraceAssertion { + kind: Some(CallKind::Create), + address: Some(ADDRESS_SECOND_INNER_NUMBER), + output: Some(Bytes::from(ADDRESS_SECOND_INNER_NUMBER.to_h256().0)), + ..Default::default() + }, + TraceAssertion { + kind: Some(CallKind::Call), + address: Some(ADDRESS_SECOND_INNER_NUMBER), + data: Some(SELECTOR_INNER_FIVE), + output: Some(VALUE_FIVE), + ..Default::default() + }, + ], + }, + TraceAssertion { + kind: Some(CallKind::Call), + address: Some(ADDRESS_CONSOLE), + data: Some(VALUE_LOG_UINT_TEN), + ..Default::default() + }, + ], + ..Default::default() + }], + ..Default::default() + }], + traces, + ) +} + +#[derive(Default, Debug)] +struct TraceAssertion { + kind: Option, + address: Option
, + data: Option, + output: Option, + children: Vec, +} + +/// Assert that the execution trace matches the actual trace. +fn assert_execution_trace(expected: Vec, traces: &[(TraceKind, CallTraceArena)]) { + #[allow(dead_code)] + #[derive(Debug)] + struct AssertionFailure { + field: String, + expected: String, + actual: String, + path: Vec, + } + + fn assert_recursive( + expected: &[TraceAssertion], + actual: &[DecodedTrace], + ) -> Option { + for (idx, expected_node) in expected.iter().enumerate() { + let actual_node = match actual.get(idx) { + Some(actual) => actual, + None => { + return Some(AssertionFailure { + field: "".to_string(), + expected: "".to_string(), + actual: "".to_string(), + path: vec![idx], + }) + } + }; + if let Some(kind) = expected_node.kind { + if kind != actual_node.kind { + return Some(AssertionFailure { + field: "kind".to_string(), + expected: format!("{kind:?}"), + actual: format!("{:?}", actual_node.kind), + path: vec![idx], + }); + } + } + if let Some(address) = expected_node.address { + if address != actual_node.address { + return Some(AssertionFailure { + field: "address".to_string(), + expected: format!("{address:?}"), + actual: format!("{:?}", actual_node.address), + path: vec![idx], + }); + } + } + if let Some(data) = &expected_node.data { + if data != &actual_node.data { + return Some(AssertionFailure { + field: "data".to_string(), + expected: format!("{data:?}"), + actual: format!("{:?}", actual_node.data), + path: vec![idx], + }); + } + } + if let Some(output) = &expected_node.output { + if output != &actual_node.output { + return Some(AssertionFailure { + field: "output".to_string(), + expected: format!("{output:?}"), + actual: format!("{:?}", actual_node.output), + path: vec![idx], + }); + } + } + + if let Some(mut failure) = + assert_recursive(&expected_node.children, &actual_node.children) + { + failure.path.insert(0, idx); + return Some(failure) + } + } + None + } + + let actual = decode_first_execution_trace(traces); + if let Some(failure) = assert_recursive(&expected, &actual) { + println!("---"); + println!("{failure:#?}"); + println!("---"); + println!("Trace:"); + let mut actual = &actual; + for (depth, idx) in failure.path.iter().enumerate() { + let trace = &actual[*idx]; + println!( + "{}{:?} {:?} {:?} {:?}", + " ".repeat(depth), + trace.kind, + trace.address, + trace.data, + trace.output + ); + + actual = &trace.children; + } + println!("---\n"); + panic!("trace assertion failure"); + } +} + +/// Represents the decoded trace. +#[allow(dead_code)] +#[derive(Debug, Default)] +struct DecodedTrace { + idx: usize, + depth: usize, + kind: CallKind, + address: Address, + data: Bytes, + output: Bytes, + children: Vec, +} + +/// Decodes and returns the first execution trace. +fn decode_first_execution_trace(traces: &[(TraceKind, CallTraceArena)]) -> Vec { + fn decode_recursive(nodes: &[CallTraceNode], node: &CallTraceNode) -> DecodedTrace { + let children = + node.children.iter().map(|idx| decode_recursive(nodes, &nodes[*idx])).collect_vec(); + DecodedTrace { + idx: node.idx, + depth: node.trace.depth, + kind: node.trace.kind, + address: node.trace.address, + data: node.trace.data.clone(), + output: node.trace.output.clone(), + children, + } + } + + traces + .iter() + .find(|(kind, _)| matches!(kind, TraceKind::Execution)) + .map(|(_, trace)| { + let mut decoded_nodes = vec![]; + let nodes = trace.nodes(); + nodes + .iter() + .filter(|node| node.parent.is_none()) + .for_each(|node| decoded_nodes.push(decode_recursive(nodes, node))); + + decoded_nodes + }) + .unwrap_or_default() +} diff --git a/crates/zksync/core/src/lib.rs b/crates/zksync/core/src/lib.rs index 60bc51c8f..5173be1e0 100644 --- a/crates/zksync/core/src/lib.rs +++ b/crates/zksync/core/src/lib.rs @@ -35,6 +35,7 @@ pub use vm::{balance, encode_create_params, nonce}; use zksync_types::utils::storage_key_for_eth_balance; pub use zksync_types::{ + vm_trace::{Call, CallType}, ACCOUNT_CODE_STORAGE_ADDRESS, CONTRACT_DEPLOYER_ADDRESS, H256, L2_BASE_TOKEN_ADDRESS, NONCE_HOLDER_ADDRESS, }; diff --git a/crates/zksync/core/src/vm/inspect.rs b/crates/zksync/core/src/vm/inspect.rs index 83268f88d..48d1e0876 100644 --- a/crates/zksync/core/src/vm/inspect.rs +++ b/crates/zksync/core/src/vm/inspect.rs @@ -59,6 +59,8 @@ pub struct ZKVMExecutionResult { pub logs: Vec, /// The result of a given execution pub execution_result: rExecutionResult, + /// Call traces + pub call_traces: Vec, } /// Revm-style result with ZKVM Execution @@ -99,15 +101,23 @@ where match (&mut aggregated_result, result.execution_result) { (_, exec @ rExecutionResult::Revert { .. } | exec @ rExecutionResult::Halt { .. }) => { - return Ok(ZKVMExecutionResult { logs: result.logs, execution_result: exec }); + return Ok(ZKVMExecutionResult { + logs: result.logs, + call_traces: result.call_traces, + execution_result: exec, + }); } (None, exec) => { - aggregated_result - .replace(ZKVMExecutionResult { logs: result.logs, execution_result: exec }); + aggregated_result.replace(ZKVMExecutionResult { + logs: result.logs, + call_traces: result.call_traces, + execution_result: exec, + }); } ( Some(ZKVMExecutionResult { logs: aggregated_logs, + call_traces: aggregated_call_traces, execution_result: rExecutionResult::Success { reason: agg_reason, @@ -120,6 +130,7 @@ where rExecutionResult::Success { reason, gas_used, gas_refunded, logs, output }, ) => { aggregated_logs.append(&mut result.logs); + aggregated_call_traces.append(&mut result.call_traces); *agg_reason = reason; *agg_gas_used += gas_used; *agg_gas_refunded += gas_refunded; @@ -177,7 +188,7 @@ where let storage_ptr = StorageView::new(&mut era_db, modified_storage_keys, tx.common_data.initiator_address) .into_rc_ptr(); - let (tx_result, bytecodes, modified_storage) = + let InnerZkVmResult { tx_result, bytecodes, modified_storage, call_traces } = inspect_inner(tx, storage_ptr, chain_id, ccx, call_ctx); if let Some(record) = &mut era_db.accesses { @@ -222,6 +233,7 @@ where ZKVMExecutionResult { logs: logs.clone(), + call_traces, execution_result: rExecutionResult::Success { reason: SuccessReason::Return, gas_used: tx_result.statistics.gas_used, @@ -240,6 +252,7 @@ where ZKVMExecutionResult { logs, + call_traces, execution_result: rExecutionResult::Revert { gas_used: env_tx_gas_limit - tx_result.refunds.gas_refunded, output: Bytes::from(output), @@ -255,6 +268,7 @@ where ZKVMExecutionResult { logs, + call_traces, execution_result: rExecutionResult::Halt { reason: mapped_reason, gas_used: env_tx_gas_limit - tx_result.refunds.gas_refunded, @@ -331,13 +345,20 @@ where Ok(execution_result) } +struct InnerZkVmResult { + tx_result: VmExecutionResultAndLogs, + bytecodes: HashMap>, + modified_storage: HashMap, + call_traces: Vec, +} + fn inspect_inner( l2_tx: L2Tx, storage: StoragePtr>, chain_id: L2ChainId, ccx: &mut CheatcodeTracerContext, call_ctx: CallContext, -) -> (VmExecutionResultAndLogs, HashMap>, HashMap) { +) -> InnerZkVmResult { let l1_gas_price = call_ctx.block_basefee.to::().max(MAX_L1_GAS_PRICE); let fair_l2_gas_price = call_ctx.block_basefee.saturating_to::(); let batch_env = create_l1_batch_env(storage.clone(), l1_gas_price, fair_l2_gas_price); @@ -426,12 +447,13 @@ fn inspect_inner( .expect("failed converting bytecode to factory dep") }) .collect(); - let modified_keys = if is_static { + let modified_storage = if is_static { Default::default() } else { storage.borrow().modified_storage_keys().clone() }; - (tx_result, bytecodes, modified_keys) + + InnerZkVmResult { tx_result, bytecodes, modified_storage, call_traces } } /// Parse solidity's `console.log` events diff --git a/testdata/zk/Console.t.sol b/testdata/zk/Console.t.sol index df3f4e7c1..9bc93301f 100644 --- a/testdata/zk/Console.t.sol +++ b/testdata/zk/Console.t.sol @@ -11,10 +11,22 @@ contract Printer { } } +contract ConstructorPrinter { + constructor() { + Printer printer = new Printer(); + printer.print(); + console.log("outer print"); + console.logAddress(address(this)); + printer.print(); + console.logBytes1(0xff); + printer.print(); + } +} + contract ZkConsoleTest is DSTest { Vm constant vm = Vm(HEVM_ADDRESS); - function testZkConsoleOutput() public { + function testZkConsoleOutputDuringCall() public { vm.zkVm(true); Printer printer = new Printer(); @@ -25,4 +37,10 @@ contract ZkConsoleTest is DSTest { console.logBytes1(0xff); printer.print(); } + + function testZkConsoleOutputDuringCreate() public { + vm.zkVm(true); + + new ConstructorPrinter(); + } } diff --git a/testdata/zk/Trace.t.sol b/testdata/zk/Trace.t.sol new file mode 100644 index 000000000..42600c5ce --- /dev/null +++ b/testdata/zk/Trace.t.sol @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.18; + +import "ds-test/test.sol"; +import "../default/logs/console.sol"; + +contract InnerNumber { + event Value(uint8); + + function innerFive() public returns (uint8) { + emit Value(5); + return 5; + } +} + +contract Number { + function five() public returns (uint8) { + InnerNumber num = new InnerNumber(); + return num.innerFive(); + } +} + +contract Adder { + function add() public returns (uint8) { + Number num = new Number(); + return num.five() + num.five(); + } +} + +contract ConstructorAdder { + constructor() { + Number num = new Number(); + uint8 value = num.five() + num.five(); + console.log(value); + assert(value == 10); + } +} + +contract ZkTraceTest is DSTest { + // The test must be run with parameter `-vvv` to print traces + function testZkTraceOutputDuringCall() public { + Adder adder = new Adder(); + uint8 value = adder.add(); + assert(value == 10); + console.log(value); + } + + function testZkTraceOutputDuringCreate() public { + new ConstructorAdder(); + } +}