From 55b9b909c790507e59e04fa22480c2745feb8741 Mon Sep 17 00:00:00 2001 From: Nisheeth Barthwal Date: Fri, 13 Sep 2024 17:45:41 +0200 Subject: [PATCH] refactor code for readability --- crates/cheatcodes/src/inspector.rs | 419 +++++++++++++++-------------- testdata/zk/Cheatcodes.t.sol | 6 +- 2 files changed, 227 insertions(+), 198 deletions(-) diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index a390c954f..13179ef00 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -907,118 +907,134 @@ impl Cheatcodes { } if self.use_zk_vm { - 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"); - if input.init_code().0 == DEFAULT_CREATE2_DEPLOYER_CODE { - info!("ignoring DEFAULT_CREATE2_DEPLOYER_CODE for zk"); - return None - } + if let Some(result) = self.zk_handle_create(ecx, input, executor) { + return Some(result); + } + } - let zk_contract = self - .dual_compiled_contracts - .find_by_evm_bytecode(&input.init_code().0) - .unwrap_or_else(|| { - panic!("failed finding contract for {:?}", input.init_code()) - }); + None + } - let factory_deps = self.dual_compiled_contracts.fetch_all_factory_deps(zk_contract); - tracing::debug!(contract = zk_contract.name, "using dual compiled contract"); + // Handle the `CREATE` within zkEVM and return the result + fn zk_handle_create( + &mut self, + ecx: &mut EvmContext, + input: Input, + executor: &mut impl CheatcodesExecutor, + ) -> Option + where + DB: DatabaseExt, + Input: CommonCreateInput, + { + if self.skip_zk_vm { + self.skip_zk_vm = false; // handled the skip, reset flag + self.record_next_create_address = true; + info!("running create in EVM, instead of zkEVM"); + return None + } - let ccx = foundry_zksync_core::vm::CheatcodeTracerContext { - mocked_calls: self.mocked_calls.clone(), - expected_calls: Some(&mut self.expected_calls), - accesses: self.accesses.as_mut(), - persisted_factory_deps: Some(&mut self.persisted_factory_deps), - }; - let create_inputs = CreateInputs { - scheme: input.scheme().unwrap_or(CreateScheme::Create), - init_code: input.init_code(), - value: input.value(), - caller: input.caller(), - gas_limit: input.gas_limit(), - }; - if let Ok(result) = foundry_zksync_core::vm::create::<_, DatabaseError>( - &create_inputs, - zk_contract, - factory_deps, - ecx, - ccx, - ) { - if let Some(recorded_logs) = &mut self.recorded_logs { - recorded_logs.extend(result.logs.clone().into_iter().map(|log| Vm::Log { - topics: log.data.topics().to_vec(), - data: log.data.data.clone(), - emitter: log.address, - })); - } + info!("running create in zkEVM"); + if input.init_code().0 == DEFAULT_CREATE2_DEPLOYER_CODE { + info!("ignoring DEFAULT_CREATE2_DEPLOYER_CODE for zk"); + return None + } - // append console logs from zkEVM to the current executor's LogTracer - result.logs.iter().filter_map(decode_console_log).for_each(|decoded_log| { - executor.console_log( - &mut CheatsCtxt { - state: self, - ecx: &mut ecx.inner, - precompiles: &mut ecx.precompiles, - gas_limit: create_inputs.gas_limit, - caller: create_inputs.caller, - }, - decoded_log, - ); - }); + let zk_contract = self + .dual_compiled_contracts + .find_by_evm_bytecode(&input.init_code().0) + .unwrap_or_else(|| panic!("failed finding contract for {:?}", input.init_code())); - // append traces - executor.trace(self, ecx, result.call_traces); + let factory_deps = self.dual_compiled_contracts.fetch_all_factory_deps(zk_contract); + tracing::debug!(contract = zk_contract.name, "using dual compiled contract"); - // for each log in cloned logs call handle_expect_emit - if !self.expected_emits.is_empty() { - for log in result.logs { - expect::handle_expect_emit(self, &log); - } - } + let ccx = foundry_zksync_core::vm::CheatcodeTracerContext { + mocked_calls: self.mocked_calls.clone(), + expected_calls: Some(&mut self.expected_calls), + accesses: self.accesses.as_mut(), + persisted_factory_deps: Some(&mut self.persisted_factory_deps), + }; + let create_inputs = CreateInputs { + scheme: input.scheme().unwrap_or(CreateScheme::Create), + init_code: input.init_code(), + value: input.value(), + caller: input.caller(), + gas_limit: input.gas_limit(), + }; + if let Ok(result) = foundry_zksync_core::vm::create::<_, DatabaseError>( + &create_inputs, + zk_contract, + factory_deps, + ecx, + ccx, + ) { + if let Some(recorded_logs) = &mut self.recorded_logs { + recorded_logs.extend(result.logs.clone().into_iter().map(|log| Vm::Log { + topics: log.data.topics().to_vec(), + data: log.data.data.clone(), + emitter: log.address, + })); + } - return match result.execution_result { - ExecutionResult::Success { output, .. } => match output { - Output::Create(bytes, address) => Some(CreateOutcome { - result: InterpreterResult { - result: InstructionResult::Return, - output: bytes, - gas, - }, - address, - }), - _ => Some(CreateOutcome { - result: InterpreterResult { - result: InstructionResult::Revert, - output: Bytes::new(), - gas, - }, - address: None, - }), - }, - ExecutionResult::Revert { output, .. } => Some(CreateOutcome { - result: InterpreterResult { - result: InstructionResult::Revert, - output, - gas, - }, - address: None, - }), - ExecutionResult::Halt { .. } => Some(CreateOutcome { - result: InterpreterResult { - result: InstructionResult::Revert, - output: Bytes::from_iter(String::from("zk vm halted").as_bytes()), - gas, - }, - address: None, - }), - } + // append console logs from zkEVM to the current executor's LogTracer + result.logs.iter().filter_map(decode_console_log).for_each(|decoded_log| { + executor.console_log( + &mut CheatsCtxt { + state: self, + ecx: &mut ecx.inner, + precompiles: &mut ecx.precompiles, + gas_limit: create_inputs.gas_limit, + caller: create_inputs.caller, + }, + decoded_log, + ); + }); + + // 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 { + expect::handle_expect_emit(self, &log); } } + + // We currently exhaust the entire gas for the call as zkEVM returns a very high amount + // of gas that OOGs revm. + let gas = Gas::new(input.gas_limit()); + + return match result.execution_result { + ExecutionResult::Success { output, .. } => match output { + Output::Create(bytes, address) => Some(CreateOutcome { + result: InterpreterResult { + result: InstructionResult::Return, + output: bytes, + gas, + }, + address, + }), + _ => Some(CreateOutcome { + result: InterpreterResult { + result: InstructionResult::Revert, + output: Bytes::new(), + gas, + }, + address: None, + }), + }, + ExecutionResult::Revert { output, .. } => Some(CreateOutcome { + result: InterpreterResult { result: InstructionResult::Revert, output, gas }, + address: None, + }), + ExecutionResult::Halt { .. } => Some(CreateOutcome { + result: InterpreterResult { + result: InstructionResult::Revert, + output: Bytes::from_iter(String::from("zk vm halted").as_bytes()), + gas, + }, + address: None, + }), + } } None @@ -1431,107 +1447,120 @@ 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; - info!("running call in EVM, instead of zkEVM {:#?}", call); - } else { - if let TransactTo::Call(test_contract) = ecx.env.tx.transact_to { - if call.bytecode_address == test_contract { - info!("using evm for calls to test contract {:?}", ecx.env); - return None - } - } + if let Some(result) = self.zk_handle_call(ecx, call, executor) { + return Some(result); + } + } - info!("running call in zkEVM {:#?}", call); + None + } - let ccx = foundry_zksync_core::vm::CheatcodeTracerContext { - mocked_calls: self.mocked_calls.clone(), - expected_calls: Some(&mut self.expected_calls), - accesses: self.accesses.as_mut(), - persisted_factory_deps: Some(&mut self.persisted_factory_deps), - }; - if let Ok(result) = - foundry_zksync_core::vm::call::<_, DatabaseError>(call, ecx, ccx) - { - // append console logs from zkEVM to the current executor's LogTracer - result.logs.iter().filter_map(decode_console_log).for_each(|decoded_log| { - executor.console_log( - &mut CheatsCtxt { - state: self, - ecx: &mut ecx.inner, - precompiles: &mut ecx.precompiles, - gas_limit: call.gas_limit, - caller: call.caller, - }, - decoded_log, - ); - }); + // Handle the `CALL` within zkEVM and return the result + fn zk_handle_call( + &mut self, + ecx: &mut EvmContext, + call: &mut CallInputs, + executor: &mut impl CheatcodesExecutor, + ) -> Option + where + DB: DatabaseExt, + { + // 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 { + self.skip_zk_vm = false; // handled the skip, reset flag + info!("running call in EVM, instead of zkEVM {:#?}", call); + return None; + } - // skip log processing for static calls - if !call.is_static { - if let Some(recorded_logs) = &mut self.recorded_logs { - recorded_logs.extend(result.logs.clone().into_iter().map(|log| { - Vm::Log { - topics: log.data.topics().to_vec(), - data: log.data.data.clone(), - emitter: log.address, - } - })); - } + if let TransactTo::Call(test_contract) = ecx.env.tx.transact_to { + if call.bytecode_address == test_contract { + info!("using evm for calls to test contract {:?}", ecx.env); + return None + } + } - // append traces - executor.trace(self, ecx, result.call_traces); + info!("running call in zkEVM {:#?}", call); - // for each log in cloned logs call handle_expect_emit - if !self.expected_emits.is_empty() { - for log in result.logs { - expect::handle_expect_emit(self, &log); - } - } - } + let ccx = foundry_zksync_core::vm::CheatcodeTracerContext { + mocked_calls: self.mocked_calls.clone(), + expected_calls: Some(&mut self.expected_calls), + accesses: self.accesses.as_mut(), + persisted_factory_deps: Some(&mut self.persisted_factory_deps), + }; + if let Ok(result) = foundry_zksync_core::vm::call::<_, DatabaseError>(call, ecx, ccx) { + // append console logs from zkEVM to the current executor's LogTracer + result.logs.iter().filter_map(decode_console_log).for_each(|decoded_log| { + executor.console_log( + &mut CheatsCtxt { + state: self, + ecx: &mut ecx.inner, + precompiles: &mut ecx.precompiles, + gas_limit: call.gas_limit, + caller: call.caller, + }, + decoded_log, + ); + }); - return match result.execution_result { - ExecutionResult::Success { output, .. } => match output { - Output::Call(bytes) => Some(CallOutcome { - result: InterpreterResult { - result: InstructionResult::Return, - output: bytes, - gas, - }, - memory_offset: call.return_memory_offset.clone(), - }), - _ => Some(CallOutcome { - result: InterpreterResult { - result: InstructionResult::Revert, - output: Bytes::new(), - gas, - }, - memory_offset: call.return_memory_offset.clone(), - }), - }, - ExecutionResult::Revert { output, .. } => Some(CallOutcome { - result: InterpreterResult { - result: InstructionResult::Revert, - output, - gas, - }, - memory_offset: call.return_memory_offset.clone(), - }), - ExecutionResult::Halt { .. } => Some(CallOutcome { - result: InterpreterResult { - result: InstructionResult::Revert, - output: Bytes::from_iter(String::from("zk vm halted").as_bytes()), - gas, - }, - memory_offset: call.return_memory_offset.clone(), - }), + // skip log processing for static calls + if !call.is_static { + if let Some(recorded_logs) = &mut self.recorded_logs { + recorded_logs.extend(result.logs.clone().into_iter().map(|log| Vm::Log { + topics: log.data.topics().to_vec(), + data: log.data.data.clone(), + emitter: log.address, + })); + } + + // 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 { + expect::handle_expect_emit(self, &log); } } } + + // We currently exhaust the entire gas for the call as zkEVM returns a very high amount + // of gas that OOGs revm. + let gas = Gas::new(call.gas_limit); + + return match result.execution_result { + ExecutionResult::Success { output, .. } => match output { + Output::Call(bytes) => Some(CallOutcome { + result: InterpreterResult { + result: InstructionResult::Return, + output: bytes, + gas, + }, + memory_offset: call.return_memory_offset.clone(), + }), + _ => Some(CallOutcome { + result: InterpreterResult { + result: InstructionResult::Revert, + output: Bytes::new(), + gas, + }, + memory_offset: call.return_memory_offset.clone(), + }), + }, + ExecutionResult::Revert { output, .. } => Some(CallOutcome { + result: InterpreterResult { result: InstructionResult::Revert, output, gas }, + memory_offset: call.return_memory_offset.clone(), + }), + ExecutionResult::Halt { .. } => Some(CallOutcome { + result: InterpreterResult { + result: InstructionResult::Revert, + output: Bytes::from_iter(String::from("zk vm halted").as_bytes()), + gas, + }, + memory_offset: call.return_memory_offset.clone(), + }), + } } None diff --git a/testdata/zk/Cheatcodes.t.sol b/testdata/zk/Cheatcodes.t.sol index 332a829ed..d19736726 100644 --- a/testdata/zk/Cheatcodes.t.sol +++ b/testdata/zk/Cheatcodes.t.sol @@ -354,16 +354,16 @@ contract ZkCheatcodeZkVmSkipTest is DSTest { vm.makePersistent(address(helper)); } - function testFail_ContractUsesCheatcodesInEvmWithSkip() external { + function testFail_UseCheatcodesInZkVmWithoutSkip() external { helper.exec(); } - function testContractUsesCheatcodesInEvmWithSkip() external { + function testUseCheatcodesInEvmWithSkip() external { vm.zkVmSkip(); helper.exec(); } - function testContractDeploysInEvmWithSkipAndCallsAutoSkip() external { + function testAutoSkipAfterDeployInEvmWithSkip() external { vm.zkVmSkip(); EvmTargetContract helper2 = new EvmTargetContract();