diff --git a/crates/evm/coverage/src/inspector.rs b/crates/evm/coverage/src/inspector.rs index 73d2ff148fa9..07c24ae9e5ef 100644 --- a/crates/evm/coverage/src/inspector.rs +++ b/crates/evm/coverage/src/inspector.rs @@ -2,35 +2,41 @@ use crate::{HitMap, HitMaps}; use alloy_primitives::B256; use revm::{interpreter::Interpreter, Database, EvmContext, Inspector}; +/// Inspector implementation for collecting coverage information. #[derive(Clone, Debug, Default)] pub struct CoverageCollector { - /// Maps that track instruction hit data. - pub maps: HitMaps, + maps: HitMaps, } impl Inspector for CoverageCollector { - #[inline] - fn initialize_interp(&mut self, interp: &mut Interpreter, _context: &mut EvmContext) { + fn initialize_interp(&mut self, interpreter: &mut Interpreter, _context: &mut EvmContext) { self.maps - .entry(get_contract_hash(interp)) - .or_insert_with(|| HitMap::new(interp.contract.bytecode.original_bytes())); + .entry(*get_contract_hash(interpreter)) + .or_insert_with(|| HitMap::new(interpreter.contract.bytecode.original_bytes())); } #[inline] - fn step(&mut self, interp: &mut Interpreter, _context: &mut EvmContext) { + fn step(&mut self, interpreter: &mut Interpreter, _context: &mut EvmContext) { + if let Some(map) = self.maps.get_mut(get_contract_hash(interpreter)) { + map.hit(interpreter.program_counter()); + } + } +} + +impl CoverageCollector { + /// Finish collecting coverage information and return the [`HitMaps`]. + pub fn finish(self) -> HitMaps { self.maps - .entry(get_contract_hash(interp)) - .and_modify(|map| map.hit(interp.program_counter())); } } /// Helper function for extracting contract hash used to record coverage hit map. /// If contract hash available in interpreter contract is zero (contract not yet created but going /// to be created in current tx) then it hash is calculated from contract bytecode. -fn get_contract_hash(interp: &mut Interpreter) -> B256 { - let mut hash = interp.contract.hash.expect("Contract hash is None"); - if hash == B256::ZERO { - hash = interp.contract.bytecode.hash_slow(); +fn get_contract_hash(interpreter: &mut Interpreter) -> &B256 { + let hash = interpreter.contract.hash.as_mut().expect("coverage does not support EOF"); + if *hash == B256::ZERO { + *hash = interpreter.contract.bytecode.hash_slow(); } hash } diff --git a/crates/evm/evm/src/inspectors/stack.rs b/crates/evm/evm/src/inspectors/stack.rs index accbca4fd59c..94ec349b1f53 100644 --- a/crates/evm/evm/src/inspectors/stack.rs +++ b/crates/evm/evm/src/inspectors/stack.rs @@ -470,7 +470,7 @@ impl InspectorStack { .map(|cheatcodes| cheatcodes.labels.clone()) .unwrap_or_default(), traces, - coverage: coverage.map(|coverage| coverage.maps), + coverage: coverage.map(|coverage| coverage.finish()), cheatcodes, chisel_state: chisel_state.and_then(|state| state.state), }