dashmap/mapref/
entry.rs

1use super::one::RefMut;
2use crate::lock::RwLockWriteGuard;
3use crate::util::SharedValue;
4use crate::HashMap;
5use core::hash::Hash;
6use core::mem;
7
8pub enum Entry<'a, K, V> {
9    Occupied(OccupiedEntry<'a, K, V>),
10    Vacant(VacantEntry<'a, K, V>),
11}
12
13impl<'a, K: Eq + Hash, V> Entry<'a, K, V> {
14    /// Apply a function to the stored value if it exists.
15    pub fn and_modify(self, f: impl FnOnce(&mut V)) -> Self {
16        match self {
17            Entry::Occupied(mut entry) => {
18                f(entry.get_mut());
19
20                Entry::Occupied(entry)
21            }
22
23            Entry::Vacant(entry) => Entry::Vacant(entry),
24        }
25    }
26
27    /// Get the key of the entry.
28    pub fn key(&self) -> &K {
29        match *self {
30            Entry::Occupied(ref entry) => entry.key(),
31            Entry::Vacant(ref entry) => entry.key(),
32        }
33    }
34
35    /// Into the key of the entry.
36    pub fn into_key(self) -> K {
37        match self {
38            Entry::Occupied(entry) => entry.into_key(),
39            Entry::Vacant(entry) => entry.into_key(),
40        }
41    }
42
43    /// Return a mutable reference to the element if it exists,
44    /// otherwise insert the default and return a mutable reference to that.
45    pub fn or_default(self) -> RefMut<'a, K, V>
46    where
47        V: Default,
48    {
49        match self {
50            Entry::Occupied(entry) => entry.into_ref(),
51            Entry::Vacant(entry) => entry.insert(V::default()),
52        }
53    }
54
55    /// Return a mutable reference to the element if it exists,
56    /// otherwise a provided value and return a mutable reference to that.
57    pub fn or_insert(self, value: V) -> RefMut<'a, K, V> {
58        match self {
59            Entry::Occupied(entry) => entry.into_ref(),
60            Entry::Vacant(entry) => entry.insert(value),
61        }
62    }
63
64    /// Return a mutable reference to the element if it exists,
65    /// otherwise insert the result of a provided function and return a mutable reference to that.
66    pub fn or_insert_with(self, value: impl FnOnce() -> V) -> RefMut<'a, K, V> {
67        match self {
68            Entry::Occupied(entry) => entry.into_ref(),
69            Entry::Vacant(entry) => entry.insert(value()),
70        }
71    }
72
73    pub fn or_try_insert_with<E>(
74        self,
75        value: impl FnOnce() -> Result<V, E>,
76    ) -> Result<RefMut<'a, K, V>, E> {
77        match self {
78            Entry::Occupied(entry) => Ok(entry.into_ref()),
79            Entry::Vacant(entry) => Ok(entry.insert(value()?)),
80        }
81    }
82
83    /// Sets the value of the entry, and returns a reference to the inserted value.
84    pub fn insert(self, value: V) -> RefMut<'a, K, V> {
85        match self {
86            Entry::Occupied(mut entry) => {
87                entry.insert(value);
88                entry.into_ref()
89            }
90            Entry::Vacant(entry) => entry.insert(value),
91        }
92    }
93
94    /// Sets the value of the entry, and returns an OccupiedEntry.
95    ///
96    /// If you are not interested in the occupied entry,
97    /// consider [`insert`] as it doesn't need to clone the key.
98    ///
99    /// [`insert`]: Entry::insert
100    pub fn insert_entry(self, value: V) -> OccupiedEntry<'a, K, V>
101    where
102        K: Clone,
103    {
104        match self {
105            Entry::Occupied(mut entry) => {
106                entry.insert(value);
107                entry
108            }
109            Entry::Vacant(entry) => entry.insert_entry(value),
110        }
111    }
112}
113
114pub struct VacantEntry<'a, K, V> {
115    shard: RwLockWriteGuard<'a, HashMap<K, V>>,
116    key: K,
117    hash: u64,
118    slot: hashbrown::raw::InsertSlot,
119}
120
121unsafe impl<'a, K: Eq + Hash + Sync, V: Sync> Send for VacantEntry<'a, K, V> {}
122unsafe impl<'a, K: Eq + Hash + Sync, V: Sync> Sync for VacantEntry<'a, K, V> {}
123
124impl<'a, K: Eq + Hash, V> VacantEntry<'a, K, V> {
125    pub(crate) unsafe fn new(
126        shard: RwLockWriteGuard<'a, HashMap<K, V>>,
127        key: K,
128        hash: u64,
129        slot: hashbrown::raw::InsertSlot,
130    ) -> Self {
131        Self {
132            shard,
133            key,
134            hash,
135            slot,
136        }
137    }
138
139    pub fn insert(mut self, value: V) -> RefMut<'a, K, V> {
140        unsafe {
141            let occupied = self.shard.insert_in_slot(
142                self.hash,
143                self.slot,
144                (self.key, SharedValue::new(value)),
145            );
146
147            let (k, v) = occupied.as_ref();
148
149            RefMut::new(self.shard, k, v.as_ptr())
150        }
151    }
152
153    /// Sets the value of the entry with the VacantEntry’s key, and returns an OccupiedEntry.
154    pub fn insert_entry(mut self, value: V) -> OccupiedEntry<'a, K, V>
155    where
156        K: Clone,
157    {
158        unsafe {
159            let bucket = self.shard.insert_in_slot(
160                self.hash,
161                self.slot,
162                (self.key.clone(), SharedValue::new(value)),
163            );
164
165            OccupiedEntry::new(self.shard, self.key, bucket)
166        }
167    }
168
169    pub fn into_key(self) -> K {
170        self.key
171    }
172
173    pub fn key(&self) -> &K {
174        &self.key
175    }
176}
177
178pub struct OccupiedEntry<'a, K, V> {
179    shard: RwLockWriteGuard<'a, HashMap<K, V>>,
180    bucket: hashbrown::raw::Bucket<(K, SharedValue<V>)>,
181    key: K,
182}
183
184unsafe impl<'a, K: Eq + Hash + Sync, V: Sync> Send for OccupiedEntry<'a, K, V> {}
185unsafe impl<'a, K: Eq + Hash + Sync, V: Sync> Sync for OccupiedEntry<'a, K, V> {}
186
187impl<'a, K: Eq + Hash, V> OccupiedEntry<'a, K, V> {
188    pub(crate) unsafe fn new(
189        shard: RwLockWriteGuard<'a, HashMap<K, V>>,
190        key: K,
191        bucket: hashbrown::raw::Bucket<(K, SharedValue<V>)>,
192    ) -> Self {
193        Self { shard, bucket, key }
194    }
195
196    pub fn get(&self) -> &V {
197        unsafe { self.bucket.as_ref().1.get() }
198    }
199
200    pub fn get_mut(&mut self) -> &mut V {
201        unsafe { self.bucket.as_mut().1.get_mut() }
202    }
203
204    pub fn insert(&mut self, value: V) -> V {
205        mem::replace(self.get_mut(), value)
206    }
207
208    pub fn into_ref(self) -> RefMut<'a, K, V> {
209        unsafe {
210            let (k, v) = self.bucket.as_ref();
211            RefMut::new(self.shard, k, v.as_ptr())
212        }
213    }
214
215    pub fn into_key(self) -> K {
216        self.key
217    }
218
219    pub fn key(&self) -> &K {
220        unsafe { &self.bucket.as_ref().0 }
221    }
222
223    pub fn remove(mut self) -> V {
224        let ((_k, v), _) = unsafe { self.shard.remove(self.bucket) };
225        v.into_inner()
226    }
227
228    pub fn remove_entry(mut self) -> (K, V) {
229        let ((k, v), _) = unsafe { self.shard.remove(self.bucket) };
230        (k, v.into_inner())
231    }
232
233    pub fn replace_entry(self, value: V) -> (K, V) {
234        let (k, v) = mem::replace(
235            unsafe { self.bucket.as_mut() },
236            (self.key, SharedValue::new(value)),
237        );
238        (k, v.into_inner())
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use crate::DashMap;
245
246    use super::*;
247
248    #[test]
249    fn test_insert_entry_into_vacant() {
250        let map: DashMap<u32, u32> = DashMap::new();
251
252        let entry = map.entry(1);
253
254        assert!(matches!(entry, Entry::Vacant(_)));
255
256        let entry = entry.insert_entry(2);
257
258        assert_eq!(*entry.get(), 2);
259
260        drop(entry);
261
262        assert_eq!(*map.get(&1).unwrap(), 2);
263    }
264
265    #[test]
266    fn test_insert_entry_into_occupied() {
267        let map: DashMap<u32, u32> = DashMap::new();
268
269        map.insert(1, 1000);
270
271        let entry = map.entry(1);
272
273        assert!(matches!(&entry, Entry::Occupied(entry) if *entry.get() == 1000));
274
275        let entry = entry.insert_entry(2);
276
277        assert_eq!(*entry.get(), 2);
278
279        drop(entry);
280
281        assert_eq!(*map.get(&1).unwrap(), 2);
282    }
283}
dashmap/mapref/entry.rs

dashmap/mapref/
entry.rs
