MOpt scheduler (#161)

* add the struct for MOpt globals

* constants

* RAND_C

* more comments & reorder class members

* select_algorithm

* no_std fixes

* clippy fixes

* MOptMutator

* MutatorsTuple has HasLen

* MOptStage

* pso_update

* HasMOpt trait

* ScheduledMutator, core_fuzzing

* clippy fix

* fmt

* core_fuzzing

* core_fuzzing done

* fix

* pilot_mutate

* pilot_fuzzing

* pilot_fuzzing done

* MOpt metadata

* Make MOptMutator into a trait

* initialize_mopt

* No getter/setters

* fmt

* fixed compiler warnings & clippy warnings

* Comments

* fix type paramter, integrate into libpng

* fmt

* fmt

* No HasMOpt

* fmt

* improve

* pso_initialize, various fixes

* clippy

* fmt

* always pacemaker mode

* fmt

* fix

* less noisy fmt::Debug

Co-authored-by: Dominik Maier <domenukk@gmail.com>
Co-authored-by: Andrea Fioraldi <andreafioraldi@gmail.com>

fuzzers/libfuzzer_libpng/src/lib.rs

@@ -17,10 +17,11 @@ use libafl::{
     feedbacks::{CrashFeedback, MapFeedbackState, MaxMapFeedback, TimeFeedback, TimeoutFeedback},
     fuzzer::{Fuzzer, StdFuzzer},
     inputs::{BytesInput, HasTargetBytes},
-    mutators::scheduled::{havoc_mutations, tokens_mutations, StdScheduledMutator},
+    mutators::mopt_mutator::StdMOptMutator,
+    mutators::scheduled::{havoc_mutations, tokens_mutations},
     mutators::token_mutations::Tokens,
     observers::{HitcountsMapObserver, StdMapObserver, TimeObserver},
-    stages::mutational::StdMutationalStage,
+    stages::mopt::MOptStage,
     state::{HasCorpus, HasMetadata, StdState},
     stats::MultiStats,
     Error,
@@ -119,8 +120,14 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
     }
 
     // Setup a basic mutator with a mutational stage
+
+    let mutator = StdMOptMutator::new(havoc_mutations().merge(tokens_mutations()));
+    let mut stages = tuple_list!(MOptStage::new(mutator, &mut state, 5)?);
+
+    /*
     let mutator = StdScheduledMutator::new(havoc_mutations().merge(tokens_mutations()));
     let mut stages = tuple_list!(StdMutationalStage::new(mutator));
+    */
 
     // A minimization+queue policy to get testcasess from the corpus
     let scheduler = IndexesLenTimeMinimizerCorpusScheduler::new(QueueCorpusScheduler::new());

libafl/src/lib.rs

@@ -72,6 +72,8 @@ pub enum Error {
     IllegalArgument(String),
     /// Forkserver related Error
     Forkserver(String),
+    /// MOpt related Error
+    MOpt(String),
     /// Shutting down, not really an error.
     ShuttingDown,
     /// Something else happened
@@ -96,6 +98,7 @@ impl fmt::Display for Error {
             Self::IllegalState(s) => write!(f, "Illegal state: {0}", &s),
             Self::IllegalArgument(s) => write!(f, "Illegal argument: {0}", &s),
             Self::Forkserver(s) => write!(f, "Forkserver : {0}", &s),
+            Self::MOpt(s) => write!(f, "MOpt: {0}", &s),
             Self::ShuttingDown => write!(f, "Shutting down!"),
             Self::Unknown(s) => write!(f, "Unknown error: {0}", &s),
         }

libafl/src/mutators/mod.rs

@@ -6,6 +6,8 @@ pub mod mutations;
 pub use mutations::*;
 pub mod token_mutations;
 pub use token_mutations::*;
+pub mod mopt_mutator;
+pub use mopt_mutator::*;
 
 use crate::{
     bolts::tuples::{HasLen, Named},

libafl/src/mutators/mopt_mutator.rs

@@ -0,0 +1,620 @@
+//! The `MOpt` mutator scheduler, see <https://github.com/puppet-meteor/MOpt-AFL> and <https://www.usenix.org/conference/usenixsecurity19/presentation/lyu>
+use alloc::{string::ToString, vec::Vec};
+
+use crate::{
+    bolts::{rands::Rand, rands::StdRand},
+    inputs::Input,
+    mutators::{ComposedByMutations, MutationResult, Mutator, MutatorsTuple, ScheduledMutator},
+    state::{HasMetadata, HasRand},
+    Error,
+};
+use core::{
+    fmt::{self, Debug},
+    marker::PhantomData,
+};
+use serde::{Deserialize, Serialize};
+
+/// A Struct for managing MOpt-mutator parameters
+/// There are 2 modes for `MOpt` scheduler, the core fuzzing mode and the pilot fuzzing mode
+/// In short, in the pilot fuzzing mode, the fuzzer employs several `swarms` to compute the probability to choose the mutation operator
+/// On the other hand, in the core fuzzing mode, the fuzzer chooses the best `swarms`, which was determined during the pilot fuzzing mode, to compute the probability to choose the operation operator
+/// With the current implementation we are always in the pacemaker fuzzing mode.
+#[derive(Serialize, Deserialize, Clone)]
+pub struct MOpt {
+    /// Random number generator
+    pub rand: StdRand,
+    /// The number of total findings (unique crashes and unique interesting paths). This is equivalent to `state.corpus().count() + state.solutions().count()`;
+    pub total_finds: usize,
+    /// The number of finds before switching to this mode.
+    pub finds_before_switch: usize,
+    /// The MOpt mode that we are currently using the pilot fuzzing mode or the core_fuzzing mode
+    pub key_module: MOptMode,
+    /// These w_* and g_* values are the coefficients for updating variables according to the PSO algorithms
+    pub w_init: f64,
+    pub w_end: f64,
+    pub w_now: f64,
+    pub g_now: i32,
+    pub g_max: i32,
+    /// The number of mutation operators
+    pub operator_num: usize,
+    /// The number of swarms that we want to employ during the pilot fuzzing mode
+    pub swarm_num: usize,
+    /// We'll generate testcases for `period_pilot` times before we call pso_update in core fuzzing module
+    pub period_pilot: usize,
+    /// We'll generate testcases for `period_core` times before we call pso_update in core fuzzing module
+    pub period_core: usize,
+    /// The number of testcases generated during this pilot fuzzing mode
+    pub pilot_time: usize,
+    /// The number of testcases generated during this core fuzzing mode
+    pub core_time: usize,
+    /// The swarm identifier that we are currently using in the pilot fuzzing mode
+    pub swarm_now: usize,
+    /// These are the parameters for the PSO algorithm
+    x_now: Vec<Vec<f64>>,
+    l_best: Vec<Vec<f64>>,
+    eff_best: Vec<Vec<f64>>,
+    g_best: Vec<f64>,
+    v_now: Vec<Vec<f64>>,
+    /// The probability that we want to use to choose the mutation operator.
+    probability_now: Vec<Vec<f64>>,
+    /// The fitness for each swarm, we'll calculate the fitness in the pilot fuzzing mode and use the best one in the core fuzzing mode
+    pub swarm_fitness: Vec<f64>,
+    /// (Pilot Mode) Finds by each operators. This vector is used in pso_update
+    pub pilot_operator_finds_pso: Vec<Vec<usize>>,
+    /// (Pilot Mode) Finds by each operator till now.
+    pub pilot_operator_finds_this: Vec<Vec<usize>>,
+    /// (Pilot Mode) The number of mutation operator used. This vector is used in pso_update
+    pub pilot_operator_ctr_pso: Vec<Vec<usize>>,
+    /// (Pilot Mode) The number of mutation operator used till now
+    pub pilot_operator_ctr_this: Vec<Vec<usize>>,
+    /// (Pilot Mode) The number of mutation operator used till last execution
+    pub pilot_operator_ctr_last: Vec<Vec<usize>>,
+    /// Vector used in pso_update
+    pub operator_finds_puppet: Vec<usize>,
+    /// (Core Mode) Finds by each operators. This vector is used in pso_update
+    pub core_operator_finds_pso: Vec<usize>,
+    /// (Core Mode) Finds by each operator till now.
+    pub core_operator_finds_this: Vec<usize>,
+    /// (Core Mode) The number of mutation operator used. This vector is used in pso_update
+    pub core_operator_ctr_pso: Vec<usize>,
+    /// (Core Mode) The number of mutation operator used till now
+    pub core_operator_ctr_this: Vec<usize>,
+    /// (Core Mode) The number of mutation operator used till last execution
+    pub core_operator_ctr_last: Vec<usize>,
+}
+
+crate::impl_serdeany!(MOpt);
+
+impl fmt::Debug for MOpt {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("MOpt")
+            .field("\ntotal_finds", &self.total_finds)
+            .field("\nfinds_before_switch", &self.finds_before_switch)
+            .field("\nkey_module", &self.key_module)
+            .field("\nw_init", &self.w_init)
+            .field("\nw_end", &self.w_end)
+            .field("\nw_now", &self.g_now)
+            .field("\ng_now", &self.g_max)
+            .field("\npilot_time", &self.pilot_time)
+            .field("\ncore_time", &self.core_time)
+            .field("\n\nx_now", &self.x_now)
+            .field("\n\nl_best", &self.l_best)
+            .field("\n\neff_best", &self.eff_best)
+            .field("\n\ng_best", &self.g_best)
+            .field("\n\nv_now", &self.v_now)
+            .field("\n\nprobability_now", &self.probability_now)
+            .field("\n\nswarm_fitness", &self.swarm_fitness)
+            .field(
+                "\n\npilot_operator_finds_pso",
+                &self.pilot_operator_finds_pso,
+            )
+            .field(
+                "\n\npilot_operator_finds_this",
+                &self.pilot_operator_finds_this,
+            )
+            .field("\n\npilot_operator_ctr_pso", &self.pilot_operator_ctr_pso)
+            .field("\n\npilot_operator_ctr_this", &self.pilot_operator_ctr_this)
+            .field("\n\npilot_operator_ctr_last", &self.pilot_operator_ctr_last)
+            .field("\n\noperator_finds_puppuet", &self.operator_finds_puppet)
+            .field("\n\ncore_operator_finds_pso", &self.core_operator_finds_pso)
+            .field(
+                "\n\ncore_operator_finds_this",
+                &self.core_operator_finds_this,
+            )
+            .field("\n\ncore_operator_ctr_pso", &self.core_operator_ctr_pso)
+            .field("\n\ncore_operator_ctr_this", &self.core_operator_ctr_this)
+            .field("\n\ncore_operator_ctr_last", &self.core_operator_ctr_last)
+            .finish()
+    }
+}
+
+impl MOpt {
+    pub fn new(operator_num: usize, swarm_num: usize) -> Result<Self, Error> {
+        let mut mopt = Self {
+            rand: StdRand::with_seed(0),
+            total_finds: 0,
+            finds_before_switch: 0,
+            key_module: MOptMode::Pilotfuzzing,
+            w_init: 0.9,
+            w_end: 0.3,
+            w_now: 0.0,
+            g_now: 0,
+            g_max: 5000,
+            operator_num,
+            swarm_num,
+            period_pilot: 50000,
+            period_core: 500000,
+            pilot_time: 0,
+            core_time: 0,
+            swarm_now: 0,
+            x_now: vec![vec![0.0; operator_num]; swarm_num],
+            l_best: vec![vec![0.0; operator_num]; swarm_num],
+            eff_best: vec![vec![0.0; operator_num]; swarm_num],
+            g_best: vec![0.0; operator_num],
+            v_now: vec![vec![0.0; operator_num]; swarm_num],
+            probability_now: vec![vec![0.0; operator_num]; swarm_num],
+            swarm_fitness: vec![0.0; swarm_num],
+            pilot_operator_finds_pso: vec![vec![0; operator_num]; swarm_num],
+            pilot_operator_finds_this: vec![vec![0; operator_num]; swarm_num],
+            pilot_operator_ctr_pso: vec![vec![0; operator_num]; swarm_num],
+            pilot_operator_ctr_this: vec![vec![0; operator_num]; swarm_num],
+            pilot_operator_ctr_last: vec![vec![0; operator_num]; swarm_num],
+            operator_finds_puppet: vec![0; operator_num],
+            core_operator_finds_pso: vec![0; operator_num],
+            core_operator_finds_this: vec![0; operator_num],
+            core_operator_ctr_pso: vec![0; operator_num],
+            core_operator_ctr_this: vec![0; operator_num],
+            core_operator_ctr_last: vec![0; operator_num],
+        };
+        mopt.pso_initialize()?;
+        Ok(mopt)
+    }
+
+    /// Initialize `core_operator_*` values
+    pub fn init_core_module(&mut self) -> Result<(), Error> {
+        for i in 0..self.operator_num {
+            self.core_operator_ctr_this[i] = self.core_operator_ctr_pso[i];
+            self.core_operator_ctr_last[i] = self.core_operator_ctr_pso[i];
+            self.core_operator_finds_this[i] = self.core_operator_finds_pso[i]
+        }
+
+        let mut swarm_eff = 0.0;
+        let mut best_swarm = 0;
+        for i in 0..self.swarm_num {
+            if self.swarm_fitness[i] > swarm_eff {
+                swarm_eff = self.swarm_fitness[i];
+                best_swarm = i;
+            }
+        }
+
+        self.swarm_now = best_swarm;
+        Ok(())
+    }
+
+    #[inline]
+    pub fn update_pilot_operator_ctr_last(&mut self, swarm_now: usize) {
+        for i in 0..self.operator_num {
+            self.pilot_operator_ctr_last[swarm_now][i] = self.pilot_operator_ctr_this[swarm_now][i]
+        }
+    }
+
+    #[inline]
+    pub fn update_core_operator_ctr_last(&mut self) {
+        for i in 0..self.operator_num {
+            self.core_operator_ctr_last[i] = self.core_operator_ctr_this[i];
+        }
+    }
+
+    /// Finds the local optimum for each operator
+    /// See <https://github.com/puppet-meteor/MOpt-AFL/blob/master/MOpt/afl-fuzz.c#L8709>
+
+    #[allow(clippy::cast_precision_loss)]
+    pub fn update_pilot_operator_ctr_pso(&mut self, swarm_now: usize) {
+        let mut eff = 0.0;
+        for i in 0..self.operator_num {
+            if self.pilot_operator_ctr_this[swarm_now][i]
+                > self.pilot_operator_ctr_pso[swarm_now][i]
+            {
+                eff = ((self.pilot_operator_finds_this[swarm_now][i]
+                    - self.pilot_operator_finds_pso[swarm_now][i]) as f64)
+                    / ((self.pilot_operator_ctr_this[swarm_now][i]
+                        - self.pilot_operator_ctr_pso[swarm_now][i]) as f64)
+            }
+
+            if self.eff_best[swarm_now][i] < eff {
+                self.eff_best[swarm_now][i] = eff;
+                self.l_best[swarm_now][i] = self.x_now[swarm_now][i];
+            }
+
+            self.pilot_operator_finds_pso[swarm_now][i] =
+                self.pilot_operator_finds_this[swarm_now][i];
+            self.pilot_operator_ctr_pso[swarm_now][i] = self.pilot_operator_ctr_this[swarm_now][i];
+        }
+    }
+
+    #[inline]
+    pub fn update_core_operator_ctr_pso(&mut self) {
+        for i in 0..self.operator_num {
+            self.core_operator_finds_pso[i] = self.core_operator_finds_this[i];
+            self.core_operator_ctr_pso[i] = self.core_operator_ctr_this[i];
+        }
+    }
+
+    #[allow(clippy::cast_precision_loss)]
+    pub fn pso_initialize(&mut self) -> Result<(), Error> {
+        if self.g_now > self.g_max {
+            self.g_now = 0;
+        }
+        self.w_now = (self.w_init - self.w_end) * f64::from(self.g_max - self.g_now)
+            / f64::from(self.g_max)
+            + self.w_end;
+
+        for swarm in 0..self.swarm_num {
+            let mut total_x_now = 0.0;
+            let mut x_sum = 0.0;
+            for i in 0..self.operator_num {
+                self.x_now[swarm][i] = (self.rand.below(7000) as f64) * 0.0001 + 0.1;
+                total_x_now += self.x_now[swarm][i];
+                self.v_now[swarm][i] = 0.1;
+                self.l_best[swarm][i] = 0.5;
+                self.g_best[i] = 0.5;
+            }
+
+            for i in 0..self.operator_num {
+                self.x_now[swarm][i] /= total_x_now
+            }
+
+            for i in 0..self.operator_num {
+                self.v_now[swarm][i] = self.w_now * self.v_now[swarm][i]
+                    + (self.rand.below(1000) as f64)
+                        * 0.001
+                        * (self.l_best[swarm][i] - self.x_now[swarm][i])
+                    + (self.rand.below(1000) as f64)
+                        * 0.001
+                        * (self.g_best[i] - self.x_now[swarm][i]);
+                self.x_now[swarm][i] += self.v_now[swarm][i];
+
+                if self.x_now[swarm][i] > V_MAX {
+                    self.x_now[swarm][i] = V_MAX;
+                } else if self.x_now[swarm][i] < V_MIN {
+                    self.x_now[swarm][i] = V_MIN;
+                }
+
+                x_sum += self.x_now[swarm][i]
+            }
+
+            for i in 0..self.operator_num {
+                self.x_now[swarm][i] /= x_sum;
+                if i == 0 {
+                    self.probability_now[swarm][i] = self.x_now[swarm][i];
+                } else {
+                    self.probability_now[swarm][i] =
+                        self.probability_now[swarm][i - 1] + self.x_now[swarm][i];
+                }
+            }
+            if self.probability_now[swarm][self.operator_num - 1] < 0.99
+                || self.probability_now[swarm][self.operator_num - 1] > 1.01
+            {
+                return Err(Error::MOpt("Error in pso_update".to_string()));
+            }
+        }
+        Ok(())
+    }
+
+    /// Update the PSO algorithm parameters
+    /// See <https://github.com/puppet-meteor/MOpt-AFL/blob/master/MOpt/afl-fuzz.c#L10623>
+    #[allow(clippy::cast_precision_loss)]
+    pub fn pso_update(&mut self) -> Result<(), Error> {
+        self.g_now += 1;
+        if self.g_now > self.g_max {
+            self.g_now = 0;
+        }
+        self.w_now = (self.w_init - self.w_end) * f64::from(self.g_max - self.g_now)
+            / f64::from(self.g_max)
+            + self.w_end;
+
+        let mut operator_find_sum = 0;
+
+        for i in 0..self.operator_num {
+            self.operator_finds_puppet[i] = self.core_operator_ctr_pso[i];
+
+            for j in 0..self.swarm_num {
+                self.operator_finds_puppet[i] += self.pilot_operator_finds_pso[j][i];
+            }
+            operator_find_sum += self.operator_finds_puppet[i];
+        }
+
+        for i in 0..self.operator_num {
+            if self.operator_finds_puppet[i] > 0 {
+                self.g_best[i] =
+                    (self.operator_finds_puppet[i] as f64) / (operator_find_sum as f64);
+            }
+        }
+
+        for swarm in 0..self.swarm_num {
+            let mut x_sum = 0.0;
+            for i in 0..self.operator_num {
+                self.probability_now[swarm][i] = 0.0;
+                self.v_now[swarm][i] = self.w_now * self.v_now[swarm][i]
+                    + (self.rand.below(1000) as f64)
+                        * 0.001
+                        * (self.l_best[swarm][i] - self.x_now[swarm][i])
+                    + (self.rand.below(1000) as f64)
+                        * 0.001
+                        * (self.g_best[i] - self.x_now[swarm][i]);
+                self.x_now[swarm][i] += self.v_now[swarm][i];
+
+                if self.x_now[swarm][i] > V_MAX {
+                    self.x_now[swarm][i] = V_MAX;
+                } else if self.x_now[swarm][i] < V_MIN {
+                    self.x_now[swarm][i] = V_MIN;
+                }
+                x_sum += self.x_now[swarm][i];
+            }
+
+            for i in 0..self.operator_num {
+                self.x_now[swarm][i] /= x_sum;
+                if i == 0 {
+                    self.probability_now[swarm][i] = self.x_now[swarm][i];
+                } else {
+                    self.probability_now[swarm][i] =
+                        self.probability_now[swarm][i - 1] + self.x_now[swarm][i];
+                }
+            }
+            if self.probability_now[swarm][self.operator_num - 1] < 0.99
+                || self.probability_now[swarm][self.operator_num - 1] > 1.01
+            {
+                return Err(Error::MOpt("Error in pso_update".to_string()));
+            }
+        }
+        self.swarm_now = 0;
+
+        self.key_module = MOptMode::Pilotfuzzing;
+        //println!("Mopt struct:\n{:?}", self);
+        Ok(())
+    }
+
+    /// This function is used to decide the operator that we want to apply next
+    /// see <https://github.com/puppet-meteor/MOpt-AFL/blob/master/MOpt/afl-fuzz.c#L397>
+    #[allow(clippy::cast_precision_loss)]
+    pub fn select_algorithm(&mut self) -> Result<usize, Error> {
+        let mut res = 0;
+        let mut sentry = 0;
+
+        let operator_num = self.operator_num;
+
+        // Fetch a random sele value
+        let select_prob: f64 = self.probability_now[self.swarm_now][operator_num - 1]
+            * ((self.rand.below(10000) as f64) * 0.0001);
+
+        for i in 0..operator_num {
+            if i == 0 {
+                if select_prob < self.probability_now[self.swarm_now][i] {
+                    res = i;
+                    break;
+                }
+            } else if select_prob < self.probability_now[self.swarm_now][i] {
+                res = i;
+                sentry = 1;
+                break;
+            }
+        }
+
+        if (sentry == 1 && select_prob < self.probability_now[self.swarm_now][res - 1])
+            || (res + 1 < operator_num
+                && select_prob > self.probability_now[self.swarm_now][res + 1])
+        {
+            return Err(Error::MOpt("Error in select_algorithm".to_string()));
+        }
+        Ok(res)
+    }
+}
+
+const V_MAX: f64 = 1.0;
+const V_MIN: f64 = 0.05;
+
+#[derive(Serialize, Deserialize, Clone, Copy, Debug)]
+pub enum MOptMode {
+    Pilotfuzzing,
+    Corefuzzing,
+}
+
+pub struct StdMOptMutator<I, MT, R, S>
+where
+    I: Input,
+    MT: MutatorsTuple<I, S>,
+    R: Rand,
+    S: HasRand<R> + HasMetadata,
+{
+    mutations: MT,
+    phantom: PhantomData<(I, R, S)>,
+}
+
+impl<I, MT, R, S> Debug for StdMOptMutator<I, MT, R, S>
+where
+    I: Input,
+    MT: MutatorsTuple<I, S>,
+    R: Rand,
+    S: HasRand<R> + HasMetadata,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "StdMOptMutator with {} mutations for Input type {}",
+            self.mutations.len(),
+            core::any::type_name::<I>()
+        )
+    }
+}
+
+impl<I, MT, R, S> Mutator<I, S> for StdMOptMutator<I, MT, R, S>
+where
+    I: Input,
+    MT: MutatorsTuple<I, S>,
+    R: Rand,
+    S: HasRand<R> + HasMetadata,
+{
+    #[inline]
+    fn mutate(
+        &mut self,
+        state: &mut S,
+        input: &mut I,
+        stage_idx: i32,
+    ) -> Result<MutationResult, Error> {
+        self.scheduled_mutate(state, input, stage_idx)
+    }
+}
+
+impl<I, MT, R, S> StdMOptMutator<I, MT, R, S>
+where
+    I: Input,
+    MT: MutatorsTuple<I, S>,
+    R: Rand,
+    S: HasRand<R> + HasMetadata,
+{
+    pub fn new(mutations: MT) -> Self {
+        Self {
+            mutations,
+            phantom: PhantomData,
+        }
+    }
+    fn core_mutate(
+        &mut self,
+        state: &mut S,
+        input: &mut I,
+        stage_idx: i32,
+    ) -> Result<MutationResult, Error> {
+        // TODO
+        let mut r = MutationResult::Skipped;
+        state
+            .metadata_mut()
+            .get_mut::<MOpt>()
+            .unwrap()
+            .update_core_operator_ctr_last();
+
+        for _i in 0..self.iterations(state, input) {
+            let idx = self.schedule(state, input);
+            let outcome = self
+                .mutations_mut()
+                .get_and_mutate(idx, state, input, stage_idx)?;
+            if outcome == MutationResult::Mutated {
+                r = MutationResult::Mutated;
+            }
+
+            state
+                .metadata_mut()
+                .get_mut::<MOpt>()
+                .unwrap()
+                .core_operator_ctr_this[idx] += 1;
+        }
+
+        Ok(r)
+    }
+
+    fn pilot_mutate(
+        &mut self,
+        state: &mut S,
+        input: &mut I,
+        stage_idx: i32,
+    ) -> Result<MutationResult, Error> {
+        let mut r = MutationResult::Skipped;
+        let swarm_now;
+        {
+            let mopt = state.metadata_mut().get_mut::<MOpt>().unwrap();
+            swarm_now = mopt.swarm_now;
+            mopt.update_pilot_operator_ctr_last(swarm_now);
+        }
+
+        for _i in 0..self.iterations(state, input) {
+            let idx = self.schedule(state, input);
+            let outcome = self
+                .mutations_mut()
+                .get_and_mutate(idx, state, input, stage_idx)?;
+            if outcome == MutationResult::Mutated {
+                r = MutationResult::Mutated;
+            }
+
+            state
+                .metadata_mut()
+                .get_mut::<MOpt>()
+                .unwrap()
+                .pilot_operator_ctr_this[swarm_now][idx] += 1;
+        }
+
+        Ok(r)
+    }
+}
+
+impl<I, MT, R, S> ComposedByMutations<I, MT, S> for StdMOptMutator<I, MT, R, S>
+where
+    I: Input,
+    MT: MutatorsTuple<I, S>,
+    R: Rand,
+    S: HasRand<R> + HasMetadata,
+{
+    /// Get the mutations
+    #[inline]
+    fn mutations(&self) -> &MT {
+        &self.mutations
+    }
+
+    // Get the mutations (mut)
+    #[inline]
+    fn mutations_mut(&mut self) -> &mut MT {
+        &mut self.mutations
+    }
+}
+
+impl<I, MT, R, S> ScheduledMutator<I, MT, S> for StdMOptMutator<I, MT, R, S>
+where
+    I: Input,
+    MT: MutatorsTuple<I, S>,
+    R: Rand,
+    S: HasRand<R> + HasMetadata,
+{
+    /// Compute the number of iterations used to apply stacked mutations
+    fn iterations(&self, state: &mut S, _: &I) -> u64 {
+        1 << (1 + state.rand_mut().below(6))
+    }
+
+    /// Get the next mutation to apply
+    fn schedule(&self, state: &mut S, _: &I) -> usize {
+        state
+            .metadata_mut()
+            .get_mut::<MOpt>()
+            .unwrap()
+            .select_algorithm()
+            .unwrap()
+    }
+
+    fn scheduled_mutate(
+        &mut self,
+        state: &mut S,
+        input: &mut I,
+        stage_idx: i32,
+    ) -> Result<MutationResult, Error> {
+        let mode = state.metadata().get::<MOpt>().unwrap().key_module;
+        match mode {
+            MOptMode::Corefuzzing => self.core_mutate(state, input, stage_idx),
+            MOptMode::Pilotfuzzing => self.pilot_mutate(state, input, stage_idx),
+        }
+    }
+}
+
+pub trait MOptMutator<I, MT, R, S>: ScheduledMutator<I, MT, S>
+where
+    I: Input,
+    MT: MutatorsTuple<I, S>,
+    R: Rand,
+    S: HasRand<R> + HasMetadata,
+{
+}
+
+impl<I, MT, R, S> MOptMutator<I, MT, R, S> for StdMOptMutator<I, MT, R, S>
+where
+    I: Input,
+    MT: MutatorsTuple<I, S>,
+    R: Rand,
+    S: HasRand<R> + HasMetadata,
+{
+}

