compress instrumentation snapshots

Cargo.toml

@@ -15,6 +15,7 @@ clap = { version = "4.4", features = ["derive"] }
 md5 = "0.7"
 rand = "0.8"
 bitvec = "1.0.1"
+roaring = "0.10.9"
 target-lexicon = "0.13"
 speedy = "0.8"
 colored = "2.1"

src/fuzzer/orc.rs

@@ -415,7 +415,7 @@ impl Orchestrator {
                     opts,
                     spec: self.module.clone(),
                     swarm: swarm.clone(),
-                    covered_functions: None,
+                    instr_only_funcs: None,
                 },
                 swarm,
                 timeout,
@@ -591,7 +591,7 @@ impl Orchestrator {
                 opts,
                 spec: self.module.clone(),
                 swarm: swarm.clone(),
-                covered_functions: self.coverage_is_saturated().then(|| {
+                instr_only_funcs: self.coverage_is_saturated().then(|| {
                     let pass = self.codecov_sess.get_pass::<FunctionCoveragePass>();
                     pass.coverage.iter_covered_keys().collect()
                 }),
@@ -691,20 +691,19 @@ pub(crate) struct OrcPassesGen {
     pub swarm: SwarmConfig,
     pub opts: FeedbackOptions,
     pub spec: Arc<ModuleSpec>,
-    pub covered_functions: Option<HashSet<FuncIdx>>,
+    pub instr_only_funcs: Option<HashSet<FuncIdx>>,
 }
 
 impl std::fmt::Debug for OrcPassesGen {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("OrcPassesGen")
-            .field("swarm", &"...")
             .field("opts", &self.opts)
             .field("spec", &self.spec)
             .field(
-                "covered_functions",
-                &self.covered_functions.as_ref().map(|x| x.len()),
+                "instr_only_funcs (cnt)",
+                &self.instr_only_funcs.as_ref().map(|x| x.len()),
             )
-            .finish()
+            .finish_non_exhaustive()
     }
 }
 
@@ -749,7 +748,7 @@ impl PassesGen for OrcPassesGen {
         // We don't really need to add coverage instrumentation to functions
         // that have not been covered yet if it looks like coverage has settled.
         // This reduces compilation time (~ -30%) and improve performance (~ -25%) by reducing bitmap sizes.
-        let key_filter = |loc: &Location| match &self.covered_functions {
+        let key_filter = |loc: &Location| match &self.instr_only_funcs {
             Some(x) => x.contains(&FuncIdx(loc.function)),
             None => true,
         };

src/fuzzer/worker.rs

@@ -199,6 +199,29 @@ impl Worker {
         self.sess.reset_pass_coverage_keep_saved();
         let _ = self.sess.run(input, &mut self.stats);
         let coverage_snapshot = InstrumentationSnapshot::from(&self.sess.passes);
+        if false {
+            let total = coverage_snapshot
+                .snapshots
+                .iter()
+                .map(|x| x.mem_usage())
+                .sum::<usize>();
+            eprintln!(
+                "Snapshot Mem Usage: {}",
+                humansize::format_size(total, humansize::DECIMAL)
+            );
+            for (pass, snapshot) in self
+                .sess
+                .passes
+                .iter()
+                .zip(coverage_snapshot.snapshots.iter())
+            {
+                eprintln!(
+                    "- {}: {}",
+                    pass.shortcode(),
+                    humansize::format_size(snapshot.mem_usage(), humansize::DECIMAL),
+                );
+            }
+        }
         testcase.metadata_map_mut().insert(coverage_snapshot);
 
         self.gather_trace_metadata(input, &mut testcase);

src/instrumentation/call_params.rs

@@ -200,7 +200,7 @@ impl GlobalsRangePass {
         Self {
             coverage: AssociatedCoverageArray::new(
                 &Self::generate_keys(spec)
-                    .filter(|x| key_filter(&Location::from(*x)))
+                    .filter(key_filter)
                     .collect::<Vec<_>>(),
             ),
         }

src/instrumentation/mem.rs

@@ -192,7 +192,7 @@ impl MemoryStorePrevValRangePass {
         Self {
             coverage: AssociatedCoverageArray::new(
                 &Self::generate_keys(spec)
-                    .filter(|x| key_filter(&Location::from(*x)))
+                    .filter(key_filter)
                     .collect::<Vec<_>>(),
             ),
         }

src/instrumentation/mod.rs

@@ -149,21 +149,22 @@ impl Passes {
 pub(crate) trait ErasedCovSnapshot {
     fn clear(&mut self);
     fn update_with(&mut self, other: &dyn ErasedCovSnapshot) -> bool;
+    fn mem_usage(&self) -> usize;
     fn as_any(&self) -> &dyn Any;
 }
 
 #[derive(libafl_bolts::SerdeAny)]
 pub(crate) enum CovSnapshot {
     Noop,
-    BitBox(bitvec::boxed::BitBox),
+    Bitset(roaring::RoaringBitmap),
     Erased(Box<dyn ErasedCovSnapshot>),
 }
 
 impl std::fmt::Debug for CovSnapshot {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Noop => write!(f, "Noop"),
-            Self::BitBox(arg0) => f.debug_tuple("BitBox").field(arg0).finish(),
+            Self::Bitset(arg0) => f.debug_tuple("Bitset").field(arg0).finish(),
             Self::Erased(_) => f.debug_tuple("Erased").finish_non_exhaustive(),
         }
     }
@@ -191,24 +192,43 @@ impl CovSnapshot {
     pub fn clear(&mut self) {
         match self {
             CovSnapshot::Noop => {}
-            CovSnapshot::BitBox(v) => v.fill(false),
+            CovSnapshot::Bitset(v) => v.clear(),
             CovSnapshot::Erased(erased) => erased.clear(),
         }
     }
 
     pub fn update_with(&mut self, other: &CovSnapshot) -> bool {
         match (self, other) {
             (CovSnapshot::Noop, CovSnapshot::Noop) => false,
-            (CovSnapshot::BitBox(s), CovSnapshot::BitBox(o)) => union_bitboxes(s, o),
+            (CovSnapshot::Bitset(s), CovSnapshot::Bitset(o)) => {
+                let res = !o.is_subset(s);
+                if res {
+                    *s |= o;
+                }
+                res
+            }
             (CovSnapshot::Erased(s), CovSnapshot::Erased(o)) => s.update_with(&**o),
             _ => unreachable!(),
         }
     }
+
+    pub fn mem_usage(&self) -> usize {
+        match self {
+            CovSnapshot::Noop => 0,
+            CovSnapshot::Bitset(bits) => {
+                let s = bits.statistics();
+                (s.n_bytes_array_containers
+                    + s.n_bytes_bitset_containers
+                    + s.n_bytes_run_containers) as usize
+            }
+            CovSnapshot::Erased(erased) => erased.mem_usage(),
+        }
+    }
 }
 
 #[derive(Debug, serde::Serialize, serde::Deserialize, libafl_bolts::SerdeAny)]
 pub(crate) struct InstrumentationSnapshot {
-    snapshots: Vec<CovSnapshot>,
+    pub snapshots: Vec<CovSnapshot>,
 }
 impl InstrumentationSnapshot {
     pub(crate) fn from(passes: &Passes) -> Self {
@@ -336,31 +356,89 @@ impl<K: Ord + Clone, V: Clone + FeedbackLattice> AssociatedCoverageArray<K, V> {
     where
         V: 'static,
     {
-        struct LatticeBox<V: FeedbackLattice + 'static>(Box<[V]>);
-        impl<V: FeedbackLattice + 'static> ErasedCovSnapshot for LatticeBox<V> {
+        enum SparseLatticeBox<V: FeedbackLattice + 'static> {
+            Full(Box<[V]>),
+            Sparse {
+                non_default_values: Box<[V]>,
+                non_default: roaring::RoaringBitmap,
+                orig_len: usize,
+            },
+        }
+        impl<V: FeedbackLattice> SparseLatticeBox<V> {
+            fn new(entries: &[V]) -> Self {
+                let non_default = roaring::RoaringBitmap::from_sorted_iter(
+                    entries
+                        .iter()
+                        .enumerate()
+                        .filter_map(|(i, v)| (!v.is_bottom()).then_some(i as u32)),
+                )
+                .unwrap();
+                let non_default_values =
+                    entries.iter().filter(|x| !x.is_bottom()).cloned().collect();
+                Self::Sparse {
+                    non_default_values,
+                    non_default,
+                    orig_len: entries.len(),
+                }
+            }
+            fn _update_with(&mut self, other: &Self) -> bool {
+                match self {
+                    SparseLatticeBox::Full(base) => match other {
+                        SparseLatticeBox::Full(_) => unimplemented!(),
+                        SparseLatticeBox::Sparse {
+                            non_default_values,
+                            non_default,
+                            ..
+                        } => {
+                            let mut res = false;
+                            for (i, ov) in non_default.iter().zip(non_default_values) {
+                                let sv = &mut base[i as usize];
+                                res |= sv.is_extended_by(ov);
+                                *sv = sv.unify(ov);
+                            }
+                            res
+                        }
+                    },
+                    SparseLatticeBox::Sparse { .. } => unimplemented!(),
+                }
+            }
+        }
+        impl<V: FeedbackLattice + 'static> ErasedCovSnapshot for SparseLatticeBox<V> {
             fn clear(&mut self) {
-                self.0.fill(V::bottom());
+                let len = match self {
+                    SparseLatticeBox::Full(x) => x.len(),
+                    SparseLatticeBox::Sparse { orig_len, .. } => *orig_len,
+                };
+                *self = Self::Full(vec![V::bottom(); len].into_boxed_slice());
             }
             fn update_with(&mut self, other: &dyn ErasedCovSnapshot) -> bool {
                 if let Some(other) = other.as_any().downcast_ref::<Self>() {
-                    other
-                        .0
-                        .iter()
-                        .zip(self.0.iter_mut())
-                        .fold(false, |mut res, (a, b)| {
-                            res |= b.is_extended_by(a);
-                            *b = b.unify(a);
-                            res
-                        })
+                    self._update_with(other)
                 } else {
                     false
                 }
             }
             fn as_any(&self) -> &dyn Any {
                 self as &dyn Any
             }
+            fn mem_usage(&self) -> usize {
+                match self {
+                    SparseLatticeBox::Full(x) => x.len() * std::mem::size_of::<V>(),
+                    SparseLatticeBox::Sparse {
+                        non_default_values,
+                        non_default,
+                        ..
+                    } => {
+                        let s = non_default.statistics();
+                        (s.n_bytes_array_containers
+                            + s.n_bytes_bitset_containers
+                            + s.n_bytes_run_containers) as usize
+                            + non_default_values.len() * std::mem::size_of::<V>()
+                    }
+                }
+            }
         }
-        CovSnapshot::Erased(Box::new(LatticeBox(self.entries.clone())))
+        CovSnapshot::Erased(Box::new(SparseLatticeBox::new(&self.entries)))
     }
 }
 

src/instrumentation/traits.rs

@@ -134,7 +134,12 @@ pub(crate) trait CodeCovInstrumentationPass {
         self.coverage_mut().reset_keep_saved()
     }
     fn snapshot_coverage(&self) -> CovSnapshot {
-        CovSnapshot::BitBox(self.coverage().entries.clone())
+        CovSnapshot::Bitset(
+            roaring::RoaringBitmap::from_sorted_iter(
+                self.coverage().entries.iter_ones().map(|x| x as u32),
+            )
+            .unwrap(),
+        )
     }
 
     fn as_any(&self) -> &dyn Any;
@@ -160,7 +165,12 @@ pub(crate) trait HashBitsetInstrumentationPass {
         self.coverage_mut().reset_keep_saved()
     }
     fn snapshot_coverage(&self) -> CovSnapshot {
-        CovSnapshot::BitBox(self.coverage().entries.clone())
+        CovSnapshot::Bitset(
+            roaring::RoaringBitmap::from_sorted_iter(
+                self.coverage().entries.iter_ones().map(|x| x as u32),
+            )
+            .unwrap(),
+        )
     }
 
     fn as_any(&self) -> &dyn Any;

src/jit/mod.rs

@@ -599,7 +599,7 @@ impl FeedbackOptions {
 impl std::fmt::Debug for FeedbackOptions {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let mut f = f.debug_struct("FeedbackOptions");
-        let v = serde_json::to_value(&self).unwrap();
+        let v = serde_json::to_value(self).unwrap();
         let mut skipped = false;
         for (k, v) in v.as_object().unwrap() {
             if v.as_bool().unwrap() {