Change when weighted scheduler make new table & don't update global metadata in on_remove/on_replace (#2011)

* fix

* backtick

libafl/src/schedulers/mod.rs

@@ -65,80 +65,6 @@ where
     }
 }
 
-/// Define the metadata operations when removing testcase from AFL-style scheduler
-pub trait HasAFLRemovableScheduler: RemovableScheduler
-where
-    Self::State: HasCorpus + HasMetadata + HasTestcase,
-{
-    #[allow(clippy::cast_precision_loss)]
-    #[allow(clippy::cast_precision_loss)]
-    /// Adjusting metadata when removing the testcase
-    fn on_remove_metadata(
-        &mut self,
-        state: &mut Self::State,
-        _idx: CorpusId,
-        prev: &Option<Testcase<<Self::State as UsesInput>::Input>>,
-    ) -> Result<(), Error> {
-        let prev = prev.as_ref().ok_or_else(|| {
-            Error::illegal_argument(
-                "Power schedulers must be aware of the removed corpus entry for reweighting.",
-            )
-        })?;
-
-        let prev_meta = prev.metadata::<SchedulerTestcaseMetadata>()?;
-
-        // Use these to adjust `SchedulerMetadata`
-        let (prev_total_time, prev_cycles) = prev_meta.cycle_and_time();
-        let prev_bitmap_size = prev_meta.bitmap_size();
-        let prev_bitmap_size_log = libm::log2(prev_bitmap_size as f64);
-
-        let psmeta = state.metadata_mut::<SchedulerMetadata>()?;
-
-        psmeta.set_exec_time(psmeta.exec_time() - prev_total_time);
-        psmeta.set_cycles(psmeta.cycles() - (prev_cycles as u64));
-        psmeta.set_bitmap_size(psmeta.bitmap_size() - prev_bitmap_size);
-        psmeta.set_bitmap_size_log(psmeta.bitmap_size_log() - prev_bitmap_size_log);
-        psmeta.set_bitmap_entries(psmeta.bitmap_entries() - 1);
-
-        Ok(())
-    }
-
-    #[allow(clippy::cast_precision_loss)]
-    /// Adjusting metadata when replacing the corpus
-    fn on_replace_metadata(
-        &mut self,
-        state: &mut Self::State,
-        idx: CorpusId,
-        prev: &Testcase<<Self::State as UsesInput>::Input>,
-    ) -> Result<(), Error> {
-        let prev_meta = prev.metadata::<SchedulerTestcaseMetadata>()?;
-
-        // Next depth is + 1
-        let prev_depth = prev_meta.depth() + 1;
-
-        // Use these to adjust `SchedulerMetadata`
-        let (prev_total_time, prev_cycles) = prev_meta.cycle_and_time();
-        let prev_bitmap_size = prev_meta.bitmap_size();
-        let prev_bitmap_size_log = libm::log2(prev_bitmap_size as f64);
-
-        let psmeta = state.metadata_mut::<SchedulerMetadata>()?;
-
-        // We won't add new one because it'll get added when it gets executed in calirbation next time.
-        psmeta.set_exec_time(psmeta.exec_time() - prev_total_time);
-        psmeta.set_cycles(psmeta.cycles() - (prev_cycles as u64));
-        psmeta.set_bitmap_size(psmeta.bitmap_size() - prev_bitmap_size);
-        psmeta.set_bitmap_size_log(psmeta.bitmap_size_log() - prev_bitmap_size_log);
-        psmeta.set_bitmap_entries(psmeta.bitmap_entries() - 1);
-
-        state
-            .corpus()
-            .get(idx)?
-            .borrow_mut()
-            .add_metadata(SchedulerTestcaseMetadata::new(prev_depth));
-        Ok(())
-    }
-}
-
 /// Defines the common metadata operations for the AFL-style schedulers
 pub trait HasAFLSchedulerMetadata<O, S>: Scheduler
 where

libafl/src/schedulers/powersched.rs

@@ -12,9 +12,7 @@ use crate::{
     corpus::{Corpus, CorpusId, HasTestcase, Testcase},
     inputs::UsesInput,
     observers::{MapObserver, ObserversTuple},
-    schedulers::{
+    schedulers::{HasAFLSchedulerMetadata, RemovableScheduler, Scheduler},
-        HasAFLRemovableScheduler, HasAFLSchedulerMetadata, RemovableScheduler, Scheduler,
-    },
     state::{HasCorpus, HasMetadata, State, UsesState},
     Error,
 };
@@ -186,34 +184,29 @@ where
     type State = S;
 }
 
