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
17pub type IndexMapGet<'a, V> = ReadGuard<'a, V>;
19
20pub 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
72pub type HashMapRefMut<'a, K, V> = IndexMapRefMut<'a, K, V>;
75
76pub 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
97pub 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
125pub 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
147unsafe 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 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 #[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 #[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
400impl<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}