From 231caf0797c216c5dc377925c8bf2562027690d5 Mon Sep 17 00:00:00 2001
From: Andrea Fioraldi <andreafioraldi@gmail.com>
Date: Mon, 6 Sep 2021 10:00:26 +0200
Subject: [PATCH] reset() method for FeedbackState

---
 libafl/src/feedbacks/map.rs |  9 +++++++--
 libafl/src/feedbacks/mod.rs | 21 ++++++++++++++++++---
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/libafl/src/feedbacks/map.rs b/libafl/src/feedbacks/map.rs
index 7d0545edf9..fa6734e2af 100644
--- a/libafl/src/feedbacks/map.rs
+++ b/libafl/src/feedbacks/map.rs
@@ -132,9 +132,14 @@ where
     pub name: String,
 }
 
-impl<T> FeedbackState for MapFeedbackState<T> where
-    T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned
+impl<T> FeedbackState for MapFeedbackState<T>
+where
+    T: Integer + Default + Copy + 'static + serde::Serialize + serde::de::DeserializeOwned,
 {
+    fn reset(&mut self) -> Result<(), Error> {
+        self.history_map.iter_mut().for_each(|x| *x = T::default());
+        Ok(())
+    }
 }
 
 impl<T> Named for MapFeedbackState<T>
diff --git a/libafl/src/feedbacks/mod.rs b/libafl/src/feedbacks/mod.rs
index 531b4494fc..eaaa551993 100644
--- a/libafl/src/feedbacks/mod.rs
+++ b/libafl/src/feedbacks/mod.rs
@@ -96,18 +96,33 @@ where
 
 /// [`FeedbackState`] is the data associated with a [`Feedback`] that must persist as part
 /// of the fuzzer State
-pub trait FeedbackState: Named + serde::Serialize + serde::de::DeserializeOwned {}
+pub trait FeedbackState: Named + serde::Serialize + serde::de::DeserializeOwned {
+    /// Reset the internal state
+    fn reset(&mut self) -> Result<(), Error> {
+        Ok(())
+    }
+}
 
 /// A haskell-style tuple of feedback states
-pub trait FeedbackStatesTuple: MatchName + serde::Serialize + serde::de::DeserializeOwned {}
+pub trait FeedbackStatesTuple: MatchName + serde::Serialize + serde::de::DeserializeOwned {
+    fn reset_all(&mut self) -> Result<(), Error>;
+}
 
-impl FeedbackStatesTuple for () {}
+impl FeedbackStatesTuple for () {
+    fn reset_all(&mut self) -> Result<(), Error> {
+        Ok(())
+    }
+}
 
 impl<Head, Tail> FeedbackStatesTuple for (Head, Tail)
 where
     Head: FeedbackState,
     Tail: FeedbackStatesTuple,
 {
+    fn reset_all(&mut self) -> Result<(), Error> {
+        self.0.reset()?;
+        self.1.reset_all()
+    }
 }
 
 pub struct CombinedFeedback<A, B, I, S, FL>