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