From c97d0fa7fca474af6132d068e3792316bd8f62ed Mon Sep 17 00:00:00 2001
From: Benjamin Beyret <beyretb@gmail.com>
Date: Wed, 8 Nov 2023 11:00:25 +0000
Subject: [PATCH] Allow MinimizerScheduler to not cleanup metadata after use
 (#1658)

* Allow MinimizerScheduler to not cleanup the metadata after use.

* Fix constructor

* cargo fmt and fix docstring.
---
 libafl/src/schedulers/minimizer.rs | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/libafl/src/schedulers/minimizer.rs b/libafl/src/schedulers/minimizer.rs
index d829af4a25..62d249b6d9 100644
--- a/libafl/src/schedulers/minimizer.rs
+++ b/libafl/src/schedulers/minimizer.rs
@@ -73,6 +73,7 @@ impl Default for TopRatedsMetadata {
 pub struct MinimizerScheduler<CS, F, M> {
     base: CS,
     skip_non_favored_prob: u64,
+    remove_metadata: bool,
     phantom: PhantomData<(F, M)>,
 }
 
@@ -293,7 +294,7 @@ where
                         old_meta.refcnt() <= 0
                     };
 
-                    if must_remove {
+                    if must_remove && self.remove_metadata {
                         drop(old.metadata_map_mut().remove::<M>());
                     }
                 }
@@ -304,7 +305,7 @@ where
             *meta.refcnt_mut() = new_favoreds.len() as isize;
         }
 
-        if new_favoreds.is_empty() {
+        if new_favoreds.is_empty() && self.remove_metadata {
             drop(
                 state
                     .corpus()
@@ -368,10 +369,25 @@ where
 
     /// Creates a new [`MinimizerScheduler`] that wraps a `base` [`Scheduler`]
     /// and has a default probability to skip non-faved [`Testcase`]s of [`DEFAULT_SKIP_NON_FAVORED_PROB`].
+    /// This will remove the metadata `M` when it is no longer needed, after consumption. This might
+    /// for example be a `MapIndexesMetadata`.
     pub fn new(base: CS) -> Self {
         Self {
             base,
             skip_non_favored_prob: DEFAULT_SKIP_NON_FAVORED_PROB,
+            remove_metadata: true,
+            phantom: PhantomData,
+        }
+    }
+
+    /// Creates a new [`MinimizerScheduler`] that wraps a `base` [`Scheduler`]
+    /// and has a default probability to skip non-faved [`Testcase`]s of [`DEFAULT_SKIP_NON_FAVORED_PROB`].
+    /// This method will prevent the metadata `M` from being removed at the end of scoring.
+    pub fn non_metadata_removing(base: CS) -> Self {
+        Self {
+            base,
+            skip_non_favored_prob: DEFAULT_SKIP_NON_FAVORED_PROB,
+            remove_metadata: false,
             phantom: PhantomData,
         }
     }
@@ -382,6 +398,7 @@ where
         Self {
             base,
             skip_non_favored_prob,
+            remove_metadata: true,
             phantom: PhantomData,
         }
     }