foundry_evm_coverage/
inspector.rs

Help
1use crate::{HitMap, HitMaps};
2use alloy_primitives::B256;
3use revm::{
4    Inspector,
5    context::ContextTr,
6    inspector::JournalExt,
7    interpreter::{Interpreter, interpreter_types::Jumps},
8};
9use std::ptr::NonNull;
10
11/// Inspector implementation for collecting coverage information.
12#[derive(Clone, Debug)]
13pub struct LineCoverageCollector {
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 LineCoverageCollector {}
25unsafe impl Sync for LineCoverageCollector {}
26
27impl Default for LineCoverageCollector {
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 LineCoverageCollector
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    fn step(&mut self, interpreter: &mut Interpreter, _context: &mut CTX) {
47        let map = self.get_or_insert_map(interpreter);
48        map.hit(interpreter.bytecode.pc() as u32);
49    }
50}
51
52impl LineCoverageCollector {
53    /// Finish collecting coverage information and return the [`HitMaps`].
54    pub fn finish(self) -> HitMaps {
55        self.maps
56    }
57
58    /// Gets the hit map for the current contract, or inserts a new one if it doesn't exist.
59    ///
60    /// The map is stored in `current_map` and returned as a mutable reference.
61    /// See comments on `current_map` for more details.
62    #[inline]
63    fn get_or_insert_map(&mut self, interpreter: &mut Interpreter) -> &mut HitMap {
64        let hash = get_or_insert_contract_hash(interpreter);
65        if self.current_hash != *hash {
66            self.insert_map(interpreter);
67        }
68        // SAFETY: See comments on `current_map`.
69        unsafe { self.current_map.as_mut() }
70    }
71
72    #[cold]
73    #[inline(never)]
74    fn insert_map(&mut self, interpreter: &mut Interpreter) {
75        let hash = interpreter.bytecode.hash().unwrap();
76        self.current_hash = hash;
77        // Converts the mutable reference to a `NonNull` pointer.
78        self.current_map = self
79            .maps
80            .entry(hash)
81            .or_insert_with(|| HitMap::new(interpreter.bytecode.original_bytes()))
82            .into();
83    }
84}
85
86/// Helper function for extracting contract hash used to record coverage hit map.
87///
88/// If the contract hash is zero (contract not yet created but it's going to be created in current
89/// tx) then the hash is calculated from the bytecode.
90#[inline]
91fn get_or_insert_contract_hash(interpreter: &mut Interpreter) -> B256 {
92    // TODO: use just `get_or_calculate_hash`
93    interpreter
94        .bytecode
95        .hash()
96        .filter(|h| !h.is_zero())
97        .unwrap_or_else(|| interpreter.bytecode.regenerate_hash())
98}
foundry_evm_coverage/inspector.rs

foundry_evm_coverage/
inspector.rs
