Skip to main content

dark_std/sync/
map_index.rs

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