pass scheduler to stage

libafl/src/events/llmp.rs

@@ -22,6 +22,7 @@ use crate::{
     executors::{Executor, HasObservers},
     inputs::Input,
     observers::ObserversTuple,
+    corpus::CorpusScheduler,
     state::IfInteresting,
     stats::Stats,
     Error,
@@ -252,14 +253,16 @@ where
     }
 
     // Handle arriving events in the client
-    fn handle_in_client<E, OT>(
+    fn handle_in_client<CS, E, OT>(
         &mut self,
         state: &mut S,
         _sender_id: u32,
         event: Event<I>,
         _executor: &mut E,
+        scheduler: &CS,
     ) -> Result<(), Error>
     where
+        CS: CorpusScheduler<I, S>,
         E: Executor<I> + HasObservers<OT>,
         OT: ObserversTuple,
     {
@@ -315,10 +318,11 @@ where
         }
     }
 
-    fn process<E, OT>(&mut self, state: &mut S, executor: &mut E) -> Result<usize, Error>
+    fn process<CS, E, OT>(&mut self, state: &mut S, executor: &mut E, scheduler: &CS) -> Result<usize, Error>
     where
+        CS: CorpusScheduler<I, S>,
         E: Executor<I> + HasObservers<OT>,
-        OT: ObserversTuple,
+        OT: ObserversTuple
     {
         // TODO: Get around local event copy by moving handle_in_client
         let mut events = vec![];
@@ -342,7 +346,7 @@ where
         };
         let count = events.len();
         events.drain(..).try_for_each(|(sender_id, event)| {
-            self.handle_in_client(state, sender_id, event, executor)
+            self.handle_in_client(state, sender_id, event, executor, scheduler)
         })?;
         Ok(count)
     }
@@ -426,12 +430,13 @@ where
             .send_buf(_LLMP_TAG_RESTART, &state_corpus_serialized)
     }
 
-    fn process<E, OT>(&mut self, state: &mut S, executor: &mut E) -> Result<usize, Error>
+    fn process<CS, E, OT>(&mut self, state: &mut S, executor: &mut E, scheduler: &CS) -> Result<usize, Error>
     where
+        CS: CorpusScheduler<I, S>,
         E: Executor<I> + HasObservers<OT>,
