foundry_evm_coverage/
inspector.rs

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        let map = self.get_or_insert_map(interpreter);
43        // Reserve some space early to avoid reallocating too often.
44        map.reserve(8192.min(interpreter.bytecode.len()));
45    }
46
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 LineCoverageCollector {
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 = interpreter.bytecode.get_or_calculate_hash();
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();
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}
foundry_evm_coverage/inspector.rs

foundry_evm_coverage/
inspector.rs