-impl<O, S> HasAFLRemovableScheduler for PowerQueueScheduler<O, S>
+impl<O, S> RemovableScheduler for PowerQueueScheduler<O, S>
 where
     S: State + HasTestcase + HasMetadata + HasCorpus,
     O: MapObserver,
 {
-}
+    /// This will *NOT* neutralize the effect of this removed testcase from the global data such as `SchedulerMetadata`
-
-impl<O, S> RemovableScheduler for PowerQueueScheduler<O, S>
-where
-    S: HasCorpus + HasMetadata + HasTestcase + State,
-    O: MapObserver,
-{
     fn on_remove(
         &mut self,
-        state: &mut Self::State,
+        _state: &mut Self::State,
-        idx: CorpusId,
+        _idx: CorpusId,
-        prev: &Option<Testcase<<Self::State as UsesInput>::Input>>,
+        _prev: &Option<Testcase<<Self::State as UsesInput>::Input>>,
     ) -> Result<(), Error> {
-        self.on_remove_metadata(state, idx, prev)
+        Ok(())
     }
 
+    /// This will *NOT* neutralize the effect of this removed testcase from the global data such as `SchedulerMetadata`
     fn on_replace(
         &mut self,
-        state: &mut Self::State,
+        _state: &mut Self::State,
-        idx: CorpusId,
+        _idx: CorpusId,
-        prev: &Testcase<<Self::State as UsesInput>::Input>,
+        _prev: &Testcase<<Self::State as UsesInput>::Input>,
     ) -> Result<(), Error> {
-        self.on_replace_metadata(state, idx, prev)
+        Ok(())
     }
 }
 

libafl/src/schedulers/weighted.rs

@@ -16,7 +16,7 @@ use crate::{
     schedulers::{
         powersched::{PowerSchedule, SchedulerMetadata},
         testcase_score::{CorpusWeightTestcaseScore, TestcaseScore},
-        HasAFLRemovableScheduler, HasAFLSchedulerMetadata, RemovableScheduler, Scheduler,
+        HasAFLSchedulerMetadata, RemovableScheduler, Scheduler,
     },
     state::{HasCorpus, HasMetadata, HasRand, State, UsesState},
     Error,
@@ -93,6 +93,7 @@ libafl_bolts::impl_serdeany!(WeightedScheduleMetadata);
 /// A corpus scheduler using power schedules with weighted queue item selection algo.
 #[derive(Clone, Debug)]
 pub struct WeightedScheduler<F, O, S> {
+    table_invalidated: bool,
     strat: Option<PowerSchedule>,
     map_observer_name: String,
     last_hash: usize,
@@ -121,6 +122,7 @@ where
             strat,
             map_observer_name: map_observer.name().to_string(),
             last_hash: 0,
+            table_invalidated: true,
             phantom: PhantomData,
         }
     }
@@ -223,36 +225,32 @@ where
     type State = S;
 }
 
-impl<F, O, S> HasAFLRemovableScheduler for WeightedScheduler<F, O, S>
-where
-    F: TestcaseScore<S>,
-    S: State + HasTestcase + HasMetadata + HasCorpus + HasRand,
-    O: MapObserver,
-{
-}
-
 impl<F, O, S> RemovableScheduler for WeightedScheduler<F, O, S>
 where
     F: TestcaseScore<S>,
     O: MapObserver,
     S: HasCorpus + HasMetadata + HasRand + HasTestcase + State,
 {
+    /// This will *NOT* neutralize the effect of this removed testcase from the global data such as `SchedulerMetadata`
     fn on_remove(
         &mut self,
-        state: &mut Self::State,
+        _state: &mut Self::State,
-        idx: CorpusId,
+        _idx: CorpusId,
-        prev: &Option<Testcase<<Self::State as UsesInput>::Input>>,
+        _prev: &Option<Testcase<<Self::State as UsesInput>::Input>>,
     ) -> Result<(), Error> {
-        self.on_remove_metadata(state, idx, prev)
+        self.table_invalidated = true;
+        Ok(())
     }
 
+    /// This will *NOT* neutralize the effect of this removed testcase from the global data such as `SchedulerMetadata`
     fn on_replace(
         &mut self,
-        state: &mut Self::State,
+        _state: &mut Self::State,
-        idx: CorpusId,
+        _idx: CorpusId,
-        prev: &Testcase<<Self::State as UsesInput>::Input>,
+        _prev: &Testcase<<Self::State as UsesInput>::Input>,
     ) -> Result<(), Error> {
-        self.on_replace_metadata(state, idx, prev)
+        self.table_invalidated = true;
+        Ok(())
     }
 }
 
@@ -284,7 +282,8 @@ where
     /// Called when a [`Testcase`] is added to the corpus
     fn on_add(&mut self, state: &mut S, idx: CorpusId) -> Result<(), Error> {
         self.on_add_metadata(state, idx)?;
-        self.create_alias_table(state)
+        self.table_invalidated = true;
+        Ok(())
     }
 
     fn on_evaluation<OT>(
@@ -301,6 +300,10 @@ where
 
     #[allow(clippy::similar_names, clippy::cast_precision_loss)]
     fn next(&mut self, state: &mut S) -> Result<CorpusId, Error> {
+        if self.table_invalidated {
+            self.create_alias_table(state)?;
+            self.table_invalidated = false;
+        }
         let corpus_counts = state.corpus().count();
         if corpus_counts == 0 {
             Err(Error::empty(String::from(