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