From 21ee8d2caeca6782ed491a318f437662286214b4 Mon Sep 17 00:00:00 2001
From: Elsa Granger <zeyugao@outlook.com>
Date: Sat, 8 Apr 2023 22:58:25 +0800
Subject: [PATCH] Fix on_remove of MinimizerScheduler (#1161)

* Fix cursor not step

* Update ref_cnt after remove

---------

Co-authored-by: Dongjia "toka" Zhang <tokazerkje@outlook.com>
Co-authored-by: Dominik Maier <domenukk@gmail.com>
---
 libafl/src/schedulers/minimizer.rs | 29 ++++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/libafl/src/schedulers/minimizer.rs b/libafl/src/schedulers/minimizer.rs
index 3955250fe1..66ecae2833 100644
--- a/libafl/src/schedulers/minimizer.rs
+++ b/libafl/src/schedulers/minimizer.rs
@@ -140,6 +140,8 @@ where
                                         }
                                     })
                                     .or_insert((factor, i));
+                                entry = e_iter.next();
+                                map_entry = map_iter.next();
                             }
                             Ordering::Greater => {
                                 map_entry = map_iter.next();
@@ -151,9 +153,30 @@ where
                 }
             }
         }
-        if let Some(meta) = state.metadata_map_mut().get_mut::<TopRatedsMetadata>() {
-            meta.map
-                .extend(map.into_iter().map(|(entry, (_, idx))| (entry, idx)));
+        if let Some(mut meta) = state.metadata_map_mut().remove::<TopRatedsMetadata>() {
+            let map_iter = map.iter();
+
+            let reserve = if meta.map.is_empty() {
+                map_iter.size_hint().0
+            } else {
+                (map_iter.size_hint().0 + 1) / 2
+            };
+            meta.map.reserve(reserve);
+
+            for (entry, (_, new_idx)) in map_iter {
+                let mut new = state.corpus().get(*new_idx)?.borrow_mut();
+                let new_meta = new.metadata_map_mut().get_mut::<M>().ok_or_else(|| {
+                    Error::key_not_found(format!(
+                        "{} needed for MinimizerScheduler not found in testcase #{new_idx}",
+                        type_name::<M>()
+                    ))
+                })?;
+                *new_meta.refcnt_mut() += 1;
+                meta.map.insert(*entry, *new_idx);
+            }
+
+            // Put back the metadata
+            state.metadata_map_mut().insert_boxed(meta);
         }
         Ok(())
     }