foundry_evm_coverage/
inspector.rs

1use crate::{CallData, HitMap, HitMaps};
2use alloy_primitives::B256;
3use revm::{
4    Inspector,
5    interpreter::{CreateInputs, CreateOutcome, 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 call = interpreter.input.bytecode_address.is_some().then(|| {
38            let calldata = if interpreter.input.input.is_empty() {
39                CallData::Empty
40            } else {
41                let input = interpreter.input.input.as_bytes_memory(&interpreter.memory);
42                CallData::new(&input)
43            };
44            (calldata, !interpreter.input.call_value.is_zero())
45        });
46        let map = self.get_or_insert_map(interpreter);
47        if let Some((call, with_value)) = call {
48            map.call(call, with_value);
49        }
50        // Reserve some space early to avoid reallocating too often.
51        map.reserve(8192.min(interpreter.bytecode.len()));
52    }
53
54    fn step(&mut self, interpreter: &mut Interpreter, _context: &mut CTX) {
55        let map = self.get_or_insert_map(interpreter);
56        map.hit(interpreter.bytecode.pc() as u32);
57    }
58
59    fn create_end(
60        &mut self,
61        _context: &mut CTX,
62        inputs: &CreateInputs,
63        outcome: &mut CreateOutcome,
64    ) {
65        if outcome.result.result.is_ok()
66            && let Some(map) = self.maps.get_mut(&inputs.init_code_hash())
67        {
68            map.creation();
69        }
70    }
71}
72
73impl LineCoverageCollector {
74    /// Finish collecting coverage information and return the [`HitMaps`].
75    pub fn finish(self) -> HitMaps {
76        self.maps
77    }
78
79    /// Gets the hit map for the current contract, or inserts a new one if it doesn't exist.
80    ///
81    /// The map is stored in `current_map` and returned as a mutable reference.
82    /// See comments on `current_map` for more details.
83    #[inline]
84    fn get_or_insert_map(&mut self, interpreter: &mut Interpreter) -> &mut HitMap {
85        let hash = interpreter.bytecode.get_or_calculate_hash();
86        if self.current_hash != *hash {
87            self.insert_map(interpreter);
88        }
89        // SAFETY: See comments on `current_map`.
90        unsafe { self.current_map.as_mut() }
91    }
92
93    #[cold]
94    #[inline(never)]
95    fn insert_map(&mut self, interpreter: &mut Interpreter) {
96        let hash = interpreter.bytecode.hash().unwrap();
97        self.current_hash = hash;
98        // Converts the mutable reference to a `NonNull` pointer.
99        self.current_map = self
100            .maps
101            .entry(hash)
102            .or_insert_with(|| HitMap::new(interpreter.bytecode.original_bytes()))
103            .into();
104    }
105}
foundry_evm_coverage/inspector.rs

foundry_evm_coverage/
inspector.rs
