Merge pull request #331 from tfreiberg-fastly/tfreiberg/entry-ref

Add (try_)entry_ref functions and EntryRef types

Author: xacrimon

src/lib.rs

@@ -27,6 +27,9 @@ use crate::lock::RwLock;
 
 #[cfg(feature = "raw-api")]
 pub use crate::lock::{RawRwLock, RwLock};
+use crate::mapref::entry_ref::EntryRef;
+use crate::mapref::entry_ref::OccupiedEntryRef;
+use crate::mapref::entry_ref::VacantEntryRef;
 
 use cfg_if::cfg_if;
 use core::fmt;
@@ -889,6 +892,17 @@ impl<'a, K: 'a + Eq + Hash, V: 'a, S: BuildHasher + Clone> DashMap<K, V, S> {
         self._try_entry(key)
     }
 
+    /// Advanced entry API that tries to mimic `hashbrown::HashMap::entry_ref`.
+    /// See the documentation on `dashmap::mapref::entry_ref` for more details.
+    ///
+    /// **Locking behaviour:** May deadlock if called when holding any sort of reference into the map.
+    pub fn entry_ref<'q, Q>(&'a self, key: &'q Q) -> EntryRef<'a, 'q, K, Q, V>
+    where
+        Q: Hash + Equivalent<K>,
+    {
+        self._entry_ref(key)
+    }
+
     /// Advanced entry API that tries to mimic `std::collections::HashMap::try_reserve`.
     /// Tries to reserve capacity for at least `shard * additional`
     /// and may reserve more space to avoid frequent reallocations.
@@ -1182,6 +1196,69 @@ impl<'a, K: 'a + Eq + Hash, V: 'a, S: 'a + BuildHasher + Clone> DashMap<K, V, S>
         }
     }
 
+    fn _entry_ref<'q, Q>(&'a self, key: &'q Q) -> EntryRef<'a, 'q, K, Q, V>
+    where
+        Q: Hash + Equivalent<K>,
+    {
+        let hash = self.hash_u64(&key);
+
+        let idx = self.determine_shard(hash as usize);
+
+        let shard = self.shards[idx].write();
+        // SAFETY: The data will not outlive the guard, since we pass the guard to `Entry`.
+        let (guard, shard) = unsafe { RwLockWriteGuardDetached::detach_from(shard) };
+
+        match shard.entry(
+            hash,
+            |(k, _v)| key.equivalent(k),
+            |(k, _v)| {
+                let mut hasher = self.hasher.build_hasher();
+                k.hash(&mut hasher);
+                hasher.finish()
+            },
+        ) {
+            hash_table::Entry::Occupied(entry) => {
+                EntryRef::Occupied(OccupiedEntryRef::new(guard, key, entry))
+            }
+            hash_table::Entry::Vacant(entry) => {
+                EntryRef::Vacant(VacantEntryRef::new(guard, key, entry))
+            }
+        }
+    }
+
+    fn _try_entry_ref<'q, Q>(&'a self, key: &'q Q) -> Option<EntryRef<'a, 'q, K, Q, V>>
+    where
+        Q: Hash + Equivalent<K>,
+    {
+        let hash = self.hash_u64(&key);
+
+        let idx = self.determine_shard(hash as usize);
+
+        let shard = match self.shards[idx].try_write() {
+            Some(shard) => shard,
+            None => return None,
+        };
+        // SAFETY: The data will not outlive the guard, since we pass the guard to `Entry`.
+        let (guard, shard) = unsafe { RwLockWriteGuardDetached::detach_from(shard) };
+
+        match shard.entry(
+            hash,
+            |(k, _v)| key.equivalent(k),
+            |(k, _v)| {
+                let mut hasher = self.hasher.build_hasher();
+                k.hash(&mut hasher);
+                hasher.finish()
+            },
+        ) {
+            hash_table::Entry::Occupied(entry) => {
+                Some(EntryRef::Occupied(OccupiedEntryRef::new(guard, key, entry)))
+            }
+            hash_table::Entry::Vacant(entry) => {
+                Some(EntryRef::Vacant(VacantEntryRef::new(guard, key, entry)))
+            }
+        }
+    }
+
     fn _clear(&self) {
         self._retain(|_, _| false)
     }

src/mapref/entry.rs

@@ -214,6 +214,42 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    fn test_insert_into_vacant() {
+        let map: DashMap<u32, u32> = DashMap::new();
+
+        let entry = map.entry(1);
+
+        assert!(matches!(entry, Entry::Vacant(_)));
+
+        let val = entry.insert(2);
+
+        assert_eq!(*val, 2);
+
+        drop(val);
+
+        assert_eq!(*map.get(&1).unwrap(), 2);
+    }
+
+    #[test]
+    fn test_insert_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 val = entry.insert(2);
+
+        assert_eq!(*val, 2);
+
+        drop(val);
+
+        assert_eq!(*map.get(&1).unwrap(), 2);
+    }
+
     #[test]
     fn test_insert_entry_into_vacant() {
         let map: DashMap<u32, u32> = DashMap::new();