From cd8003ee69ba2e5c67e792d557420574f01c8272 Mon Sep 17 00:00:00 2001
From: Dan Blackwell <danblackwell95@yahoo.co.uk>
Date: Tue, 5 Mar 2024 18:04:13 +0000
Subject: [PATCH] Fix for #1881, ProbabilitySamplingScheduler seems to be
 inverted (#1884)

* Fix for #1881, ProbabilitySamplingScheduler seems to be inverted

* Run 'cargo +nightly fmt'

* Added debug_assert as negative and infinite probability values would break the current implementation of next()

---------

Co-authored-by: Dan Blackwell <danblackwell@eduroam-int-dhcp-97-192-218.ucl.ac.uk>
Co-authored-by: Dan Blackwell <danblackwell@ML-W0NYVW4XCK.local>
Co-authored-by: Dongjia "toka" Zhang <tokazerkje@outlook.com>
---
 .../src/schedulers/probabilistic_sampling.rs  | 55 ++++++++++++++++---
 1 file changed, 46 insertions(+), 9 deletions(-)

diff --git a/libafl/src/schedulers/probabilistic_sampling.rs b/libafl/src/schedulers/probabilistic_sampling.rs
index a5fe0cb644..921ba40344 100644
--- a/libafl/src/schedulers/probabilistic_sampling.rs
+++ b/libafl/src/schedulers/probabilistic_sampling.rs
@@ -9,9 +9,9 @@ use libafl_bolts::rands::Rand;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    corpus::{Corpus, CorpusId, HasTestcase},
+    corpus::{Corpus, CorpusId, HasTestcase, Testcase},
     inputs::UsesInput,
-    schedulers::{Scheduler, TestcaseScore},
+    schedulers::{RemovableScheduler, Scheduler, TestcaseScore},
     state::{HasCorpus, HasMetadata, HasRand, State, UsesState},
     Error,
 };
@@ -74,23 +74,60 @@ where
     #[allow(clippy::cast_precision_loss)]
     #[allow(clippy::unused_self)]
     pub fn store_probability(&self, state: &mut S, idx: CorpusId) -> Result<(), Error> {
-        let factor = F::compute(state, &mut *state.corpus().get(idx)?.borrow_mut())?;
-        if factor == 0.0 {
-            return Err(Error::illegal_state(
-                "Infinity probability calculated for probabilistic sampling scheduler",
-            ));
-        }
+        let prob = F::compute(state, &mut *state.corpus().get(idx)?.borrow_mut())?;
+        debug_assert!(
+            prob >= 0.0 && prob.is_finite(),
+            "scheduler probability is {prob}; to work correctly it must be >= 0.0 and finite"
+        );
         let meta = state
             .metadata_map_mut()
             .get_mut::<ProbabilityMetadata>()
             .unwrap();
-        let prob = 1.0 / factor;
         meta.map.insert(idx, prob);
         meta.total_probability += prob;
         Ok(())
     }
 }
 
+impl<F, S> RemovableScheduler for ProbabilitySamplingScheduler<F, S>
+where
+    F: TestcaseScore<S>,
+    S: HasCorpus + HasMetadata + HasRand + HasTestcase + State,
+{
+    fn on_remove(
+        &mut self,
+        state: &mut Self::State,
+        idx: CorpusId,
+        _testcase: &Option<Testcase<<Self::State as UsesInput>::Input>>,
+    ) -> Result<(), Error> {
+        let meta = state
+            .metadata_map_mut()
+            .get_mut::<ProbabilityMetadata>()
+            .unwrap();
+        if let Some(prob) = meta.map.remove(&idx) {
+            meta.total_probability -= prob;
+        }
+        Ok(())
+    }
+
+    fn on_replace(
+        &mut self,
+        state: &mut Self::State,
+        idx: CorpusId,
+        _prev: &Testcase<<Self::State as UsesInput>::Input>,
+    ) -> Result<(), Error> {
+        let meta = state
+            .metadata_map_mut()
+            .get_mut::<ProbabilityMetadata>()
+            .unwrap();
+        if let Some(prob) = meta.map.remove(&idx) {
+            meta.total_probability -= prob;
+        }
+
+        self.store_probability(state, idx)
+    }
+}
+
 impl<F, S> UsesState for ProbabilitySamplingScheduler<F, S>
 where
     S: State + HasTestcase,