From 3b418466edd8918244d9ecc30ddff3b37ae73c4c Mon Sep 17 00:00:00 2001
From: Tobias Kahlert <code.databyte@gmail.com>
Date: Tue, 11 Jul 2023 05:29:24 +0200
Subject: [PATCH] Entry insert (#226)

* added Entry::insert_entry

* added Entry::insert
---
 src/mapref/entry.rs | 87 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 87 insertions(+)

diff --git a/src/mapref/entry.rs b/src/mapref/entry.rs
index 16a42ce7..3b77e850 100644
--- a/src/mapref/entry.rs
+++ b/src/mapref/entry.rs
@@ -82,6 +82,33 @@ impl<'a, K: Eq + Hash, V, S: BuildHasher> Entry<'a, K, V, S> {
             Entry::Vacant(entry) => Ok(entry.insert(value()?)),
         }
     }
+
+    /// Sets the value of the entry, and returns a reference to the inserted value.
+    pub fn insert(self, value: V) -> RefMut<'a, K, V, S> {
+        match self {
+            Entry::Occupied(mut entry) => {
+                entry.insert(value);
+                entry.into_ref()
+            }
+            Entry::Vacant(entry) => entry.insert(value),
+        }
+    }
+
+    /// Sets the value of the entry, and returns an OccupiedEntry.
+    ///
+    /// If you are not interested in the occupied entry,
+    /// consider [`insert`] as it doesn't need to clone the key.
+    /// 
+    /// [`insert`]: Entry::insert
+    pub fn insert_entry(self, value: V) -> OccupiedEntry<'a, K, V, S> where K: Clone {
+        match self {
+            Entry::Occupied(mut entry) => {
+                entry.insert(value);
+                entry
+            }
+            Entry::Vacant(entry) => entry.insert_entry(value),
+        }
+    }
 }
 
 pub struct VacantEntry<'a, K, V, S = RandomState> {
@@ -117,6 +144,21 @@ impl<'a, K: Eq + Hash, V, S: BuildHasher> VacantEntry<'a, K, V, S> {
         }
     }
 
+    /// Sets the value of the entry with the VacantEntry’s key, and returns an OccupiedEntry.
+    pub fn insert_entry(mut self, value: V) -> OccupiedEntry<'a, K, V, S> where K: Clone {
+        unsafe {
+            self.shard.insert(self.key.clone(), SharedValue::new(value));
+
+            let (k, v) = self.shard.get_key_value(&self.key).unwrap();
+
+            let kptr: *const K = k;
+            let vptr: *mut V = v.as_ptr();
+            let r = OccupiedEntry::new(self.shard, self.key, (kptr, vptr));
+
+            r
+        }
+    }
+
     pub fn into_key(self) -> K {
         self.key
     }
@@ -187,3 +229,48 @@ impl<'a, K: Eq + Hash, V, S: BuildHasher> OccupiedEntry<'a, K, V, S> {
         (k, v.into_inner())
     }
 }
+
+
+
+#[cfg(test)]
+mod tests {
+    use crate::DashMap;
+
+    use super::*;
+
+    #[test]
+    fn test_insert_entry_into_vacant() {
+        let map: DashMap<u32, u32> = DashMap::new();
+
+        let entry = map.entry(1);
+
+        assert!(matches!(entry, Entry::Vacant(_)));
+        
+        let entry = entry.insert_entry(2);
+
+        assert_eq!(*entry.get(), 2);
+
+        drop(entry);
+
+        assert_eq!(*map.get(&1).unwrap(), 2);
+    }
+
+    #[test]
+    fn test_insert_entry_into_occupied() {
+        let map: DashMap<u32, u32> = DashMap::new();
+
+        map.insert(1, 1000);
+
+        let entry = map.entry(1);
+
+        assert!(matches!(&entry, Entry::Occupied(entry) if *entry.get() == 1000));
+        
+        let entry = entry.insert_entry(2);
+
+        assert_eq!(*entry.get(), 2);
+
+        drop(entry);
+
+        assert_eq!(*map.get(&1).unwrap(), 2);
+    }
+}
\ No newline at end of file