From d3bae6a503cb482cf07fb7e22914ad5ba40823ee Mon Sep 17 00:00:00 2001
From: "Dongjia \"toka\" Zhang" <tokazerkje@outlook.com>
Date: Thu, 20 Feb 2025 19:27:35 +0100
Subject: [PATCH] Add try_insert for SerdeAnyMap (#3012)

* add

* use hashbrown method

---------

Co-authored-by: toka <toka@tokas-MacBook-Air.local>
---
 libafl/src/common/mod.rs     | 17 +++------------
 libafl_bolts/src/lib.rs      | 15 +++++++++++++
 libafl_bolts/src/serdeany.rs | 41 ++++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 14 deletions(-)

diff --git a/libafl/src/common/mod.rs b/libafl/src/common/mod.rs
index b55701208e..da1e2d4fa1 100644
--- a/libafl/src/common/mod.rs
+++ b/libafl/src/common/mod.rs
@@ -29,18 +29,11 @@ pub trait HasMetadata {
     /// Add a metadata to the metadata map
     /// Returns error if the metadata is already there
     #[inline]
-    fn add_metadata_checked<M>(&mut self, meta: M) -> Result<(), Error>
+    fn try_add_metadata<M>(&mut self, meta: M) -> Result<(), Error>
     where
         M: SerdeAny,
     {
-        if self.has_metadata::<M>() {
-            return Err(Error::illegal_argument(format!(
-                "Tried to add a metadata of {}. But this will overwrite the existing metadata",
-                type_name::<M>()
-            )));
-        }
-        self.metadata_map_mut().insert(meta);
-        Ok(())
+        self.metadata_map_mut().try_insert(meta)
     }
 
     /// Gets metadata, or inserts it using the given construction function `default`
@@ -118,11 +111,7 @@ pub trait HasNamedMetadata {
     where
         M: SerdeAny,
     {
-        if self.has_named_metadata::<M>(name) {
-            return Err(Error::illegal_argument(format!("Tried to add a metadata of {} named {}. But this will overwrite the existing metadata", type_name::<M>(), name)));
-        }
-        self.named_metadata_map_mut().insert(name, meta);
-        Ok(())
+        self.named_metadata_map_mut().try_insert(name, meta)
     }
 
     /// Add a metadata to the metadata map
diff --git a/libafl_bolts/src/lib.rs b/libafl_bolts/src/lib.rs
index ecef128cfd..95e539cf07 100644
--- a/libafl_bolts/src/lib.rs
+++ b/libafl_bolts/src/lib.rs
@@ -305,6 +305,8 @@ pub enum Error {
     EmptyOptional(String, ErrorBacktrace),
     /// Key not in Map
     KeyNotFound(String, ErrorBacktrace),
+    /// Key already exists and should not overwrite
+    KeyExists(String, ErrorBacktrace),
     /// No elements in the current item
     Empty(String, ErrorBacktrace),
     /// End of iteration
@@ -365,6 +367,15 @@ impl Error {
         Error::KeyNotFound(arg.into(), ErrorBacktrace::new())
     }
 
+    /// Key already exists in Map
+    #[must_use]
+    pub fn key_exists<S>(arg: S) -> Self
+    where
+        S: Into<String>,
+    {
+        Error::KeyExists(arg.into(), ErrorBacktrace::new())
+    }
+
     /// No elements in the current item
     #[must_use]
     pub fn empty<S>(arg: S) -> Self
@@ -508,6 +519,10 @@ impl Display for Error {
                 write!(f, "Key: `{0}` - not found", &s)?;
                 display_error_backtrace(f, b)
             }
+            Self::KeyExists(s, b) => {
+                write!(f, "Key: `{0}` - already exists", &s)?;
+                display_error_backtrace(f, b)
+            }
             Self::Empty(s, b) => {
                 write!(f, "No items in {0}", &s)?;
                 display_error_backtrace(f, b)
diff --git a/libafl_bolts/src/serdeany.rs b/libafl_bolts/src/serdeany.rs
index 6e2f8ffcb3..14be4c4864 100644
--- a/libafl_bolts/src/serdeany.rs
+++ b/libafl_bolts/src/serdeany.rs
@@ -321,6 +321,15 @@ pub mod serdeany_registry {
             self.insert_boxed(Box::new(t));
         }
 
+        /// Insert an element into the map if it doesn't exist, else return error.
+        #[inline]
+        pub fn try_insert<T>(&mut self, t: T) -> Result<(), Error>
+        where
+            T: crate::serdeany::SerdeAny,
+        {
+            self.try_insert_boxed(Box::new(t))
+        }
+
         /// Insert a boxed element into the map.
         #[inline]
         pub fn insert_boxed<T>(&mut self, value: Box<T>)
@@ -331,6 +340,21 @@ pub mod serdeany_registry {
                 .insert(type_repr_owned::<T>(), value);
         }
 
+        /// Insert a boxed element into the map if it doesn't exist, else return error.
+        #[inline]
+        pub fn try_insert_boxed<T>(&mut self, value: Box<T>) -> Result<(), Error>
+        where
+            T: crate::serdeany::SerdeAny,
+        {
+            match self.map.try_insert(type_repr_owned::<T>(), value) {
+                Ok(_) => (), // then it's fine
+                Err(hashbrown::hash_map::OccupiedError { entry: _, value }) => {
+                    return Err(Error::key_exists(format!("Tried to add a metadata of type {:?}. But this will overwrite the existing metadata value {:?}", core::any::type_name::<T>(), value)));
+                }
+            }
+            Ok(())
+        }
+
         /// Get an entry to an element in this map.
         #[inline]
         #[expect(unused_qualifications)]
@@ -600,6 +624,23 @@ pub mod serdeany_registry {
             self.entry::<T>(name.into()).insert(Box::new(val));
         }
 
+        /// Insert an element into the map if it doesn't exist, else return error.
+        #[inline]
+        #[expect(unused_qualifications)]
+        pub fn try_insert<T>(&mut self, name: &str, val: T) -> Result<(), Error>
+        where
+            T: crate::serdeany::SerdeAny,
+        {
+            let outer = self.outer_map_mut::<T>();
+            match outer.try_insert(name.into(), Box::new(val)) {
+                Ok(_) => (), // then it's fine
+                Err(hashbrown::hash_map::OccupiedError { entry, value }) => {
+                    return Err(Error::key_exists(format!("Tried to add a metadata of type {:?} named {:?}. But this will overwrite the existing metadata value {:?}", core::any::type_name::<T>(), entry.key(), value)));
+                }
+            }
+            Ok(())
+        }
+
         /// Get a reference to the type map.
         #[inline]
         #[expect(unused_qualifications)]