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    }
51
52    fn step(&mut self, interpreter: &mut Interpreter, _context: &mut CTX) {
53        let pc = interpreter.bytecode.pc();
54        self.get_or_insert_map(interpreter).record(pc);
55    }
56
57    fn create_end(
58        &mut self,
59        _context: &mut CTX,
60        inputs: &CreateInputs,
61        outcome: &mut CreateOutcome,
62    ) {
63        if outcome.result.result.is_ok()
64            && let Some(map) = self.maps.get_mut(&inputs.init_code_hash())
65        {
66            map.creation();
67        }
68    }
69}
70
71impl LineCoverageCollector {
72    /// Finish collecting coverage information and return the [`HitMaps`].
73    pub fn finish(mut self) -> HitMaps {
74        self.maps.values_mut().for_each(HitMap::flush);
75        self.maps
76    }
77
78    /// Gets the hit map for the current contract, or inserts a new one if it doesn't exist.
79    ///
80    /// The map is stored in `current_map` and returned as a mutable reference.
81    /// See comments on `current_map` for more details.
82    #[inline]
83    fn get_or_insert_map(&mut self, interpreter: &mut Interpreter) -> &mut HitMap {
84        let hash = interpreter.bytecode.get_or_calculate_hash();
85        if self.current_hash != *hash {
86            self.insert_map(interpreter);
87        }
88        // SAFETY: See comments on `current_map`.
89        unsafe { self.current_map.as_mut() }
90    }
91
92    #[cold]
93    #[inline(never)]
94    fn insert_map(&mut self, interpreter: &mut Interpreter) {
95        let hash = interpreter.bytecode.hash().unwrap();
96        self.current_hash = hash;
97        // Converts the mutable reference to a `NonNull` pointer.
98        self.current_map = self
99            .maps
100            .entry(hash)
101            .or_insert_with(|| HitMap::new(interpreter.bytecode.original_bytes()))
102            .into();
103    }
104}
foundry_evm_coverage/inspector.rs

foundry_evm_coverage/
inspector.rs
