associated type for some traits

src/corpus/mod.rs

@@ -6,29 +6,30 @@ use crate::inputs::Input;
 use crate::AflError;
 
 use std::path::PathBuf;
+use std::marker::PhantomData;
 
-pub trait HasEntriesVec<InputT: Input> {
+pub trait HasEntriesVec<I> where I: Input {
     /// Get the entries vector field
-    fn entries(&self) -> &Vec<Box<Testcase<InputT>>>;
+    fn entries(&self) -> &Vec<Box<Testcase<I>>>;
 
     /// Get the entries vector field (mutable)
-    fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<InputT>>>;
+    fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>;
 }
 
 /// Corpus with all current testcases
-pub trait Corpus<InputT: Input, RandT: Rand> : HasEntriesVec<InputT> + HasRand<RandT> {
+pub trait Corpus<I> : HasEntriesVec<I> + HasRand where I: Input {
     /// Returns the number of elements
     fn count(&self) -> usize {
         self.entries().len()
     }
 
     /// Add an entry to the corpus
-    fn add(&mut self, mut entry: Box<Testcase<InputT>>) {
+    fn add(&mut self, mut entry: Box<Testcase<I>>) {
         self.entries_mut().push(entry);
     }
 
     /// Removes an entry from the corpus, returning it if it was present.
-    fn remove(&mut self, entry: &Testcase<InputT>) -> Option<Box<Testcase<InputT>>> {
+    fn remove(&mut self, entry: &Testcase<I>) -> Option<Box<Testcase<I>>> {
         let mut i: usize = 0;
         let mut found = false;
         for x in self.entries() {
@@ -45,46 +46,49 @@ pub trait Corpus<InputT: Input, RandT: Rand> : HasEntriesVec<InputT> + HasRand<R
     }
 
     /// Gets a random entry
-    fn random_entry(&mut self) -> Result<&Box<Testcase<InputT>>, AflError> {
+    fn random_entry(&mut self) -> Result<&Box<Testcase<I>>, AflError> {
-        let id = self.rand_mut().below(self.entries().len() as u64) as usize;
+        let len = { self.entries().len() };
+        let id = self.rand_mut().below(len as u64) as usize;
         Ok(self.entries_mut().get_mut(id).unwrap())
     }
 
     /// Gets the next entry (random by default)
-    fn get(&mut self) -> Result<&Box<Testcase<InputT>>, AflError> {
+    fn get(&mut self) -> Result<&Box<Testcase<I>>, AflError> {
         self.random_entry()
     }
 }
 
-pub struct InMemoryCorpus<'a, InputT: Input, RandT: Rand> {
+pub struct InMemoryCorpus<'a, I, R> where I: Input, R: Rand {
-    rand: &'a mut RandT,
+    rand: &'a mut R,
-    entries: Vec<Box<Testcase<InputT>>>,
+    entries: Vec<Box<Testcase<I>>>
 }
 
-impl<InputT: Input, RandT: Rand> HasEntriesVec<InputT> for InMemoryCorpus<'_, InputT, RandT> {
+impl<I, R> HasEntriesVec<I> for InMemoryCorpus<'_, I, R> where I: Input, R: Rand {
-    fn entries(&self) -> &Vec<Box<Testcase<InputT>>> {
+    fn entries(&self) -> &Vec<Box<Testcase<I>>> {
         &self.entries
     }
-    fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<InputT>>>{
+    fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>{
         &mut self.entries
     }
 }
 
-impl<InputT: Input, RandT: Rand> HasRand<RandT> for InMemoryCorpus<'_, InputT, RandT> {
+impl<I, R> HasRand for InMemoryCorpus<'_, I, R> where I: Input, R: Rand {
-    fn rand(&self) -> &Box<dyn Rand> {
+    type R = R;
+
+    fn rand(&self) -> &Self::R {
         &self.rand
     }
-    fn rand_mut(&mut self) -> &mut Box<dyn Rand> {
+    fn rand_mut(&mut self) -> &mut Self::R {
         &mut self.rand
     }
 }
 
-impl<InputT: Input, RandT: Rand> Corpus<InputT, RandT> for InMemoryCorpus<'_, InputT, RandT> {
+impl<I, R> Corpus<I> for InMemoryCorpus<'_, I, R> where I: Input, R: Rand {
     // Just use the default implementation
 }
 
-impl<InputT: Input, RandT: Rand> InMemoryCorpus<'_, InputT, RandT> {
+impl<'a, I, R> InMemoryCorpus<'a, I, R> where I: Input, R: Rand {
-    pub fn new<'a>(rand: &'a mut RandT) -> Self {
+    pub fn new(rand: &'a mut R) -> Self {
         InMemoryCorpus {
             rand: rand,
             entries: vec![],
@@ -92,38 +96,40 @@ impl<InputT: Input, RandT: Rand> InMemoryCorpus<'_, InputT, RandT> {
     }
 }
 
-pub struct OnDiskCorpus<'a, InputT: Input, RandT: Rand> {
+pub struct OnDiskCorpus<'a, I, R> where I: Input, R: Rand {
-    rand: &'a mut RandT,
+    rand: &'a mut R,
-    entries: Vec<Box<Testcase<InputT>>>,
+    entries: Vec<Box<Testcase<I>>>,
     dir_path: PathBuf,
 }
 
-impl<InputT: Input, RandT: Rand> HasEntriesVec<InputT> for OnDiskCorpus<'_, InputT, RandT> {
+impl<I, R> HasEntriesVec<I> for OnDiskCorpus<'_, I, R> where I: Input, R: Rand {
-    fn entries(&self) -> &Vec<Box<Testcase<InputT>>> {
+    fn entries(&self) -> &Vec<Box<Testcase<I>>> {
         &self.entries
     }
-    fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<InputT>>>{
+    fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>{
         &mut self.entries
     }
 }
 
-impl<InputT: Input, RandT: Rand> HasRand<RandT> for OnDiskCorpus<'_, InputT, RandT> {
+impl<I, R> HasRand for OnDiskCorpus<'_, I, R> where I: Input, R: Rand {
-    fn rand(&self) -> &Box<dyn Rand> {
+    type R = R;
+
+    fn rand(&self) -> &Self::R {
         &self.rand
     }
-    fn rand_mut(&mut self) -> &mut Box<dyn Rand> {
+    fn rand_mut(&mut self) -> &mut Self::R {
         &mut self.rand
     }
 }
 
-impl<InputT: Input, RandT: Rand> Corpus<InputT, RandT> for OnDiskCorpus<'_, InputT, RandT> {
+impl<I, R> Corpus<I> for OnDiskCorpus<'_, I, R> where I: Input, R: Rand {
     /// Add an entry and save it to disk
-    fn add(&mut self, mut entry: Box<Testcase<InputT>>) {
+    fn add(&mut self, mut entry: Box<Testcase<I>>) {
-        if entry.filename() == None {
+        if *entry.filename() == None {
             // TODO walk entry metadatas to ask for pices of filename (e.g. :havoc in AFL)
             let filename = &(String::from("id:") + &self.entries.len().to_string());
             let filename = self.dir_path.join(filename);
-            entry.filename_mut() = filename;
+            *entry.filename_mut() = Some(filename);
         }
         self.entries.push(entry);
     }
@@ -131,8 +137,8 @@ impl<InputT: Input, RandT: Rand> Corpus<InputT, RandT> for OnDiskCorpus<'_, Inpu
     // TODO save and remove files, cache, etc..., ATM use just InMemoryCorpus
 }
 
-impl<InputT: Input, RandT: Rand> OnDiskCorpus<'_, InputT, RandT> {
+impl<'a, I, R> OnDiskCorpus<'a, I, R> where I: Input, R: Rand {
-    pub fn new<'a>(rand: &'a mut RandT, dir_path: PathBuf) -> Self {
+    pub fn new(rand: &'a mut R, dir_path: PathBuf) -> Self {
         OnDiskCorpus {
             dir_path: dir_path,
             entries: vec![],
@@ -142,50 +148,72 @@ impl<InputT: Input, RandT: Rand> OnDiskCorpus<'_, InputT, RandT> {
 }
 
 /// A Queue-like corpus, wrapping an existing Corpus instance
-pub struct QueueCorpus<'a, InputT: Input, RandT: Rand, CorpusT: Corpus<InputT, RandT>> {
+pub struct QueueCorpus<I, C> where I: Input, C: Corpus<I> {
-    corpus: CorpusT,
+    corpus: C,
+    phantom: PhantomData<I>,
     pos: usize,
     cycles: u64,
 }
 
-impl<'a, InputT: Input, RandT: Rand, CorpusT: Corpus<InputT, RandT>> Corpus<InputT, RandT> for QueueCorpus<'_, InputT, RandT, CorpusT> {
+impl<'a, I, C> HasEntriesVec<I> for QueueCorpus<I, C> where I: Input, C: Corpus<I> {
+    fn entries(&self) -> &Vec<Box<Testcase<I>>> {
+        self.corpus.entries()
+    }
+    fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>{
+        self.corpus.entries_mut()
+    }
+}
+
+impl<'a, I, C> HasRand for QueueCorpus<I, C> where I: Input, C: Corpus<I> {
+    type R = C::R;
+
+    fn rand(&self) -> &Self::R {
+        self.corpus.rand()
+    }
+    fn rand_mut(&mut self) -> &mut Self::R {
+        self.corpus.rand_mut()
+    }
+}
+
+impl<'a, I, C> Corpus<I> for QueueCorpus<I, C> where I: Input, C: Corpus<I> {
     /// Returns the number of elements
     fn count(&self) -> usize {
         self.corpus.count()
     }
 
-    fn add(&mut self, entry: Box<Testcase<InputT>>) {
+    fn add(&mut self, entry: Box<Testcase<I>>) {
         self.corpus.add(entry);
     }
 
     /// Removes an entry from the corpus, returning it if it was present.
-    fn remove(&mut self, entry: &Testcase<InputT>) -> Option<Box<Testcase<InputT>>> {
+    fn remove(&mut self, entry: &Testcase<I>) -> Option<Box<Testcase<I>>> {
         self.corpus.remove(entry)
     }
 
     /// Gets a random entry
-    fn random_entry(&mut self) -> Result<&Box<Testcase<InputT>>, AflError> {
+    fn random_entry(&mut self) -> Result<&Box<Testcase<I>>, AflError> {
         self.corpus.random_entry()
     }
 
     /// Gets the next entry
-    fn get(&mut self) -> Result<&Box<Testcase<InputT>>, AflError> {
+    fn get(&mut self) -> Result<&Box<Testcase<I>>, AflError> {
-        if self.count() == 0 {
+        if self.corpus.count() == 0 {
             return Err(AflError::Empty("Testcases".to_string()));
         }
         self.pos = self.pos + 1;
-        if self.pos >= self.count() {
+        if self.pos >= self.corpus.count() {
             self.cycles = self.cycles + 1;
             self.pos = 0;
         }
-        Ok(self.corpus.entries.get_mut(self.pos).unwrap())
+        Ok(self.corpus.entries_mut().get_mut(self.pos).unwrap())
     }
 }
 
-impl<'a, InputT: Input, RandT: Rand, CorpusT: Corpus<InputT, RandT>> QueueCorpus<'_, InputT, RandT, CorpusT> {
+impl<'a, I, C> QueueCorpus<I, C> where I: Input, C: Corpus<I> {
-    pub fn new(corpus: CorpusT) -> Self {
+    pub fn new(corpus: C) -> Self {
-        QueueCorpus {
+        QueueCorpus::<I, C> {
             corpus: corpus,
+            phantom: PhantomData,
             cycles: 0,
             pos: 0,
         }
@@ -217,10 +245,10 @@ mod tests {
         let mut q = QueueCorpus::new(OnDiskCorpus::new(&mut rand, PathBuf::from("fancy/path")));
         let i = Box::new(BytesInput::new(vec![0; 4]));
         let mut t = Box::new(Testcase::new(i));
-        t.set_filename(PathBuf::from("fancyfile"));
+        *t.filename_mut() = Some(PathBuf::from("fancyfile"));
         q.add(t);
-        let filename = q.get().unwrap().get_filename().unwrap().to_owned();
+        let filename = q.get().unwrap().filename().as_ref().unwrap().to_owned();
-        assert_eq!(filename, q.get().unwrap().get_filename().unwrap().to_owned());
+        assert_eq!(filename, q.get().unwrap().filename().as_ref().unwrap().to_owned());
         assert_eq!(filename, PathBuf::from("fancy/path/fancyfile"));
     }
 }

src/corpus/testcase.rs

@@ -7,15 +7,15 @@ use std::path::PathBuf;
 pub trait TestcaseMetadata {}
 
 /*
-pub trait TestcaseTrait<InputT: Input> {
+pub trait TestcaseTrait<I: Input> {
     /// Make sure to return a valid input instance loading it from disk if not in memory
-    fn load_input(&mut self) -> Result<&Box<InputT>, AflError>;
+    fn load_input(&mut self) -> Result<&Box<I>, AflError>;
 
     /// Get the input, if any
-    fn input(&self) -> &Option<Box<InputT>>;
+    fn input(&self) -> &Option<Box<I>>;
 
     /// Get the input, if any (mutable)
-    fn input_mut(&mut self) -> &mut Option<Box<InputT>>;
+    fn input_mut(&mut self) -> &mut Option<Box<I>>;
 
     /// Get the filename, if any
     fn filename(&self) -> &Option<PathBuf>;
@@ -29,38 +29,44 @@ pub trait TestcaseTrait<InputT: Input> {
 */
 
 #[derive(Default)]
-pub struct Testcase<InputT: Input> {
+pub struct Testcase<I> where I: Input {
-    input: Option<Box<InputT>>,
+    input: Option<Box<I>>,
     filename: Option<PathBuf>,
     metadatas: HashMap<String, Box<dyn TestcaseMetadata>>,
 }
 
-impl<InputT: Input> Testcase<InputT> {
+impl<I> Testcase<I> where I: Input {
     /// Make sure to return a valid input instance loading it from disk if not in memory
-    pub fn load_input(&mut self) -> Result<&Box<InputT>, AflError> {
+    pub fn load_input(&mut self) -> Result<&Box<I>, AflError> {
         // TODO: Implement cache to disk
         self.input.as_ref().ok_or(AflError::NotImplemented("load_input".to_string()))
     }
 
     /// Get the input, if any
-    pub fn input(&self) -> &Option<Box<InputT>> {
+    pub fn input(&self) -> &Option<Box<I>> {
         &self.input
     }
-
     /// Get the input, if any (mutable)
-    pub fn input_mut(&mut self) -> &mut Option<Box<InputT>> {
+    pub fn input_mut(&mut self) -> &mut Option<Box<I>> {
         &mut self.input
     }
+    /// Set the input
+    pub fn set_input(&mut self, input: Option<Box<I>>) {
+        self.input = input;
+    }
 
     /// Get the filename, if any
     pub fn filename(&self) -> &Option<PathBuf> {
         &self.filename
     }
-
     /// Get the filename, if any (mutable)
-    pub fn filename_mut(&mut self, filename: PathBuf) -> &mut Option<PathBuf> {
+    pub fn filename_mut(&mut self) -> &mut Option<PathBuf> {
         &mut self.filename
     }
+    /// Set the filename
+    pub fn set_filename(&mut self, filename: Option<PathBuf>) {
+        self.filename = filename;
+    }
 
     /// Get all the metadatas into an HashMap
     pub fn metadatas(&mut self) -> &mut HashMap<String, Box<dyn TestcaseMetadata>> {
@@ -68,7 +74,7 @@ impl<InputT: Input> Testcase<InputT> {
     }
 
     /// Create a new DefaultTestcase instace given an input
-    pub fn new(input: Box<InputT>) -> Self {
+    pub fn new(input: Box<I>) -> Self {
         Testcase {
             input: Some(input),
             filename: None,
@@ -77,10 +83,10 @@ impl<InputT: Input> Testcase<InputT> {
     }
 
     /// Create a new DefaultTestcase instace given an input and a filename
-    pub fn new_with_filename(input: Box<InputT>, filename: &PathBuf) -> Self {
+    pub fn new_with_filename(input: Box<I>, filename: PathBuf) -> Self {
         Testcase {
             input: Some(input),
-            filename: filename,
+            filename: Some(filename),
             metadatas: HashMap::default(),
         }
     }

src/engines/aflengine.rs

@@ -1,18 +1,19 @@
 use std::time;
 
 use crate::engines::Engine;
+use crate::inputs::Input;
 use crate::executors::Executor;
 use crate::feedbacks::Feedback;
 use crate::monitors::Monitor;
 use crate::stages::Stage;
 use crate::utils::Rand;
 
-pub struct AflEngine<'a> {
+pub struct AflEngine<'a, InputT: Input> {
     pub rand: &'a mut dyn Rand,
-    pub feedbacks: Vec<Box<dyn Feedback>>,
+    pub feedbacks: Vec<Box<dyn Feedback<InputT>>>,
 
-    pub stages: Vec<Box<dyn Stage>>,
+    pub stages: Vec<Box<dyn Stage<InputT>>>,
-    pub current_stage: &'a Box<dyn Stage>,
+    pub current_stage: &'a Box<dyn Stage<InputT>>,
 
     pub executor: Box<dyn Executor>,
 
@@ -25,4 +26,4 @@ pub struct AflEngine<'a> {
     pub monitors: Vec<Box<dyn Monitor>>,
 }
 
-impl Engine<'_> for AflEngine<'_> {}
+impl<InputT: Input> Engine<'_> for AflEngine<'_, InputT> {}

src/feedbacks/mod.rs

@@ -1,15 +1,22 @@
 extern crate num;
 
 use crate::corpus::Testcase;
+use crate::inputs::Input;
 use crate::executors::Executor;
 use crate::observers::MapObserver;
 
 use num::Integer;
 use std::marker::PhantomData;
 
+pub trait Feedback<InputT: Input> {
+    /// is_interesting should return the "Interestingness" from 0 to 255 (percent times 2.55)
+    fn is_interesting(&mut self, executor: &dyn Executor, entry: &Testcase<InputT>) -> u8;
+}
+
+/*
 pub trait Feedback {
     /// is_interesting should return the "Interestingness" from 0 to 255 (percent times 2.55)
-    fn is_interesting(&mut self, executor: &dyn Executor, entry: &dyn Testcase) -> u8;
+    fn is_interesting(&mut self, executor: &dyn Executor, entry: &Testcase) -> u8;
 }
 
 pub trait Reducer<T: Integer + Copy + 'static> {
@@ -69,3 +76,5 @@ impl<'a, MapT: Integer + Copy + 'static, ReducerT: Reducer<MapT>> MapFeedback<Ma
 type MaxMapFeedback<MapT> = MapFeedback<MapT, dyn MaxReducer<MapT>>;
 #[allow(dead_code)]
 type MinMapFeedback<MapT> = MapFeedback<MapT, dyn MinReducer<MapT>>;
+
+*/

src/mutators/mod.rs

@@ -1,20 +1,30 @@
 use crate::inputs::Input;
-use crate::utils::{Rand, HasRand};
+use crate::utils::HasRand;
 use crate::corpus::Corpus;
 use crate::AflError;
 
 pub mod scheduled;
 pub use scheduled::{ComposedByMutations, ScheduledMutator, HavocBytesMutator};
 
-pub trait Mutator<InputT : Input, RandT: Rand> : HasRand<RandT> {
+pub trait HasOptionCorpus<I> where I: Input {
+    type C : Corpus<I>;
+
+    /// Get the associated corpus, if any
+    fn corpus(&self) -> &Option<Box<Self::C>>;
+
+    /// Get the associated corpus, if any (mutable)
+    fn corpus_mut(&mut self) -> &mut Option<Box<Self::C>>;
+
+    /// Set the associated corpus
+    fn set_corpus(&mut self, corpus: Option<Box<Self::C>>);
+}
+
+pub trait Mutator<I> : HasRand + HasOptionCorpus<I> where I: Input {
     /// Mutate a given input
-    fn mutate(&mut self, input: &mut InputT, stage_idx: i32) -> Result<(), AflError>;
+    fn mutate(&mut self, input: &mut I, stage_idx: i32) -> Result<(), AflError>;
 
     /// Post-process given the outcome of the execution
     fn post_exec(&mut self, _is_interesting: bool, _stage_idx: i32) -> Result<(), AflError> {
         Ok(())
     }
-
-    /// Get the associated corpus, if any
-    fn corpus(&mut self) -> &mut Option<Box<dyn Corpus<InputT, RandT>>>;
 }

src/mutators/scheduled.rs

@@ -1,4 +1,4 @@
-use crate::mutators::Mutator;
+use crate::mutators::{HasOptionCorpus, Mutator};
 use crate::utils::{Rand, HasRand};
 use crate::corpus::Corpus;
 use crate::inputs::{Input, HasBytesVec};
@@ -7,27 +7,27 @@ use crate::AflError;
 use std::marker::PhantomData;
 
 /// The generic function type that identifies mutations
-type MutationFunction<MutatorT, InputT> = fn(&mut MutatorT, &mut InputT) -> Result<(), AflError>;
+type MutationFunction<M, I> = fn(&mut M, &mut I) -> Result<(), AflError>;
 
-pub trait ComposedByMutations<InputT : Input> {
+pub trait ComposedByMutations<I> where I: Input {
     /// Get a mutation by index
-    fn mutation_by_idx(&self, index: usize) -> Result<MutationFunction<Self, InputT>, AflError>;
+    fn mutation_by_idx(&self, index: usize) -> Result<MutationFunction<Self, I>, AflError>;
 
     /// Get the number of mutations
     fn mutations_count(&self) -> usize;
 
     /// Add a mutation
-    fn add_mutation(&mut self, mutation: MutationFunction<Self, InputT>);
+    fn add_mutation(&mut self, mutation: MutationFunction<Self, I>);
 }
 
-pub trait ScheduledMutator<InputT : Input, RandT: Rand>: Mutator<InputT, RandT> + ComposedByMutations<InputT> {
+pub trait ScheduledMutator<I>: Mutator<I> + ComposedByMutations<I> where I: Input {
-    /// Computer the number of iterations used to apply stacked mutations
+    /// Compute the number of iterations used to apply stacked mutations
-    fn iterations(&mut self, _input: &InputT) -> u64 {
+    fn iterations(&mut self, _input: &I) -> u64 {
         1 << (1 + self.rand_mut().below(7))
     }
 
     /// Get the next mutation to apply
-    fn schedule(&mut self, _input: &InputT) -> Result<MutationFunction<Self, InputT>, AflError> {
+    fn schedule(&mut self, _input: &I) -> Result<MutationFunction<Self, I>, AflError> {
         let count = self.mutations_count() as u64;
         if count == 0 {
             return Err(AflError::Empty("no mutations".to_string()));
@@ -41,7 +41,7 @@ pub trait ScheduledMutator<InputT : Input, RandT: Rand>: Mutator<InputT, RandT>
 
     /// New default implementation for mutate
     /// Implementations must forward mutate() to this method
-    fn scheduled_mutate(&mut self, input: &mut InputT, _stage_idx: i32) -> Result<(), AflError> {
+    fn scheduled_mutate(&mut self, input: &mut I, _stage_idx: i32) -> Result<(), AflError> {
         let num = self.iterations(input);
         for _ in 0..num {
             self.schedule(input)?(self, input)?;
@@ -50,23 +50,47 @@ pub trait ScheduledMutator<InputT : Input, RandT: Rand>: Mutator<InputT, RandT>
     }
 }
 
-pub struct DefaultScheduledMutator<InputT : Input, RandT: Rand> {
+pub struct DefaultScheduledMutator<'a, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
-    rand: Box<RandT>,
+    rand: &'a mut R,
-    corpus: Option<Box<dyn Corpus<InputT, RandT>>>,
+    corpus: Option<Box<C>>,
-    mutations: Vec<MutationFunction<Self, InputT>>
+    mutations: Vec<MutationFunction<Self, I>>
 }
 
-impl<InputT : Input, RandT: Rand> HasRand<RandT> for DefaultScheduledMutator<InputT, RandT> {
+impl<'a, I, R, C> HasRand for DefaultScheduledMutator<'_, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
-    fn rand(&self) -> &Box<dyn Rand> {
+    type R = R;
+
+    fn rand(&self) -> &Self::R {
         &self.rand
     }
-    fn rand_mut(&mut self) -> &mut Box<dyn Rand> {
+    fn rand_mut(&mut self) -> &mut Self::R {
         &mut self.rand
     }
 }
 
-impl<InputT : Input, RandT: Rand> ComposedByMutations<InputT> for DefaultScheduledMutator<InputT, RandT> {
+impl<I, R, C> HasOptionCorpus<I> for DefaultScheduledMutator<'_, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
-    fn mutation_by_idx(&self, index: usize) -> Result<MutationFunction<Self, InputT>, AflError> {
+    type C = C;
+
+    fn corpus(&self) -> &Option<Box<Self::C>> {
+        &self.corpus
+    }
+
+    fn corpus_mut(&mut self) -> &mut Option<Box<Self::C>> {
+        &mut self.corpus
+    }
+
+    fn set_corpus(&mut self, corpus: Option<Box<Self::C>>) {
+        self.corpus = corpus
+    }
+}
+
+impl<'a, I, R, C> Mutator<I> for DefaultScheduledMutator<'_, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
+    fn mutate(&mut self, input: &mut I, _stage_idx: i32) -> Result<(), AflError> {
+        self.scheduled_mutate(input, _stage_idx)
+    }
+}
+
+impl<'a, I, R, C> ComposedByMutations<I> for DefaultScheduledMutator<'_, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
+    fn mutation_by_idx(&self, index: usize) -> Result<MutationFunction<Self, I>, AflError> {
         if index >= self.mutations.len() {
             return Err(AflError::Unknown("oob".to_string()));
         }
@@ -77,28 +101,18 @@ impl<InputT : Input, RandT: Rand> ComposedByMutations<InputT> for DefaultSchedul
         self.mutations.len()
     }
 
-    fn add_mutation(&mut self, mutation: MutationFunction<Self, InputT>) {
+    fn add_mutation(&mut self, mutation: MutationFunction<Self, I>) {
         self.mutations.push(mutation)
     }
 }
 
-impl<InputT : Input, RandT: Rand> ScheduledMutator<InputT> for DefaultScheduledMutator<InputT, RandT> {
+impl<'a, I, R, C> ScheduledMutator<I> for DefaultScheduledMutator<'_, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
     // Just use the default methods
 }
 
-impl<InputT : Input, RandT: Rand> Mutator<InputT, RandT> for DefaultScheduledMutator<InputT, RandT> {
+impl<'a, I, R, C> DefaultScheduledMutator<'a, I, R, C> where I: Input, R: Rand, C: Corpus<I> {
-    fn mutate(&mut self, input: &mut InputT, _stage_idx: i32) -> Result<(), AflError> {
-        self.scheduled_mutate(input, _stage_idx)
-    }
-
-    fn corpus(&mut self) -> &mut Option<Box<dyn Corpus>> {
-        &mut self.corpus
-    }
-}
-
-impl<InputT : Input, RandT: Rand> DefaultScheduledMutator<InputT> {
     /// Create a new DefaultScheduledMutator instance without mutations and corpus
-    pub fn new(rand: Box<RandT>) -> Self {
+    pub fn new(rand: &'a mut R) -> Self {
         DefaultScheduledMutator {
             rand: rand,
             corpus: None,
@@ -107,7 +121,7 @@ impl<InputT : Input, RandT: Rand> DefaultScheduledMutator<InputT> {
     }
 
     /// Create a new DefaultScheduledMutator instance specifying mutations and corpus too
-    pub fn new_all(rand: Box<RandT>, corpus: Option<Box<dyn Corpus>>, mutations: Vec<MutationFunction<Self, InputT>>) -> Self {
+    pub fn new_all(rand: &'a mut R, corpus: Option<Box<C>>, mutations: Vec<MutationFunction<Self, I>>) -> Self {
         DefaultScheduledMutator {
             rand: rand,
             corpus: corpus,
@@ -117,56 +131,70 @@ impl<InputT : Input, RandT: Rand> DefaultScheduledMutator<InputT> {
 }
 
 /// Bitflip mutation for inputs with a bytes vector
-pub fn mutation_bitflip<MutatorT: Mutator<InputT, RandT>, InputT: Input + HasBytesVec, RandT: Rand>(mutator: &mut MutatorT, input: &mut InputT) -> Result<(), AflError> {
+pub fn mutation_bitflip<M, I>(mutator: &mut M, input: &mut I) -> Result<(), AflError> where M: Mutator<I>, I: Input + HasBytesVec {
     let bit = mutator.rand_mut().below(input.bytes().len() as u64) as usize;
     input.bytes_mut()[bit >> 3] ^= (128 >> (bit & 7)) as u8;
     Ok(())
 }
 
 /// Schedule some selected byte level mutations given a ScheduledMutator type
-pub struct HavocBytesMutator<InputT: Input + HasBytesVec, RandT: Rand, ScheduledMutatorT: ScheduledMutator<InputT, RandT>> {
+pub struct HavocBytesMutator<I, S> where I: Input + HasBytesVec, S: ScheduledMutator<I> {
-    scheduled: ScheduledMutatorT,
+    scheduled: S,
-    _phantom: PhantomData<InputT>
+    phantom: PhantomData<I>
 }
 
-impl<InputT: Input + HasBytesVec, RandT: Rand, ScheduledMutatorT: ScheduledMutator<InputT, RandT>> Mutator<InputT, RandT> for HavocBytesMutator<InputT, RandT, ScheduledMutatorT> {
+impl<I, S> HasRand for HavocBytesMutator<I, S> where I: Input + HasBytesVec, S: ScheduledMutator<I> {
-    fn mutate(&mut self, input: &mut InputT, stage_idx: i32) -> Result<(), AflError> {
+    type R = S::R;
-        self.scheduled.mutate(input, stage_idx)
-    }
 
-    fn corpus(&mut self) -> &mut Option<Box<dyn Corpus>> {
+    fn rand(&self) -> &Self::R {
-        self.scheduled.corpus()
-    }
-}
-
-impl<InputT: Input + HasBytesVec, RandT: Rand, ScheduledMutatorT: ScheduledMutator<InputT, RandT>> HasRand<RandT> for HavocBytesMutator<InputT, RandT, ScheduledMutatorT> {
-    fn rand(&self) -> &Box<dyn Rand> {
         self.scheduled.rand()
     }
-    fn rand_mut(&mut self) -> &mut Box<dyn Rand> {
+    fn rand_mut(&mut self) -> &mut Self::R {
         self.scheduled.rand_mut()
     }
 }
 
-impl<InputT: Input + HasBytesVec, RandT: Rand, ScheduledMutatorT: ScheduledMutator<InputT, RandT>> HavocBytesMutator<InputT, RandT, ScheduledMutatorT> {
+impl<I, S> HasOptionCorpus<I> for HavocBytesMutator<I, S> where I: Input + HasBytesVec, S: ScheduledMutator<I> {
+    type C = S::C;
+
+    fn corpus(&self) -> &Option<Box<Self::C>> {
+        self.scheduled.corpus()
+    }
+
+    fn corpus_mut(&mut self) -> &mut Option<Box<Self::C>> {
+        self.scheduled.corpus_mut()
+    }
+
+    fn set_corpus(&mut self, corpus: Option<Box<Self::C>>) {
+        self.scheduled.set_corpus(corpus)
+    }
+}
+
+impl<I, S> Mutator<I> for HavocBytesMutator<I, S> where I: Input + HasBytesVec, S: ScheduledMutator<I> {
+    fn mutate(&mut self, input: &mut I, stage_idx: i32) -> Result<(), AflError> {
+        self.scheduled.mutate(input, stage_idx)
+    }
+}
+
+impl<I, S> HavocBytesMutator<I, S> where I: Input + HasBytesVec, S: ScheduledMutator<I> {
     /// Create a new HavocBytesMutator instance given a ScheduledMutator to wrap
-    pub fn new(mut scheduled: ScheduledMutatorT) -> Self {
+    pub fn new(mut scheduled: S) -> Self {
         scheduled.add_mutation(mutation_bitflip);
         HavocBytesMutator {
             scheduled: scheduled,
-            _phantom: PhantomData
+            phantom: PhantomData
         }
     }
 }
 
-impl<InputT: Input + HasBytesVec, RandT: Rand> HavocBytesMutator<InputT, RandT, DefaultScheduledMutator<InputT, RandT>> {
+impl<'a, I, R, C> HavocBytesMutator<I, DefaultScheduledMutator<'a, I, R, C>> where I: Input + HasBytesVec, R: Rand, C: Corpus<I> {
     /// Create a new HavocBytesMutator instance wrapping DefaultScheduledMutator
-    pub fn new_default(rand: Box<RandT>) -> Self {
+    pub fn new_default(rand: &'a mut R) -> Self {
-        let mut scheduled = DefaultScheduledMutator::new(rand);
+        let mut scheduled = DefaultScheduledMutator::<'a, I, R, C>::new(rand);
         scheduled.add_mutation(mutation_bitflip);
         HavocBytesMutator {
             scheduled: scheduled,
-            _phantom: PhantomData
+            phantom: PhantomData
         }
     }
 }

src/stages/mod.rs

@@ -2,6 +2,6 @@ use crate::corpus::Testcase;
 use crate::inputs::Input;
 use crate::AflError;
 /// Stages
-pub trait Stage {
+pub trait Stage<InputT: Input> {
-    fn perform(&mut self, input: &dyn Input, entry: &mut Testcase) -> Result<(), AflError>;
+    fn perform(&mut self, input: &dyn Input, entry: &mut Testcase<InputT>) -> Result<(), AflError>;
 }

src/utils.rs

@@ -40,12 +40,14 @@ pub trait Rand: Debug {
 }
 
 /// Has a Rand box field
-pub trait HasRand<RandT: Rand> {
+pub trait HasRand {
+    type R : Rand;
+
     /// Get the hold Rand instance
-    fn rand(&self) -> &Box<RandT>;
+    fn rand(&self) -> &Self::R;
 
     /// Get the hold Rand instance (mutable)
-    fn rand_mut(&mut self) -> &mut Box<RandT>;
+    fn rand_mut(&mut self) -> &mut Self::R;
 }
 
 const HASH_CONST: u64 = 0xa5b35705;