From 58619beb7d6d0b09a8729215d6724e276071e06f Mon Sep 17 00:00:00 2001
From: Andrea Fioraldi <andreafioraldi@gmail.com>
Date: Fri, 5 Mar 2021 14:33:58 +0100
Subject: [PATCH] start working on power scheduling

---
 libafl/src/stages/mod.rs   |   5 +-
 libafl/src/stages/power.rs | 110 +++++++++++++++++++++++++++++++++++++
 2 files changed, 114 insertions(+), 1 deletion(-)
 create mode 100644 libafl/src/stages/power.rs

diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs
index e50a165d0b..6d91508a58 100644
--- a/libafl/src/stages/mod.rs
+++ b/libafl/src/stages/mod.rs
@@ -1,5 +1,8 @@
 pub mod mutational;
-pub use mutational::StdMutationalStage;
+pub use mutational::{MutationalStage, StdMutationalStage};
+
+//pub mod power;
+//pub use power::PowerMutationalStage;
 
 use crate::{
     bolts::tuples::TupleList, events::EventManager, executors::Executor, inputs::Input, Error,
diff --git a/libafl/src/stages/power.rs b/libafl/src/stages/power.rs
new file mode 100644
index 0000000000..3e333b52d7
--- /dev/null
+++ b/libafl/src/stages/power.rs
@@ -0,0 +1,110 @@
+use core::marker::PhantomData;
+
+use crate::{
+    corpus::{Corpus, CorpusScheduler},
+    events::EventManager,
+    executors::{Executor, HasObservers},
+    inputs::Input,
+    mutators::Mutator,
+    observers::ObserversTuple,
+    stages::{Stage, MutationalStage},
+    state::{Evaluator, HasCorpus, HasRand},
+    utils::Rand,
+    Error,
+};
+
+/// The mutational stage using power schedules
+#[derive(Clone, Debug)]
+pub struct PowerMutationalStage<C, CS, E, EM, I, M, OT, R, S>
+where
+    M: Mutator<I, S>,
+    I: Input,
+    S: HasCorpus<C, I> + Evaluator<I> + HasRand<R>,
+    C: Corpus<I>,
+    EM: EventManager<I, S>,
+    E: Executor<I> + HasObservers<OT>,
+    OT: ObserversTuple,
+    CS: CorpusScheduler<I, S>,
+    R: Rand,
+{
+    mutator: M,
+    phantom: PhantomData<(C, CS, E, EM, I, OT, R, S)>,
+}
+
+impl<C, CS, E, EM, I, M, OT, R, S> MutationalStage<C, CS, E, EM, I, M, OT, S>
+    for PowerMutationalStage<C, CS, E, EM, I, M, OT, R, S>
+where
+    M: Mutator<I, S>,
+    I: Input,
+    S: HasCorpus<C, I> + Evaluator<I> + HasRand<R>,
+    C: Corpus<I>,
+    EM: EventManager<I, S>,
+    E: Executor<I> + HasObservers<OT>,
+    OT: ObserversTuple,
+    CS: CorpusScheduler<I, S>,
+    R: Rand,
+{
+    /// The mutator, added to this stage
+    #[inline]
+    fn mutator(&self) -> &M {
+        &self.mutator
+    }
+
+    /// The list of mutators, added to this stage (as mutable ref)
+    #[inline]
+    fn mutator_mut(&mut self) -> &mut M {
+        &mut self.mutator
+    }
+
+    /// Gets the number of iterations as a random number
+    fn iterations(&self, state: &mut S) -> usize {
+        1 + state.rand_mut().below(DEFAULT_MUTATIONAL_MAX_ITERATIONS) as usize
+    }
+}
+
+impl<C, CS, E, EM, I, M, OT, R, S> Stage<CS, E, EM, I, S>
+    for PowerMutationalStage<C, CS, E, EM, I, M, OT, R, S>
+where
+    M: Mutator<I, S>,
+    I: Input,
+    S: HasCorpus<C, I> + Evaluator<I> + HasRand<R>,
+    C: Corpus<I>,
+    EM: EventManager<I, S>,
+    E: Executor<I> + HasObservers<OT>,
+    OT: ObserversTuple,
+    CS: CorpusScheduler<I, S>,
+    R: Rand,
+{
+    #[inline]
+    fn perform(
+        &self,
+        state: &mut S,
+        executor: &mut E,
+        manager: &mut EM,
+        scheduler: &CS,
+        corpus_idx: usize,
+    ) -> Result<(), Error> {
+        self.perform_mutational(state, executor, manager, scheduler, corpus_idx)
+    }
+}
+
+impl<C, CS, E, EM, I, M, OT, R, S> PowerMutationalStage<C, CS, E, EM, I, M, OT, R, S>
+where
+    M: Mutator<I, S>,
+    I: Input,
+    S: HasCorpus<C, I> + Evaluator<I> + HasRand<R>,
+    C: Corpus<I>,
+    EM: EventManager<I, S>,
+    E: Executor<I> + HasObservers<OT>,
+    OT: ObserversTuple,
+    CS: CorpusScheduler<I, S>,
+    R: Rand,
+{
+    /// Creates a new default mutational stage
+    pub fn new(mutator: M) -> Self {
+        Self {
+            mutator: mutator,
+            phantom: PhantomData,
+        }
+    }
+}