foundry_evm_coverage/
inspector.rs

1use crate::{HitMap, HitMaps};
2use alloy_primitives::B256;
3use revm::{
4    Inspector,
5    interpreter::{Interpreter, interpreter_types::Jumps},
6};
7use std::ptr::NonNull;
8
9/// Inspector implementation for collecting coverage information.
10#[derive(Clone, Debug)]
11pub struct LineCoverageCollector {
12    // NOTE: `current_map` is always a valid reference into `maps`.
13    // It is accessed only through `get_or_insert_map` which guarantees that it's valid.
14    // Both of these fields are unsafe to access directly outside of `*insert_map`.
15    current_map: NonNull<HitMap>,
16    current_hash: B256,
17
18    maps: HitMaps,
19}
20
21// SAFETY: See comments on `current_map`.
22unsafe impl Send for LineCoverageCollector {}
23unsafe impl Sync for LineCoverageCollector {}
24
25impl Default for LineCoverageCollector {
26    fn default() -> Self {
27        Self {
28            current_map: NonNull::dangling(),
29            current_hash: B256::ZERO,
30            maps: Default::default(),
31        }
32    }
33}
34
35impl<CTX> Inspector<CTX> for LineCoverageCollector {
36    fn initialize_interp(&mut self, interpreter: &mut Interpreter, _context: &mut CTX) {
37        let map = self.get_or_insert_map(interpreter);
38        // Reserve some space early to avoid reallocating too often.
39        map.reserve(8192.min(interpreter.bytecode.len()));
40    }
41
42    fn step(&mut self, interpreter: &mut Interpreter, _context: &mut CTX) {
43        let map = self.get_or_insert_map(interpreter);
44        map.hit(interpreter.bytecode.pc() as u32);
45    }
46}
47
48impl LineCoverageCollector {
49    /// Finish collecting coverage information and return the [`HitMaps`].
50    pub fn finish(self) -> HitMaps {
51        self.maps
52    }
53
54    /// Gets the hit map for the current contract, or inserts a new one if it doesn't exist.
55    ///
56    /// The map is stored in `current_map` and returned as a mutable reference.
57    /// See comments on `current_map` for more details.
58    #[inline]
59    fn get_or_insert_map(&mut self, interpreter: &mut Interpreter) -> &mut HitMap {
60        let hash = interpreter.bytecode.get_or_calculate_hash();
61        if self.current_hash != *hash {
62            self.insert_map(interpreter);
63        }
64        // SAFETY: See comments on `current_map`.
65        unsafe { self.current_map.as_mut() }
66    }
67
68    #[cold]
69    #[inline(never)]
70    fn insert_map(&mut self, interpreter: &mut Interpreter) {
71        let hash = interpreter.bytecode.hash().unwrap();
72        self.current_hash = hash;
73        // Converts the mutable reference to a `NonNull` pointer.
74        self.current_map = self
75            .maps
76            .entry(hash)
77            .or_insert_with(|| HitMap::new(interpreter.bytecode.original_bytes()))
78            .into();
79    }
80}
foundry_evm_coverage/inspector.rs

foundry_evm_coverage/
inspector.rs
