1use parking_lot::{ReentrantMutex, ReentrantMutexGuard};
2use serde::{Deserializer, Serialize, Serializer};
3use std::borrow::Borrow;
4use std::cell::UnsafeCell;
5use std::collections::{btree_map::IntoIter as MapIntoIter, btree_map::Iter as MapIter, BTreeMap};
6use std::fmt::{Debug, Display, Formatter};
7use std::hash::Hash;
8use std::ops::{Deref, DerefMut};
9use std::sync::Arc;
10
11pub struct SyncBtreeMap<K: Eq + Hash, V> {
13 locks: UnsafeCell<BTreeMap<K, ReentrantMutex<()>>>,
14 dirty: UnsafeCell<BTreeMap<K, V>>,
15 lock: ReentrantMutex<()>,
16}
17
18unsafe impl<K: Eq + Hash, V> Send for SyncBtreeMap<K, V> {}
20
21unsafe impl<K: Eq + Hash, V> Sync for SyncBtreeMap<K, V> {}
23
24impl<K, V> std::ops::Index<&K> for SyncBtreeMap<K, V>
25where
26 K: Eq + Hash + Ord,
27{
28 type Output = V;
29
30 fn index(&self, index: &K) -> &Self::Output {
31 unsafe { &(&*self.dirty.get())[index] }
32 }
33}
34
35impl<K: Eq + Hash, V> SyncBtreeMap<K, V>
36where
37 K: Eq + Hash,
38{
39 pub fn new_arc() -> Arc<Self> {
40 Arc::new(Self::new())
41 }
42
43 pub fn new() -> Self {
44 Self {
45 locks: UnsafeCell::new(BTreeMap::new()),
46 dirty: UnsafeCell::new(BTreeMap::new()),
47 lock: Default::default(),
48 }
49 }
50
51 pub fn with_capacity(_capacity: usize) -> Self {
52 Self::new()
53 }
54
55 pub fn with_map(map: BTreeMap<K, V>) -> Self {
56 Self {
57 locks: UnsafeCell::new(BTreeMap::new()),
58 dirty: UnsafeCell::new(map),
59 lock: Default::default(),
60 }
61 }
62
63 pub fn insert(&self, k: K, v: V) -> Option<V>
64 where
65 K: Ord,
66 {
67 let g = self.lock.lock();
68 let m = unsafe { &mut *self.dirty.get() };
69 let r = m.insert(k, v);
70 drop(g);
71 r
72 }
73
74 pub fn insert_mut(&mut self, k: K, v: V) -> Option<V>
75 where
76 K: Ord,
77 {
78 let m = unsafe { &mut *self.dirty.get() };
79 m.insert(k, v)
80 }
81
82 pub fn remove(&self, k: &K) -> Option<V>
83 where
84 K: Ord,
85 {
86 let g = self.lock.lock();
87 let m = unsafe { &mut *self.dirty.get() };
88 let r = m.remove(k);
89 drop(g);
90 r
91 }
92
93 pub fn remove_mut(&mut self, k: &K) -> Option<V>
94 where
95 K: Ord,
96 {
97 let m = unsafe { &mut *self.dirty.get() };
98 m.remove(k)
99 }
100
101 pub fn len(&self) -> usize {
102 unsafe { (&*self.dirty.get()).len() }
103 }
104
105 pub fn is_empty(&self) -> bool {
106 unsafe { (&*self.dirty.get()).is_empty() }
107 }
108
109 pub fn clear(&self)
110 where
111 K: Eq + Hash,
112 {
113 let g = self.lock.lock();
114 let m = unsafe { &mut *self.dirty.get() };
115 m.clear();
116 drop(g);
117 }
118
119 pub fn clear_mut(&mut self)
120 where
121 K: Eq + Hash,
122 {
123 let m = unsafe { &mut *self.dirty.get() };
124 m.clear();
125 }
126
127 pub fn shrink_to_fit(&self) {}
128
129 pub fn shrink_to_fit_mut(&mut self) {}
130
131 pub fn from(map: BTreeMap<K, V>) -> Self
132 where
133 K: Eq + Hash,
134 {
135 let s = Self::with_map(map);
136 s
137 }
138
139 #[inline]
159 pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
160 where
161 K: Borrow<Q> + Ord,
162 Q: Hash + Eq + Ord,
163 {
164 unsafe { (&*self.dirty.get()).get(k) }
165 }
166
167 #[inline]
168 pub fn get_mut(&self, k: &K) -> Option<BtreeMapRefMut<'_, K, V>>
169 where
170 K: Hash + Eq + Clone + Ord,
171 {
172 let get_mut_lock = self.lock.lock();
173 let m = unsafe { &mut *self.locks.get() };
174 if m.contains_key(k) == false {
175 let g = ReentrantMutex::new(());
176 m.insert(k.clone(), g);
177 }
178 let g = m.get(k).unwrap();
179 let m = unsafe { &mut *self.dirty.get() };
180 let v = BtreeMapRefMut {
181 k: unsafe { std::mem::transmute(&k) },
182 m: self,
183 _g: g.lock(),
184 value: m.get_mut(k)?,
185 };
186 drop(get_mut_lock);
187 Some(v)
188 }
189
190 #[inline]
191 pub fn contains_key(&self, x: &K) -> bool
192 where
193 K: PartialEq + Ord,
194 {
195 let m = unsafe { &mut *self.dirty.get() };
196 m.contains_key(x)
197 }
198
199 pub fn iter(&self) -> MapIter<'_, K, V> {
200 unsafe { (&*self.dirty.get()).iter() }
201 }
202
203 pub fn iter_mut(&self) -> BtreeIterMut<'_, K, V> {
204 let m = unsafe { &mut *self.dirty.get() };
205 return BtreeIterMut {
206 _g: self.lock.lock(),
207 inner: m.iter_mut(),
208 };
209 }
210
211 pub fn into_iter(self) -> MapIntoIter<K, V> {
212 self.dirty.into_inner().into_iter()
213 }
214
215 pub fn dirty_ref(&self) -> &BTreeMap<K, V> {
216 unsafe { &*self.dirty.get() }
217 }
218
219 pub fn into_inner(self) -> BTreeMap<K, V> {
220 self.dirty.into_inner()
221 }
222}
223
224pub struct BtreeMapRefMut<'a, K: Eq + Hash + Ord, V> {
225 k: &'a K,
226 m: &'a SyncBtreeMap<K, V>,
227 _g: ReentrantMutexGuard<'a, ()>,
228 value: &'a mut V,
229}
230
231impl<'a, K: Eq + Hash + Ord, V> Drop for BtreeMapRefMut<'a, K, V> {
232 fn drop(&mut self) {
233 let m = unsafe { &mut *self.m.locks.get() };
234 _ = m.remove(self.k);
235 }
236}
237
238impl<'a, K: Eq + Hash + Ord, V> Deref for BtreeMapRefMut<'_, K, V> {
239 type Target = V;
240
241 fn deref(&self) -> &Self::Target {
242 self.value
243 }
244}
245
246impl<'a, K: Eq + Hash + Ord, V> DerefMut for BtreeMapRefMut<'_, K, V> {
247 fn deref_mut(&mut self) -> &mut Self::Target {
248 self.value
249 }
250}
251
252impl<'a, K: Eq + Hash + Ord, V> Debug for BtreeMapRefMut<'_, K, V>
253where
254 V: Debug,
255{
256 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
257 self.value.fmt(f)
258 }
259}
260
261impl<'a, K: Eq + Hash + Ord, V> Display for BtreeMapRefMut<'_, K, V>
262where
263 V: Display,
264{
265 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
266 self.value.fmt(f)
267 }
268}
269
270impl<'a, K: Eq + Hash + Ord, V> PartialEq<Self> for BtreeMapRefMut<'_, K, V>
271where
272 V: Eq,
273{
274 fn eq(&self, other: &Self) -> bool {
275 self.value.eq(&other.value)
276 }
277}
278
279impl<'a, K: Eq + Hash + Ord, V> Eq for BtreeMapRefMut<'_, K, V> where V: Eq {}
280
281pub struct BtreeIterMut<'a, K, V> {
282 _g: ReentrantMutexGuard<'a, ()>,
283 inner: std::collections::btree_map::IterMut<'a, K, V>,
284}
285
286impl<'a, K, V> Deref for BtreeIterMut<'a, K, V> {
287 type Target = std::collections::btree_map::IterMut<'a, K, V>;
288
289 fn deref(&self) -> &Self::Target {
290 &self.inner
291 }
292}
293
294impl<'a, K, V> DerefMut for BtreeIterMut<'a, K, V> {
295 fn deref_mut(&mut self) -> &mut Self::Target {
296 &mut self.inner
297 }
298}
299
300impl<'a, K, V> Iterator for BtreeIterMut<'a, K, V> {
301 type Item = (&'a K, &'a mut V);
302
303 fn next(&mut self) -> Option<Self::Item> {
304 self.inner.next()
305 }
306}
307
308impl<'a, K: Eq + Hash, V> IntoIterator for &'a SyncBtreeMap<K, V> {
309 type Item = (&'a K, &'a V);
310 type IntoIter = MapIter<'a, K, V>;
311
312 fn into_iter(self) -> Self::IntoIter {
313 self.iter()
314 }
315}
316
317impl<K: Eq + Hash, V> IntoIterator for SyncBtreeMap<K, V> {
318 type Item = (K, V);
319 type IntoIter = MapIntoIter<K, V>;
320
321 fn into_iter(self) -> Self::IntoIter {
322 self.into_iter()
323 }
324}
325
326impl<K: Eq + Hash, V> From<BTreeMap<K, V>> for SyncBtreeMap<K, V> {
327 fn from(arg: BTreeMap<K, V>) -> Self {
328 Self::from(arg)
329 }
330}
331
332impl<K: Eq + Hash, V> serde::Serialize for SyncBtreeMap<K, V>
333where
334 K: Eq + Hash + Serialize + Ord,
335 V: Serialize,
336{
337 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
338 where
339 S: Serializer,
340 {
341 self.dirty_ref().serialize(serializer)
342 }
343}
344
345impl<'de, K, V> serde::Deserialize<'de> for SyncBtreeMap<K, V>
346where
347 K: Eq + Hash + Ord + serde::Deserialize<'de>,
348 V: serde::Deserialize<'de>,
349{
350 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
351 where
352 D: Deserializer<'de>,
353 {
354 let m = BTreeMap::deserialize(deserializer)?;
355 Ok(Self::from(m))
356 }
357}
358
359impl<K: Eq + Hash, V> Debug for SyncBtreeMap<K, V>
360where
361 K: Eq + Hash + Debug,
362 V: Debug,
363{
364 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
365 self.dirty_ref().fmt(f)
366 }
367}
368
369impl<K: Eq + Hash, V> Display for SyncBtreeMap<K, V>
370where
371 K: Eq + Hash + Display,
372 V: Display,
373{
374 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
375 use std::fmt::Pointer;
376 self.dirty_ref().fmt(f)
377 }
378}
379
380pub struct BtreeIter<'a, K, V> {
381 inner: MapIter<'a, K, *const V>,
382}
383
384impl<'a, K, V> Iterator for BtreeIter<'a, K, V> {
385 type Item = (&'a K, &'a V);
386
387 fn next(&mut self) -> Option<Self::Item> {
388 match self.inner.next() {
389 None => None,
390 Some((k, v)) => Some((k, unsafe { v.as_ref().unwrap() })),
391 }
392 }
393}
394
395impl<K: Clone + Eq + Hash, V: Clone> Clone for SyncBtreeMap<K, V> {
396 fn clone(&self) -> Self {
397 let c = (*self.dirty_ref()).clone();
398 SyncBtreeMap::from(c)
399 }
400}
401
402impl<K: Eq + Hash, V> Default for SyncBtreeMap<K, V> {
403 fn default() -> Self {
404 SyncBtreeMap::new()
405 }
406}