-        OT: ObserversTuple,
+        OT: ObserversTuple
     {
-        self.llmp_mgr.process(state, executor)
+        self.llmp_mgr.process(state, executor, scheduler)
     }
 
     fn fire(&mut self, state: &mut S, event: Event<I>) -> Result<(), Error> {

libafl/src/events/mod.rs

@@ -12,6 +12,7 @@ use serde::{Deserialize, Serialize};
 use crate::{
     executors::{Executor, HasObservers},
     inputs::Input,
+    corpus::{ CorpusScheduler},
     observers::ObserversTuple,
     Error,
 };
@@ -159,8 +160,9 @@ where
 
     /// Lookup for incoming events and process them.
     /// Return the number of processes events or an error
-    fn process<E, OT>(&mut self, state: &mut S, executor: &mut E) -> Result<usize, Error>
+    fn process<CS, E, OT>(&mut self, state: &mut S, executor: &mut E, scheduler: &CS) -> Result<usize, Error>
     where
+        CS: CorpusScheduler<I, S>,
         E: Executor<I> + HasObservers<OT>,
         OT: ObserversTuple;
 
@@ -203,8 +205,9 @@ impl<I, S> EventManager<I, S> for NopEventManager<I, S>
 where
     I: Input,
 {
-    fn process<E, OT>(&mut self, _state: &mut S, _executor: &mut E) -> Result<usize, Error>
+    fn process<CS, E, OT>(&mut self, _state: &mut S, _executor: &mut E, _scheduler: &CS) -> Result<usize, Error>
     where
+        CS: CorpusScheduler<I, S>,
         E: Executor<I> + HasObservers<OT>,
         OT: ObserversTuple,
     {

libafl/src/events/simple.rs

@@ -7,6 +7,7 @@ use crate::{
     executors::{Executor, HasObservers},
     inputs::Input,
     observers::ObserversTuple,
+    corpus::CorpusScheduler,
     stats::Stats,
     Error,
 };
@@ -30,10 +31,11 @@ where
     I: Input,
     ST: Stats, //CE: CustomEvent<I, OT>,
 {
-    fn process<E, OT>(&mut self, state: &mut S, _executor: &mut E) -> Result<usize, Error>
+    fn process<CS, E, OT>(&mut self, state: &mut S, _executor: &mut E, _scheduler: &CS) -> Result<usize, Error>
     where
+        CS: CorpusScheduler<I, S>,
         E: Executor<I> + HasObservers<OT>,
-        OT: ObserversTuple,
+        OT: ObserversTuple
     {
         let count = self.events.len();
         while self.events.len() > 0 {

libafl/src/fuzzer.rs

@@ -12,12 +12,13 @@ use crate::{
 use core::marker::PhantomData;
 
 /// Holds a set of stages
-pub trait HasStages<ST, E, EM, I, S>
+pub trait HasStages<CS, E, EM, I, S, ST>
 where
-    ST: StagesTuple<E, EM, I, S>,
+    ST: StagesTuple<CS, E, EM, I, S>,
     E: Executor<I>,
     EM: EventManager<I, S>,
     I: Input,
+    Self: Sized,
 {
     fn stages(&self) -> &ST;
 
@@ -47,7 +48,7 @@ pub trait Fuzzer<E, EM, S> {
 pub struct StdFuzzer<CS, ST, E, EM, I, OT, S>
 where
     CS: CorpusScheduler<I, S>,
-    ST: StagesTuple<E, EM, I, S>,
+    ST: StagesTuple<CS, E, EM, I, S>,
     E: Executor<I>,
     EM: EventManager<I, S>,
     I: Input,
@@ -57,10 +58,10 @@ where
     phantom: PhantomData<(E, EM, I, OT, S)>,
 }
 
-impl<CS, ST, E, EM, I, OT, S> HasStages<ST, E, EM, I, S> for StdFuzzer<CS, ST, E, EM, I, OT, S>
+impl<CS, ST, E, EM, I, OT, S> HasStages<CS, E, EM, I, S, ST> for StdFuzzer<CS, ST, E, EM, I, OT, S>
 where
     CS: CorpusScheduler<I, S>,
-    ST: StagesTuple<E, EM, I, S>,
+    ST: StagesTuple<CS, E, EM, I, S>,
     E: Executor<I>,
     EM: EventManager<I, S>,
     I: Input,
@@ -77,7 +78,7 @@ where
 impl<CS, ST, E, EM, I, OT, S> HasCorpusScheduler<CS, I, S> for StdFuzzer<CS, ST, E, EM, I, OT, S>
 where
     CS: CorpusScheduler<I, S>,
-    ST: StagesTuple<E, EM, I, S>,
+    ST: StagesTuple<CS, E, EM, I, S>,
     E: Executor<I>,
     EM: EventManager<I, S>,
     I: Input,
@@ -95,7 +96,7 @@ impl<CS, ST, E, EM, I, OT, S> Fuzzer<E, EM, S> for StdFuzzer<CS, ST, E, EM, I, O
 where
     CS: CorpusScheduler<I, S>,
     S: HasExecutions,
-    ST: StagesTuple<E, EM, I, S>,
+    ST: StagesTuple<CS, E, EM, I, S>,
     EM: EventManager<I, S>,
     E: Executor<I> + HasObservers<OT>,
     OT: ObserversTuple,
@@ -104,9 +105,9 @@ where
     fn fuzz_one(&self, state: &mut S, executor: &mut E, manager: &mut EM) -> Result<usize, Error> {
         let idx = self.scheduler().next(state)?;
 
-        self.stages().perform_all(state, executor, manager, idx)?;
+        self.stages().perform_all(state, executor, manager, self.scheduler(), idx)?;
 
-        manager.process(state, executor)?;
+        manager.process(state, executor, self.scheduler())?;
         Ok(idx)
     }
 
@@ -133,7 +134,7 @@ where
 impl<CS, ST, E, EM, I, OT, S> StdFuzzer<CS, ST, E, EM, I, OT, S>
 where
     CS: CorpusScheduler<I, S>,
-    ST: StagesTuple<E, EM, I, S>,
+    ST: StagesTuple<CS, E, EM, I, S>,
     E: Executor<I>,
     EM: EventManager<I, S>,
     I: Input,

libafl/src/stages/mod.rs

@@ -7,7 +7,7 @@ use crate::{
 
 /// A stage is one step in the fuzzing process.
 /// Multiple stages will be scheduled one by one for each input.
-pub trait Stage<E, EM, I, S>
+pub trait Stage<CS, E, EM,  I, S>
 where
     EM: EventManager<I, S>,
     E: Executor<I>,
@@ -19,11 +19,12 @@ where
         state: &mut S,
         executor: &mut E,
         manager: &mut EM,
+        scheduler: &CS,
         corpus_idx: usize,
     ) -> Result<(), Error>;
 }
 
-pub trait StagesTuple<E, EM, I, S>
+pub trait StagesTuple<CS, E, EM,  I, S>
 where
     EM: EventManager<I, S>,
     E: Executor<I>,
@@ -34,25 +35,26 @@ where
         state: &mut S,
         executor: &mut E,
         manager: &mut EM,
+        scheduler: &CS,
         corpus_idx: usize,
     ) -> Result<(), Error>;
 }
 
-impl<E, EM, I, S> StagesTuple<E, EM, I, S> for ()
+impl<CS, E, EM,  I, S> StagesTuple<CS, E, EM,  I, S> for ()
 where
     EM: EventManager<I, S>,
     E: Executor<I>,
     I: Input,
 {
-    fn perform_all(&self, _: &mut S, _: &mut E, _: &mut EM, _: usize) -> Result<(), Error> {
+    fn perform_all(&self, _: &mut S, _: &mut E, _: &mut EM, _: &CS, _: usize) -> Result<(), Error> {
         Ok(())
     }
 }
 
-impl<Head, Tail, E, EM, I, S> StagesTuple<E, EM, I, S> for (Head, Tail)
+impl<Head, Tail, CS, E, EM, I, S> StagesTuple<CS, E, EM,  I, S> for (Head, Tail)
 where
-    Head: Stage<E, EM, I, S>,
+    Head: Stage<CS, E, EM,  I, S>,
-    Tail: StagesTuple<E, EM, I, S> + TupleList,
+    Tail: StagesTuple<CS, E, EM,  I, S> + TupleList,
     EM: EventManager<I, S>,
     E: Executor<I>,
     I: Input,
@@ -62,9 +64,10 @@ where
         state: &mut S,
         executor: &mut E,
         manager: &mut EM,
+        scheduler: &CS,
         corpus_idx: usize,
     ) -> Result<(), Error> {
-        self.0.perform(state, executor, manager, corpus_idx)?;
+        self.0.perform(state, executor, manager, scheduler, corpus_idx)?;
-        self.1.perform_all(state, executor, manager, corpus_idx)
+        self.1.perform_all(state, executor, manager, scheduler, corpus_idx)
     }
 }

libafl/src/stages/mutational.rs

@@ -1,7 +1,7 @@
 use core::marker::PhantomData;
 
 use crate::{
-    corpus::Corpus,
+    corpus::{Corpus, CorpusScheduler},
     events::EventManager,
     executors::{Executor, HasObservers},
     inputs::Input,
@@ -18,7 +18,7 @@ use crate::{
 /// A Mutational stage is the stage in a fuzzing run that mutates inputs.
 /// Mutational stages will usually have a range of mutations that are
 /// being applied to the input one by one, between executions.
-pub trait MutationalStage<C, E, EM, I, M, OT, S>: Stage<E, EM, I, S>
+pub trait MutationalStage<C, CS, E, EM, I, M, OT, S>: Stage<CS, E, EM, I, S>
 where
     M: Mutator<I, S>,
     I: Input,
@@ -27,6 +27,7 @@ where
     EM: EventManager<I, S>,
     E: Executor<I> + HasObservers<OT>,
     OT: ObserversTuple,
+    CS: CorpusScheduler<I, S>
 {
     /// The mutator registered for this stage
     fn mutator(&self) -> &M;
@@ -43,6 +44,7 @@ where
         state: &mut S,
         executor: &mut E,
         manager: &mut EM,
+        scheduler: &CS,
         corpus_idx: usize,
     ) -> Result<(), Error> {
         let num = self.iterations(state);
@@ -67,7 +69,7 @@ pub static DEFAULT_MUTATIONAL_MAX_ITERATIONS: u64 = 128;
 
 /// The default mutational stage
 #[derive(Clone, Debug)]
-pub struct StdMutationalStage<C, E, EM, I, M, OT, R, S>
+pub struct StdMutationalStage<C, CS, E, EM, I, M, OT, R, S>
 where
     M: Mutator<I, S>,
     I: Input,
@@ -76,14 +78,15 @@ where
     EM: EventManager<I, S>,
     E: Executor<I> + HasObservers<OT>,
     OT: ObserversTuple,
+    CS: CorpusScheduler<I, S>,
     R: Rand,
 {
     mutator: M,
-    phantom: PhantomData<(C, E, EM, I, OT, R, S)>,
+    phantom: PhantomData<(C, CS, E, EM, I, OT, R, S)>,
 }
 
-impl<C, E, EM, I, M, OT, R, S> MutationalStage<C, E, EM, I, M, OT, S>
+impl<C, CS, E, EM, I, M, OT, R, S> MutationalStage<C, CS, E, EM, I, M, OT, S>
-    for StdMutationalStage<C, E, EM, I, M, OT, R, S>
+    for StdMutationalStage<C, CS, E, EM, I, M, OT, R, S>
 where
     M: Mutator<I, S>,
     I: Input,
@@ -92,6 +95,7 @@ where
     EM: EventManager<I, S>,
     E: Executor<I> + HasObservers<OT>,
     OT: ObserversTuple,
+    CS: CorpusScheduler<I, S>,
     R: Rand,
 {
     /// The mutator, added to this stage
@@ -112,7 +116,7 @@ where
     }
 }
 
-impl<C, E, EM, I, M, OT, R, S> Stage<E, EM, I, S> for StdMutationalStage<C, E, EM, I, M, OT, R, S>
+impl<C, CS, E, EM, I, M, OT, R, S> Stage<CS, E, EM, I, S> for StdMutationalStage<C, CS, E, EM, I, M, OT, R, S>
 where
     M: Mutator<I, S>,
     I: Input,
@@ -121,6 +125,7 @@ where
     EM: EventManager<I, S>,
     E: Executor<I> + HasObservers<OT>,
     OT: ObserversTuple,
+    CS: CorpusScheduler<I, S>,
     R: Rand,
 {
     #[inline]
@@ -129,13 +134,14 @@ where
         state: &mut S,
         executor: &mut E,
         manager: &mut EM,
+        scheduler: &CS,
         corpus_idx: usize,
     ) -> Result<(), Error> {
-        self.perform_mutational(state, executor, manager, corpus_idx)
+        self.perform_mutational(state, executor, manager, scheduler, corpus_idx)
     }
 }
 
-impl<C, E, EM, I, M, OT, R, S> StdMutationalStage<C, E, EM, I, M, OT, R, S>
+impl<C, CS, E, EM, I, M, OT, R, S> StdMutationalStage<C, CS, E, EM, I, M, OT, R, S>
 where
     M: Mutator<I, S>,
     I: Input,
@@ -144,6 +150,7 @@ where
     EM: EventManager<I, S>,
     E: Executor<I> + HasObservers<OT>,
     OT: ObserversTuple,
+    CS: CorpusScheduler<I, S>,
     R: Rand,
 {
     /// Creates a new default mutational stage

libafl/src/state/mod.rs

@@ -600,7 +600,7 @@ where
                 phantom: PhantomData,
             },
         )?;
-        manager.process(self, executor)?;
+        manager.process(self, executor, scheduler)?;
         Ok(())
     }
 }
@@ -680,7 +680,7 @@ where
                 phantom: PhantomData,
             },
         )?;
-        manager.process(self, executor)?;
+        manager.process(self, executor, scheduler)?;
         Ok(())
     }