foundry_evm_coverage/
inspector.rs

Help
1use crate::{HitMap, HitMaps};
2use alloy_primitives::B256;
3use revm::{interpreter::Interpreter, Database, EvmContext, Inspector};
4use std::ptr::NonNull;
5
6/// Inspector implementation for collecting coverage information.
7#[derive(Clone, Debug)]
8pub struct CoverageCollector {
9    // NOTE: `current_map` is always a valid reference into `maps`.
10    // It is accessed only through `get_or_insert_map` which guarantees that it's valid.
11    // Both of these fields are unsafe to access directly outside of `*insert_map`.
12    current_map: NonNull<HitMap>,
13    current_hash: B256,
14
15    maps: HitMaps,
16}
17
18// SAFETY: See comments on `current_map`.
19unsafe impl Send for CoverageCollector {}
20unsafe impl Sync for CoverageCollector {}
21
22impl Default for CoverageCollector {
23    fn default() -> Self {
24        Self {
25            current_map: NonNull::dangling(),
26            current_hash: B256::ZERO,
27            maps: Default::default(),
28        }
29    }
30}
31
32impl<DB: Database> Inspector<DB> for CoverageCollector {
33    fn initialize_interp(&mut self, interpreter: &mut Interpreter, _context: &mut EvmContext<DB>) {
34        get_or_insert_contract_hash(interpreter);
35        self.insert_map(interpreter);
36    }
37
38    #[inline]
39    fn step(&mut self, interpreter: &mut Interpreter, _context: &mut EvmContext<DB>) {
40        let map = self.get_or_insert_map(interpreter);
41        map.hit(interpreter.program_counter() as u32);
42    }
43}
44
45impl CoverageCollector {
46    /// Finish collecting coverage information and return the [`HitMaps`].
47    pub fn finish(self) -> HitMaps {
48        self.maps
49    }
50
51    /// Gets the hit map for the current contract, or inserts a new one if it doesn't exist.
52    ///
53    /// The map is stored in `current_map` and returned as a mutable reference.
54    /// See comments on `current_map` for more details.
55    #[inline]
56    fn get_or_insert_map(&mut self, interpreter: &mut Interpreter) -> &mut HitMap {
57        let hash = get_or_insert_contract_hash(interpreter);
58        if self.current_hash != *hash {
59            self.insert_map(interpreter);
60        }
61        // SAFETY: See comments on `current_map`.
62        unsafe { self.current_map.as_mut() }
63    }
64
65    #[cold]
66    #[inline(never)]
67    fn insert_map(&mut self, interpreter: &Interpreter) {
68        let Some(hash) = interpreter.contract.hash else { eof_panic() };
69        self.current_hash = hash;
70        // Converts the mutable reference to a `NonNull` pointer.
71        self.current_map = self
72            .maps
73            .entry(hash)
74            .or_insert_with(|| HitMap::new(interpreter.contract.bytecode.original_bytes()))
75            .into();
76    }
77}
78
79/// Helper function for extracting contract hash used to record coverage hit map.
80///
81/// If the contract hash is zero (contract not yet created but it's going to be created in current
82/// tx) then the hash is calculated from the bytecode.
83#[inline]
84fn get_or_insert_contract_hash(interpreter: &mut Interpreter) -> &B256 {
85    let Some(hash) = interpreter.contract.hash.as_mut() else { eof_panic() };
86    if hash.is_zero() {
87        set_contract_hash(hash, &interpreter.contract.bytecode);
88    }
89    hash
90}
91
92#[cold]
93#[inline(never)]
94fn set_contract_hash(hash: &mut B256, bytecode: &revm::primitives::Bytecode) {
95    *hash = bytecode.hash_slow();
96}
97
98#[cold]
99#[inline(never)]
100fn eof_panic() -> ! {
101    panic!("coverage does not support EOF");
102}
foundry_evm_coverage/inspector.rs

foundry_evm_coverage/
inspector.rs