libafl/src/stages/mod.rs

@@ -11,6 +11,9 @@ pub use mutational::{MutationalStage, StdMutationalStage};
 pub mod tracing;
 pub use tracing::{ShadowTracingStage, TracingStage};
 
+pub mod mopt;
+pub use mopt::*;
+
 //pub mod power;
 //pub use power::PowerMutationalStage;
 use crate::Error;

libafl/src/stages/mopt.rs

@@ -0,0 +1,237 @@
+use core::marker::PhantomData;
+
+use crate::{
+    bolts::rands::Rand,
+    corpus::Corpus,
+    fuzzer::Evaluator,
+    inputs::Input,
+    mutators::{MOpt, MOptMode, MOptMutator, MutatorsTuple},
+    stages::{MutationalStage, Stage},
+    state::{HasClientPerfStats, HasCorpus, HasMetadata, HasRand, HasSolutions},
+    Error,
+};
+
+const PERIOD_PILOT_COEF: f64 = 5000.0;
+
+#[derive(Clone, Debug)]
+pub struct MOptStage<C, E, EM, I, M, MT, R, S, SC, Z>
+where
+    C: Corpus<I>,
+    M: MOptMutator<I, MT, R, S>,
+    MT: MutatorsTuple<I, S>,
+    I: Input,
+    R: Rand,
+    S: HasClientPerfStats + HasCorpus<C, I> + HasSolutions<SC, I> + HasRand<R> + HasMetadata,
+    SC: Corpus<I>,
+    Z: Evaluator<E, EM, I, S>,
+{
+    mutator: M,
+    #[allow(clippy::type_complexity)]
+    phantom: PhantomData<(C, E, EM, I, MT, R, S, SC, Z)>,
+}
+
+impl<C, E, EM, I, M, MT, R, S, SC, Z> MutationalStage<C, E, EM, I, M, S, Z>
+    for MOptStage<C, E, EM, I, M, MT, R, S, SC, Z>
+where
+    C: Corpus<I>,
+    M: MOptMutator<I, MT, R, S>,
+    MT: MutatorsTuple<I, S>,
+    I: Input,
+    R: Rand,
+    S: HasClientPerfStats + HasCorpus<C, I> + HasSolutions<SC, I> + HasRand<R> + HasMetadata,
+    SC: Corpus<I>,
+    Z: Evaluator<E, EM, I, S>,
+{
+    /// 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 {
+        // TODO: we want to use calculate_score here
+
+        1 + state.rand_mut().below(128) as usize
+    }
+
+    #[allow(
+        clippy::cast_possible_wrap,
+        clippy::cast_precision_loss,
+        clippy::too_many_lines
+    )]
+    fn perform_mutational(
+        &mut self,
+        fuzzer: &mut Z,
+        executor: &mut E,
+        state: &mut S,
+        manager: &mut EM,
+        corpus_idx: usize,
+    ) -> Result<(), Error> {
+        let key_module = state.metadata().get::<MOpt>().unwrap().key_module;
+
+        match key_module {
+            MOptMode::Corefuzzing => {
+                let num = self.iterations(state);
+
+                for stage_id in 0..num {
+                    let mut input = state
+                        .corpus()
+                        .get(corpus_idx)?
+                        .borrow_mut()
+                        .load_input()?
+                        .clone();
+
+                    self.mutator_mut()
+                        .mutate(state, &mut input, stage_id as i32)?;
+
+                    let finds_before = state.corpus().count() + state.solutions().count();
+
+                    let (_, corpus_idx) = fuzzer.evaluate_input(state, executor, manager, input)?;
+
+                    self.mutator_mut()
+                        .post_exec(state, stage_id as i32, corpus_idx)?;
+
+                    let finds_after = state.corpus().count() + state.solutions().count();
+
+                    let mopt = state.metadata_mut().get_mut::<MOpt>().unwrap();
+
+                    mopt.core_time += 1;
+
+                    if finds_after > finds_before {
+                        let diff = finds_after - finds_before;
+                        mopt.total_finds += diff;
+                        for i in 0..mopt.operator_num {
+                            if mopt.core_operator_ctr_this[i] > mopt.core_operator_ctr_last[i] {
+                                mopt.core_operator_finds_this[i] += diff;
+                            }
+                        }
+                    }
+
+                    if mopt.core_time > mopt.period_core {
+                        // Make a call to pso_update()
+                        mopt.core_time = 0;
+                        let total_finds = mopt.total_finds;
+                        mopt.finds_before_switch = total_finds;
+                        mopt.update_core_operator_ctr_pso();
+                        mopt.pso_update()?;
+                    }
+                }
+            }
+            MOptMode::Pilotfuzzing => {
+                let num = self.iterations(state);
+                for stage_id in 0..num {
+                    let mut input = state
+                        .corpus()
+                        .get(corpus_idx)?
+                        .borrow_mut()
+                        .load_input()?
+                        .clone();
+
+                    self.mutator_mut()
+                        .mutate(state, &mut input, stage_id as i32)?;
+
+                    let finds_before = state.corpus().count() + state.solutions().count();
+
+                    let (_, corpus_idx) = fuzzer.evaluate_input(state, executor, manager, input)?;
+
+                    self.mutator_mut()
+                        .post_exec(state, stage_id as i32, corpus_idx)?;
+
+                    let finds_after = state.corpus().count() + state.solutions().count();
+
+                    let mopt = state.metadata_mut().get_mut::<MOpt>().unwrap();
+
+                    mopt.pilot_time += 1;
+                    let swarm_now = mopt.swarm_now;
+
+                    if finds_after > finds_before {
+                        let diff = finds_after - finds_before;
+                        mopt.total_finds += diff;
+                        for i in 0..mopt.operator_num {
+                            if mopt.pilot_operator_ctr_this[swarm_now][i]
+                                > mopt.pilot_operator_ctr_last[swarm_now][i]
+                            {
+                                mopt.pilot_operator_finds_this[swarm_now][i] += diff;
+                            }
+                        }
+                    }
+
+                    if mopt.pilot_time > mopt.period_pilot {
+                        let new_finds = mopt.total_finds - mopt.finds_before_switch;
+                        let f =
+                            (new_finds as f64) / ((mopt.pilot_time as f64) / (PERIOD_PILOT_COEF));
+                        mopt.swarm_fitness[swarm_now] = f;
+                        mopt.pilot_time = 0;
+                        let total_finds = mopt.total_finds;
+                        mopt.finds_before_switch = total_finds;
+                        mopt.update_pilot_operator_ctr_pso(swarm_now);
+
+                        mopt.swarm_now += 1;
+
+                        if mopt.swarm_now == mopt.swarm_num {
+                            // Move to CORE_FUZING mode
+                            mopt.key_module = MOptMode::Corefuzzing;
+
+                            mopt.init_core_module()?;
+                        }
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
+}
+
+impl<C, E, EM, I, M, MT, R, S, SC, Z> Stage<E, EM, S, Z>
+    for MOptStage<C, E, EM, I, M, MT, R, S, SC, Z>
+where
+    C: Corpus<I>,
+    M: MOptMutator<I, MT, R, S>,
+    MT: MutatorsTuple<I, S>,
+    I: Input,
+    R: Rand,
+    S: HasClientPerfStats + HasCorpus<C, I> + HasSolutions<SC, I> + HasRand<R> + HasMetadata,
+    SC: Corpus<I>,
+    Z: Evaluator<E, EM, I, S>,
+{
+    #[inline]
+    fn perform(
+        &mut self,
+        fuzzer: &mut Z,
+        executor: &mut E,
+        state: &mut S,
+        manager: &mut EM,
+        corpus_idx: usize,
+    ) -> Result<(), Error> {
+        self.perform_mutational(fuzzer, executor, state, manager, corpus_idx)
+    }
+}
+
+impl<C, E, EM, I, M, MT, R, S, SC, Z> MOptStage<C, E, EM, I, M, MT, R, S, SC, Z>
+where
+    C: Corpus<I>,
+    M: MOptMutator<I, MT, R, S>,
+    MT: MutatorsTuple<I, S>,
+    I: Input,
+    R: Rand,
+    S: HasClientPerfStats + HasCorpus<C, I> + HasSolutions<SC, I> + HasRand<R> + HasMetadata,
+    SC: Corpus<I>,
+    Z: Evaluator<E, EM, I, S>,
+{
+    /// Creates a new default mutational stage
+    pub fn new(mutator: M, state: &mut S, swarm_num: usize) -> Result<Self, Error> {
+        state.add_metadata::<MOpt>(MOpt::new(mutator.mutations().len(), swarm_num)?);
+        Ok(Self {
+            mutator,
+            phantom: PhantomData,
+        })
+    }
+}