perf(coverage): cache current HitMap, reserve when merging

crates/evm/coverage/src/inspector.rs

@@ -1,25 +1,40 @@
 use crate::{HitMap, HitMaps};
 use alloy_primitives::B256;
 use revm::{interpreter::Interpreter, Database, EvmContext, Inspector};
+use std::ptr::NonNull;
 
 /// Inspector implementation for collecting coverage information.
-#[derive(Clone, Debug, Default)]
+#[derive(Clone, Debug)]
 pub struct CoverageCollector {
+    current_map: NonNull<HitMap>,
+    current_hash: B256,
     maps: HitMaps,
 }
 
+// SAFETY: `current_map` is always valid and points into an allocation managed by self.
+unsafe impl Send for CoverageCollector {}
+unsafe impl Sync for CoverageCollector {}
+
+impl Default for CoverageCollector {
+    fn default() -> Self {
+        Self {
+            current_map: NonNull::dangling(),
+            current_hash: B256::ZERO,
+            maps: Default::default(),
+        }
+    }
+}
+
 impl<DB: Database> Inspector<DB> for CoverageCollector {
     fn initialize_interp(&mut self, interpreter: &mut Interpreter, _context: &mut EvmContext<DB>) {
-        self.maps
-            .entry(*get_contract_hash(interpreter))
-            .or_insert_with(|| HitMap::new(interpreter.contract.bytecode.original_bytes()));
+        get_or_insert_contract_hash(interpreter);
+        self.insert_map(interpreter);
     }
 
     #[inline]
     fn step(&mut self, interpreter: &mut Interpreter, _context: &mut EvmContext<DB>) {
-        if let Some(map) = self.maps.get_mut(get_contract_hash(interpreter)) {
-            map.hit(interpreter.program_counter());
-        }
+        let map = self.get_or_insert_map(interpreter);
+        map.hit(interpreter.program_counter());
     }
 }
 
@@ -28,15 +43,50 @@ impl CoverageCollector {
     pub fn finish(self) -> HitMaps {
         self.maps
     }
+
+    #[inline]
+    fn get_or_insert_map(&mut self, interpreter: &mut Interpreter) -> &mut HitMap {
+        let hash = get_or_insert_contract_hash(interpreter);
+        if self.current_hash != *hash {
+            self.insert_map(interpreter);
+        }
+        unsafe { self.current_map.as_mut() }
+    }
+
+    #[cold]
+    #[inline(never)]
+    fn insert_map(&mut self, interpreter: &Interpreter) {
+        let Some(hash) = interpreter.contract.hash else { eof_panic() };
+        self.current_hash = hash;
+        self.current_map = self
+            .maps
+            .entry(hash)
+            .or_insert_with(|| HitMap::new(interpreter.contract.bytecode.original_bytes()))
+            .into();
+    }
 }
 
 /// Helper function for extracting contract hash used to record coverage hit map.
-/// If contract hash available in interpreter contract is zero (contract not yet created but going
-/// to be created in current tx) then it hash is calculated from contract bytecode.
-fn get_contract_hash(interpreter: &mut Interpreter) -> &B256 {
-    let hash = interpreter.contract.hash.as_mut().expect("coverage does not support EOF");
-    if *hash == B256::ZERO {
-        *hash = interpreter.contract.bytecode.hash_slow();
+///
+/// If the contract hash is zero (contract not yet created but it's going to be created in current
+/// tx) then the hash is calculated from the bytecode.
+#[inline]
+fn get_or_insert_contract_hash(interpreter: &mut Interpreter) -> &B256 {
+    let Some(hash) = interpreter.contract.hash.as_mut() else { eof_panic() };
+    if hash.is_zero() {
+        set_contract_hash(hash, &interpreter.contract.bytecode);
     }
     hash
 }
+
+#[cold]
+#[inline(never)]
+fn set_contract_hash(hash: &mut B256, bytecode: &revm::primitives::Bytecode) {
+    *hash = bytecode.hash_slow();
+}
+
+#[cold]
+#[inline(never)]
+fn eof_panic() -> ! {
+    panic!("coverage does not support EOF");
+}

crates/evm/coverage/src/lib.rs

@@ -174,6 +174,7 @@ impl HitMaps {
 
     /// Merges two `HitMaps`.
     pub fn merge(&mut self, other: Self) {
+        self.reserve(other.len());
         for (code_hash, other) in other.0 {
             self.entry(code_hash).and_modify(|e| e.merge(&other)).or_insert(other);
         }
@@ -211,37 +212,61 @@ pub struct HitMap {
 
 impl HitMap {
     /// Create a new hitmap with the given bytecode.
+    #[inline]
     pub fn new(bytecode: Bytes) -> Self {
-        Self { bytecode, hits: Default::default() }
+        Self { bytecode, hits: HashMap::with_capacity_and_hasher(1024, Default::default()) }
     }
 
     /// Returns the bytecode.
+    #[inline]
     pub fn bytecode(&self) -> &Bytes {
         &self.bytecode
     }
 
     /// Returns the number of hits for the given program counter.
+    #[inline]
     pub fn get(&self, pc: usize) -> Option<NonZeroU32> {
         NonZeroU32::new(self.hits.get(&Self::cvt_pc(pc)).copied().unwrap_or(0))
     }
 
     /// Increase the hit counter by 1 for the given program counter.
+    #[inline]
     pub fn hit(&mut self, pc: usize) {
         self.hits(pc, 1)
     }
 
     /// Increase the hit counter by `hits` for the given program counter.
+    #[inline]
     pub fn hits(&mut self, pc: usize, hits: u32) {
         *self.hits.entry(Self::cvt_pc(pc)).or_default() += hits;
     }
 
     /// Merge another hitmap into this, assuming the bytecode is consistent
     pub fn merge(&mut self, other: &Self) {
-        for (&pc, &hits) in &other.hits {
-            self.hits(pc as usize, hits);
+        self.hits.reserve(other.len());
+        for (pc, hits) in other.iter() {
+            self.hits(pc, hits);
         }
     }
 
+    /// Returns an iterator over all the program counters and their hit counts.
+    #[inline]
+    pub fn iter(&self) -> impl Iterator<Item = (usize, u32)> + '_ {
+        self.hits.iter().map(|(&pc, &hits)| (pc as usize, hits))
+    }
+
+    /// Returns the number of program counters hit in the hitmap.
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.hits.len()
+    }
+
+    /// Returns `true` if the hitmap is empty.
+    #[inline]
+    pub fn is_empty(&self) -> bool {
+        self.hits.is_empty()
+    }
+
     #[inline]
     fn cvt_pc(pc: usize) -> u32 {
         pc.try_into().expect("4GiB bytecode")