Skip to main content

dark_std/sync/
map_hash.rs

1use parking_lot::Mutex;
2use serde::{Deserializer, Serialize, Serializer};
3use std::borrow::Borrow;
4use std::cell::UnsafeCell;
5use std::collections::{
6    hash_map::IntoIter as MapIntoIter, hash_map::Iter as MapIter,
7    hash_map::IterMut as MapIterMut, HashMap as Map,
8};
9use std::fmt::{Debug, Display, Formatter};
10use std::hash::Hash;
11use std::marker::PhantomData;
12use std::ops::{Deref, DerefMut};
13use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
14use std::sync::Arc;
15
16use super::{ReadGuard, ReadMapGuard, WriteGuard, WriteLock};
17
18/// Read guard returned by [`SyncHashMap::get`].
19pub type HashMapGet<'a, V> = ReadGuard<'a, V>;
20
21/// Write guard returned by [`SyncHashMap::get_mut`].
22pub struct HashMapRefMut<'a, K, V> {
23    inner: WriteGuard<'a, V>,
24    _k: PhantomData<&'a K>,
25}
26
27impl<'a, K, V> HashMapRefMut<'a, K, V> {
28    #[inline]
29    pub(crate) fn new(inner: WriteGuard<'a, V>) -> Self {
30        HashMapRefMut {
31            inner,
32            _k: PhantomData,
33        }
34    }
35}
36
37impl<'a, K, V> Deref for HashMapRefMut<'a, K, V> {
38    type Target = V;
39
40    #[inline]
41    fn deref(&self) -> &Self::Target {
42        &self.inner
43    }
44}
45
46impl<'a, K, V> DerefMut for HashMapRefMut<'a, K, V> {
47    #[inline]
48    fn deref_mut(&mut self) -> &mut Self::Target {
49        &mut self.inner
50    }
51}
52
53impl<'a, K, V: Debug> Debug for HashMapRefMut<'a, K, V> {
54    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
55        Debug::fmt(&*self.inner, f)
56    }
57}
58
59impl<'a, K, V: Display> Display for HashMapRefMut<'a, K, V> {
60    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
61        Display::fmt(&*self.inner, f)
62    }
63}
64
65impl<'a, K, V: PartialEq> PartialEq for HashMapRefMut<'a, K, V> {
66    fn eq(&self, other: &Self) -> bool {
67        *self.inner == *other.inner
68    }
69}
70
71impl<'a, K, V: Eq> Eq for HashMapRefMut<'a, K, V> {}
72
73/// Read iterator returned by [`SyncHashMap::iter`].
74pub struct HashMapIter<'a, K, V> {
75    count: &'a AtomicUsize,
76    inner: MapIter<'a, K, V>,
77    _not_send: PhantomData<*const ()>,
78}
79
80impl<'a, K, V> Drop for HashMapIter<'a, K, V> {
81    fn drop(&mut self) {
82        self.count.fetch_sub(1, Ordering::Release);
83    }
84}
85
86impl<'a, K, V> Iterator for HashMapIter<'a, K, V> {
87    type Item = (&'a K, &'a V);
88
89    fn next(&mut self) -> Option<Self::Item> {
90        self.inner.next()
91    }
92}
93
94/// Write iterator returned by [`SyncHashMap::iter_mut`].
95pub struct HashMapIterMut<'a, K, V> {
96    _w: WriteLock<'a>,
97    inner: MapIterMut<'a, K, V>,
98}
99
100impl<'a, K, V> Iterator for HashMapIterMut<'a, K, V> {
101    type Item = (&'a K, &'a mut V);
102
103    fn next(&mut self) -> Option<Self::Item> {
104        self.inner.next()
105    }
106}
107
108/// this sync map used to many reader,writer less.space-for-time strategy
109///
110/// Reads are lock-free: `get`/`iter`/`dirty_ref`/`len`/`contains_key` only
111/// register a reader slot with an atomic counter and then read the map without
112/// any lock (readers never block each other and never touch a lock word).
113/// Writes take a mutex, raise a `writing` flag and wait until all in-flight
114/// readers are gone before mutating the map in place — O(1), no whole-container
115/// copy and no `Clone` requirement on `K`/`V`.
116///
117/// # Deadlock note
118/// A read guard makes writers wait until it is dropped. Do not call a write
119/// method while a read/write guard is alive in the same scope: drop the guard
120/// first (e.g. `drop(g)` before `insert`/`remove`/`get_mut`), otherwise the
121/// writer waits for its own guard and deadlocks.
122pub struct SyncHashMap<K: Eq + Hash, V> {
123    dirty: UnsafeCell<Map<K, V>>,
124    write: Mutex<()>,
125    id: usize,
126    writing: AtomicBool,
127    registry: Mutex<Vec<std::boxed::Box<AtomicUsize>>>,
128}
129
130// SAFETY: all writers hold `write` and wait for `readers` to drain before
131// touching `dirty`; readers either see a consistent snapshot or retry while a
132// writer is active, so concurrent access to `dirty` is race-free.
133unsafe impl<K: Eq + Hash, V: Send> Send for SyncHashMap<K, V> {}
134unsafe impl<K: Eq + Hash, V: Sync> Sync for SyncHashMap<K, V> {}
135
136impl<K, V> SyncHashMap<K, V>
137where
138    K: Eq + Hash,
139{
140    #[inline]
141    fn begin_read(&self) -> &AtomicUsize {
142        // The counter lives in thread-local storage: concurrent readers only
143        // touch their own cache line and never contend with each other. SeqCst
144        // closes the store-buffering window with the writer's all-zero scan.
145        let count = super::reader_count_for(self.id, &self.registry);
146        loop {
147            count.fetch_add(1, Ordering::SeqCst);
148            if !self.writing.load(Ordering::SeqCst) {
149                return count;
150            }
151            count.fetch_sub(1, Ordering::SeqCst);
152            std::thread::yield_now();
153        }
154    }
155
156    #[inline]
157    fn begin_write(&self) -> WriteLock<'_> {
158        let lock = self.write.lock();
159        self.writing.store(true, Ordering::SeqCst);
160        loop {
161            let registry = self.registry.lock();
162            let all_zero = registry.iter().all(|c| c.load(Ordering::SeqCst) == 0);
163            if all_zero {
164                break;
165            }
166            drop(registry);
167            std::thread::yield_now();
168        }
169        WriteLock::new(lock, &self.writing)
170    }
171
172    pub fn new_arc() -> Arc<Self> {
173        Arc::new(Self::new())
174    }
175
176    pub fn new() -> Self {
177        Self {
178            dirty: UnsafeCell::new(Map::new()),
179            write: Mutex::new(()),
180            id: super::CONTAINER_ID.fetch_add(1, Ordering::Relaxed),
181            writing: AtomicBool::new(false),
182            registry: Mutex::new(Vec::new()),
183        }
184    }
185
186    pub fn with_capacity(capacity: usize) -> Self {
187        Self {
188            dirty: UnsafeCell::new(Map::with_capacity(capacity)),
189            write: Mutex::new(()),
190            id: super::CONTAINER_ID.fetch_add(1, Ordering::Relaxed),
191            writing: AtomicBool::new(false),
192            registry: Mutex::new(Vec::new()),
193        }
194    }
195
196    pub fn with_map(map: Map<K, V>) -> Self {
197        Self {
198            dirty: UnsafeCell::new(map),
199            write: Mutex::new(()),
200            id: super::CONTAINER_ID.fetch_add(1, Ordering::Relaxed),
201            writing: AtomicBool::new(false),
202            registry: Mutex::new(Vec::new()),
203        }
204    }
205
206    pub fn insert(&self, k: K, v: V) -> Option<V> {
207        let _w = self.begin_write();
208        unsafe { &mut *self.dirty.get() }.insert(k, v)
209    }
210
211    pub fn insert_mut(&mut self, k: K, v: V) -> Option<V> {
212        unsafe { &mut *self.dirty.get() }.insert(k, v)
213    }
214
215    pub fn remove(&self, k: &K) -> Option<V> {
216        let _w = self.begin_write();
217        unsafe { &mut *self.dirty.get() }.remove(k)
218    }
219
220    pub fn remove_mut(&mut self, k: &K) -> Option<V> {
221        unsafe { &mut *self.dirty.get() }.remove(k)
222    }
223
224    pub fn len(&self) -> usize {
225        let count = self.begin_read();
226        let n = unsafe { &*self.dirty.get() }.len();
227        count.fetch_sub(1, Ordering::Release);
228        n
229    }
230
231    pub fn is_empty(&self) -> bool {
232        let count = self.begin_read();
233        let b = unsafe { &*self.dirty.get() }.is_empty();
234        count.fetch_sub(1, Ordering::Release);
235        b
236    }
237
238    pub fn clear(&self) {
239        let _w = self.begin_write();
240        unsafe { &mut *self.dirty.get() }.clear();
241    }
242
243    pub fn clear_mut(&mut self) {
244        unsafe { &mut *self.dirty.get() }.clear();
245    }
246
247    pub fn shrink_to_fit(&self) {
248        let _w = self.begin_write();
249        unsafe { &mut *self.dirty.get() }.shrink_to_fit();
250    }
251
252    pub fn shrink_to_fit_mut(&mut self) {
253        unsafe { &mut *self.dirty.get() }.shrink_to_fit()
254    }
255
256    pub fn from(map: Map<K, V>) -> Self
257    where
258        K: Eq + Hash,
259    {
260        Self::with_map(map)
261    }
262
263    /// Returns a read-guarded reference to the value corresponding to the key.
264    ///
265    /// The key may be any borrowed form of the map's key type, but
266    /// [`Hash`] and [`Eq`] on the borrowed form *must* match those for
267    /// the key type.
268    ///
269    /// The read is lock-free: it only registers a reader slot, so concurrent
270    /// reads never block each other and never take a lock. Writers wait for
271    /// the returned guard to be dropped before mutating the map.
272    ///
273    /// # Examples
274    ///
275    /// ```
276    /// use dark_std::sync::{SyncHashMap};
277    ///
278    /// let mut map = SyncHashMap::new();
279    /// map.insert_mut(1, "a");
280    /// assert_eq!(*map.get(&1).unwrap(), "a");
281    /// assert_eq!(map.get(&2).is_none(), true);
282    /// ```
283    #[inline]
284    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<HashMapGet<'_, V>>
285    where
286        K: Borrow<Q>,
287        Q: Hash + Eq,
288    {
289        let count = self.begin_read();
290        let m = unsafe { &*self.dirty.get() };
291        match m.get(k) {
292            Some(v) => Some(ReadGuard::new(count, v)),
293            None => {
294                count.fetch_sub(1, Ordering::Release);
295                None
296            }
297        }
298    }
299
300    /// Returns a write-guarded mutable reference to the value of the key.
301    ///
302    /// The guard holds the writer lock (writers are mutually exclusive and
303    /// wait for in-flight readers) until it is dropped, so the mutable
304    /// reference can never race with concurrent readers or writers. Drop it
305    /// before calling another method from the same scope.
306    #[inline]
307    pub fn get_mut(&self, k: &K) -> Option<HashMapRefMut<'_, K, V>> {
308        let w = self.begin_write();
309        let m = unsafe { &mut *self.dirty.get() };
310        match m.get_mut(k) {
311            Some(v) => Some(HashMapRefMut::new(WriteGuard::new(w, v))),
312            None => None,
313        }
314    }
315
316    #[inline]
317    pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
318    where
319        K: Borrow<Q>,
320        Q: Hash + Eq,
321    {
322        let count = self.begin_read();
323        let b = unsafe { &*self.dirty.get() }.contains_key(k);
324        count.fetch_sub(1, Ordering::Release);
325        b
326    }
327
328    pub fn iter(&self) -> HashMapIter<'_, K, V> {
329        let count = self.begin_read();
330        let m = unsafe { &*self.dirty.get() };
331        HashMapIter {
332            count,
333            inner: m.iter(),
334            _not_send: PhantomData,
335        }
336    }
337
338    pub fn iter_mut(&self) -> HashMapIterMut<'_, K, V> {
339        let w = self.begin_write();
340        let m = unsafe { &mut *self.dirty.get() };
341        HashMapIterMut {
342            _w: w,
343            inner: m.iter_mut(),
344        }
345    }
346
347    pub fn into_iter(self) -> MapIntoIter<K, V> {
348        self.into_inner().into_iter()
349    }
350
351    pub fn dirty_ref(&self) -> ReadMapGuard<'_, Map<K, V>> {
352        let count = self.begin_read();
353        let m = unsafe { &*self.dirty.get() };
354        ReadMapGuard::new(count, m)
355    }
356
357    pub fn into_inner(self) -> Map<K, V> {
358        self.dirty.into_inner()
359    }
360}
361
362impl<K, V> IntoIterator for SyncHashMap<K, V>
363where
364    K: Eq + Hash,
365{
366    type Item = (K, V);
367    type IntoIter = MapIntoIter<K, V>;
368
369    fn into_iter(self) -> Self::IntoIter {
370        self.into_iter()
371    }
372}
373
374impl<'a, K: Eq + Hash, V> IntoIterator for &'a SyncHashMap<K, V> {
375    type Item = (&'a K, &'a V);
376    type IntoIter = HashMapIter<'a, K, V>;
377
378    fn into_iter(self) -> Self::IntoIter {
379        self.iter()
380    }
381}
382
383impl<K: Eq + Hash, V> From<Map<K, V>> for SyncHashMap<K, V> {
384    fn from(arg: Map<K, V>) -> Self {
385        Self::from(arg)
386    }
387}
388
389impl<K, V> serde::Serialize for SyncHashMap<K, V>
390where
391    K: Eq + Hash + Serialize,
392    V: Serialize,
393{
394    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
395    where
396        S: Serializer,
397    {
398        self.dirty_ref().serialize(serializer)
399    }
400}
401
402impl<'de, K, V> serde::Deserialize<'de> for SyncHashMap<K, V>
403where
404    K: Eq + Hash + serde::Deserialize<'de>,
405    V: serde::Deserialize<'de>,
406{
407    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
408    where
409        D: Deserializer<'de>,
410    {
411        let m = Map::deserialize(deserializer)?;
412        Ok(Self::from(m))
413    }
414}
415
416impl<K, V> Debug for SyncHashMap<K, V>
417where
418    K: Eq + Hash + Debug,
419    V: Debug,
420{
421    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
422        Debug::fmt(&*self.dirty_ref(), f)
423    }
424}
425
426impl<K, V> Display for SyncHashMap<K, V>
427where
428    K: Eq + Hash + Debug,
429    V: Debug,
430{
431    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
432        Debug::fmt(&*self.dirty_ref(), f)
433    }
434}
435
436impl<K: Clone + Eq + Hash, V: Clone> Clone for SyncHashMap<K, V> {
437    fn clone(&self) -> Self {
438        let c = (*self.dirty_ref()).clone();
439        SyncHashMap::from(c)
440    }
441}
442
443impl<K: Eq + Hash, V> Default for SyncHashMap<K, V> {
444    fn default() -> Self {
445        SyncHashMap::new()
446    }
447}