foundry_evm_coverage/
inspector.rs

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

foundry_evm_coverage/
inspector.rs
