diff --git a/src/corpus/mod.rs b/src/corpus/mod.rs index f2bfc182c0..6015646a8e 100644 --- a/src/corpus/mod.rs +++ b/src/corpus/mod.rs @@ -30,23 +30,14 @@ pub trait Corpus: Debug { fn get(&mut self) -> Result<&Box, Error>; } -/// A queue-like corpus -pub trait Queue: Corpus { - fn get_cycles(&self) -> u64; - - fn get_pos(&self) -> usize; -} - #[derive(Debug)] -pub struct DefaultQueue<'a> { +pub struct RandomCorpus<'a> { rand: &'a mut dyn Rand, - pos: usize, - cycles: u64, entries: Vec>, dir_path: String, } -impl Corpus for DefaultQueue<'_> { +impl Corpus for RandomCorpus<'_> { /// Returns the number of elements fn count(&self) -> usize { self.entries.len() @@ -62,7 +53,7 @@ impl Corpus for DefaultQueue<'_> { let mut found = false; for x in &self.entries { i = i + 1; - if x.as_ref().get_filename() == entry.get_filename() { + if x.as_ref() as *const _ == entry as *const _ { found = true; break; } @@ -81,40 +72,80 @@ impl Corpus for DefaultQueue<'_> { /// Gets the next entry fn get(&mut self) -> Result<&Box, Error> { - if self.entries.len() == 0 { - return Err(Error::Unknown); + self.random_entry() + } +} + +impl RandomCorpus<'_> { + pub fn new<'a>(rand: &'a mut dyn Rand, dir_path: &str) -> RandomCorpus<'a> { + RandomCorpus { + dir_path: dir_path.to_owned(), + entries: vec![], + rand: rand, + } + } +} + +/// A queue-like corpus +#[derive(Debug)] +pub struct QueueCorpus<'a> { + random_corpus: RandomCorpus<'a>, + pos: usize, + cycles: u64, +} + +impl Corpus for QueueCorpus<'_> { + /// Returns the number of elements + fn count(&self) -> usize { + self.random_corpus.count() + } + + fn add(&mut self, entry: Box) { + self.random_corpus.add(entry); + } + + /// Removes an entry from the corpus, returning it if it was present. + fn remove(&mut self, entry: &dyn Testcase) -> Option> { + self.random_corpus.remove(entry) + } + + /// Gets a random entry + fn random_entry(&mut self) -> Result<&Box, Error> { + self.random_corpus.random_entry() + } + + /// Gets the next entry + fn get(&mut self) -> Result<&Box, Error> { + if self.count() == 0 { + return Err(Error::Unknown); // TODO(andrea) why unknown? use EmptyContainerError or similar } self.pos = self.pos + 1; - if self.pos >= self.entries.len() { + if self.pos >= self.count() { self.cycles = self.cycles + 1; self.pos = 0; } - Ok(self.entries.get_mut(self.pos).unwrap()) + Ok(self.random_corpus.entries.get_mut(self.pos).unwrap()) } } -impl Queue for DefaultQueue<'_> { - fn get_cycles(&self) -> u64 { +impl QueueCorpus<'_> { + pub fn new<'a>(rand: &'a mut dyn Rand, dir_path: &str) -> QueueCorpus<'a> { + QueueCorpus { + random_corpus: RandomCorpus::new(rand, dir_path), + cycles: 0, + pos: 0, + } + } + + pub fn get_cycles(&self) -> u64 { self.cycles } - fn get_pos(&self) -> usize { + pub fn get_pos(&self) -> usize { self.pos } } -impl DefaultQueue<'_> { - pub fn new<'a>(rand: &'a mut dyn Rand, dir_path: &str) -> DefaultQueue<'a> { - DefaultQueue { - cycles: 0, - dir_path: dir_path.to_owned(), - entries: vec![], - pos: 0, - rand: rand, - } - } -} - #[derive(Debug, Default)] struct SimpleTestcase { is_on_disk: bool, @@ -154,14 +185,14 @@ impl SimpleTestcase { #[cfg(test)] mod tests { use crate::corpus::Corpus; - use crate::corpus::DefaultQueue; + use crate::corpus::QueueCorpus; use crate::corpus::SimpleTestcase; use crate::utils::Xoshiro256StarRand; #[test] - fn test_defaultqueue() { + fn test_queuecorpus() { let mut rand = Xoshiro256StarRand::new(); - let mut q = DefaultQueue::new(&mut rand, "fancy/path"); + let mut q = QueueCorpus::new(&mut rand, "fancy/path"); q.add(Box::new(SimpleTestcase::new("fancyfile"))); let filename = q.get().unwrap().get_filename().to_owned(); assert_eq!(filename, q.get().unwrap().get_filename());