From b2a6ddb0dd9cc8c8b54860ae1a5b202537635d24 Mon Sep 17 00:00:00 2001
From: Andrea Fioraldi <andreafioraldi@gmail.com>
Date: Wed, 4 Nov 2020 15:29:46 +0100
Subject: [PATCH] finish DefaultMutationalStage

---
 src/stages/mod.rs        |  5 ++++-
 src/stages/mutational.rs | 28 ++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/stages/mod.rs b/src/stages/mod.rs
index 06c25616c3..55a1ade03d 100644
--- a/src/stages/mod.rs
+++ b/src/stages/mod.rs
@@ -5,6 +5,9 @@ use crate::inputs::Input;
 use crate::engines::Engine;
 use crate::AflError;
 
+use std::cell::RefCell;
+use std::rc::Rc;
+
 pub trait HasEngine<'a, I> where I: Input {
     type E : Engine<'a, I>;
 
@@ -15,5 +18,5 @@ pub trait HasEngine<'a, I> where I: Input {
 
 pub trait Stage<'a, I> : HasEngine<'a, I> where I: Input {
     /// Run the stage
-    fn perform(&mut self, entry: &mut Testcase<I>) -> Result<(), AflError>;
+    fn perform(&mut self, entry: Rc<RefCell<Testcase<I>>>) -> Result<(), AflError>;
 }
diff --git a/src/stages/mutational.rs b/src/stages/mutational.rs
index 5de64cdf8e..713c45ff85 100644
--- a/src/stages/mutational.rs
+++ b/src/stages/mutational.rs
@@ -9,6 +9,8 @@ use crate::engines::Engine;
 use std::cell::RefCell;
 use std::rc::Rc;
 
+// TODO create HasMutatorsVec trait
+
 pub trait MutationalStage<'a, I> : Stage<'a, I> + HasRand where I: Input {
     fn mutators(&self) -> &Vec<Box<dyn Mutator<I, R = Self::R>>>;
 
@@ -71,4 +73,30 @@ impl<'a, I, R, E> HasEngine<'a, I> for DefaultMutationalStage<'a, I, R, E> where
     fn engine_mut(&mut self) -> &mut Self::E {
         self.engine
     }
+}
+
+impl<'a, I, R, E> MutationalStage<'a, I> for DefaultMutationalStage<'a, I, R, E> where I: Input, R: Rand, E: Engine<'a, I> {
+    fn mutators(&self) -> &Vec<Box<dyn Mutator<I, R = Self::R>>> {
+        &self.mutators
+    }
+
+    fn mutators_mut(&mut self) -> &mut Vec<Box<dyn Mutator<I, R = Self::R>>> {
+        &mut self.mutators
+    }
+}
+
+impl<'a, I, R, E> Stage<'a, I> for DefaultMutationalStage<'a, I, R, E> where I: Input, R: Rand, E: Engine<'a, I> {
+    fn perform(&mut self, entry: Rc<RefCell<Testcase<I>>>) -> Result<(), AflError> {
+        self.perform_mutational(entry)
+    }
+}
+
+impl<'a, I, R, E> DefaultMutationalStage<'a, I, R, E> where I: Input, R: Rand, E: Engine<'a, I> {
+    pub fn new(rand: &'a mut R, engine: &'a mut E) -> Self {
+        DefaultMutationalStage {
+            rand: rand,
+            engine: engine,
+            mutators: vec![]
+        }
+    }
 }
\ No newline at end of file