From e75f65080e6687cc5d9c95a96eb7805b8b387749 Mon Sep 17 00:00:00 2001
From: Langston Barrett <langston.barrett@gmail.com>
Date: Sun, 5 Feb 2023 19:44:57 -0500
Subject: [PATCH] libafl: with_capacity method for NewHashFeedback (#1034)

Most of the time, fuzzing campaigns are reasonably long. Accordingly, when
using `NewHashFeedback`, you might have good reason to believe that you'll
find thousands (or more) different observations with different hashes. When
the `HashSet` outgrows its capacity, it can cause reallocation, which is slow.
See the following link for more details:

https://doc.rust-lang.org/std/vec/struct.Vec.html#capacity-and-reallocation
---
 libafl/src/feedbacks/new_hash_feedback.rs | 24 ++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/libafl/src/feedbacks/new_hash_feedback.rs b/libafl/src/feedbacks/new_hash_feedback.rs
index 77f4241e87..0b74cc49fa 100644
--- a/libafl/src/feedbacks/new_hash_feedback.rs
+++ b/libafl/src/feedbacks/new_hash_feedback.rs
@@ -45,6 +45,14 @@ impl NewHashFeedbackMetadata {
         Self::default()
     }
 
+    /// Create a new [`NewHashFeedbackMetadata`] with the given initial capacity
+    #[must_use]
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            hash_set: HashSet::with_capacity(capacity),
+        }
+    }
+
     /// Reset the internal state
     pub fn reset(&mut self) -> Result<(), Error> {
         self.hash_set.clear();
@@ -71,6 +79,8 @@ impl HashSetState<u64> for NewHashFeedbackMetadata {
 pub struct NewHashFeedback<O, S> {
     name: String,
     observer_name: String,
+    /// Initial capacity of hash set
+    capacity: usize,
     o_type: PhantomData<(O, S)>,
 }
 
@@ -80,7 +90,10 @@ where
     S: UsesInput + Debug + HasNamedMetadata + HasClientPerfMonitor,
 {
     fn init_state(&mut self, state: &mut S) -> Result<(), Error> {
-        state.add_named_metadata(NewHashFeedbackMetadata::default(), &self.name);
+        state.add_named_metadata(
+            NewHashFeedbackMetadata::with_capacity(self.capacity),
+            &self.name,
+        );
         Ok(())
     }
 
@@ -146,6 +159,7 @@ where
         Self {
             name: name.to_string(),
             observer_name: observer_name.to_string(),
+            capacity: 0,
             o_type: PhantomData,
         }
     }
@@ -153,9 +167,17 @@ where
     /// Returns a new [`NewHashFeedback`].
     #[must_use]
     pub fn new(observer: &O) -> Self {
+        Self::with_capacity(observer, 0)
+    }
+
+    /// Returns a new [`NewHashFeedback`] that will create a hash set with the
+    /// given initial capacity.
+    #[must_use]
+    pub fn with_capacity(observer: &O, capacity: usize) -> Self {
         Self {
             name: NEWHASHFEEDBACK_PREFIX.to_string() + observer.name(),
             observer_name: observer.name().to_string(),
+            capacity,
             o_type: PhantomData,
         }
     }