A better Rust ATProto crate
1use core::{
2 fmt::{Debug, Formatter},
3 marker::PhantomData,
4 ops::{Deref, DerefMut, Index, IndexMut},
5 slice::SliceIndex,
6};
7use std::collections::VecDeque;
8
9use maitake_sync::blocking::{RwLock, RwLockReadGuard, RwLockWriteGuard};
10use serde::{Deserialize, Serialize};
11
12use crate::{
13 Handle,
14 io::{FromEmbed, Read, Write},
15 new_handle,
16 storage::Storage,
17};
18
19#[cfg(not(feature = "alloc"))]
20pub type VecDeque<T> = heapless::Deque<T, 2>;
21
22pub struct LMap<K, V, S> {
23 handle: [u8; 4],
24 loaded: RwLock<VecDeque<(K, V)>>,
25 lock: RwLock<()>,
26 len: usize,
27 buf_size: usize,
28 storage: S,
29}
30
31impl<K, V, S> LMap<K, V, S> {
32 pub fn new_with(handle: Handle, buf_size: usize, storage: S) -> Self {
33 LMap {
34 handle: handle.0.to_be_bytes(),
35 loaded: RwLock::new(VecDeque::new()),
36 lock: RwLock::new(()),
37 len: 0,
38 buf_size,
39 storage,
40 }
41 }
42
43 pub fn new(storage: S) -> Self {
44 LMap::new_with(new_handle(), 2, storage)
45 }
46}
47
48impl<K, V, S> LMap<K, V, S>
49where
50 K: Ord,
51 S: Storage + Read + Write,
52{
53 pub fn len(&self) -> usize {
54 self.len
55 }
56
57 fn get_key(&self, index: &K) -> Vec<u8>
58 where
59 K: Serialize,
60 {
61 let mut key = postcard::to_allocvec(index).expect("Failed to serialize key");
62 key.extend_from_slice(&self.handle);
63 key
64 }
65
66 /// NOTE: Do not hold these references for an extended period of time.
67 /// This does NOT spin waiting for a write lock if it misses the cache, but
68 /// then it will not cache what it found. I consider this an acceptable failure
69 /// mode for the intended use case.
70 pub fn get_elem<'de, 's: 'de>(&'s self, key: &K) -> Option<Ref<'s, K, V>>
71 where
72 K: Serialize + Clone,
73 V: Deserialize<'de>,
74 S: 'de,
75 {
76 let loaded = self.loaded.read();
77 let mut found = None;
78 for (i, (idx, _)) in loaded.iter().enumerate() {
79 if idx == key {
80 found = Some(i);
81 }
82 }
83 if let Some(i) = found {
84 let guard = self.lock.read();
85 return loaded.get(i).map(|(k, v)| unsafe { Ref::new(guard, k, v) });
86 }
87 drop(loaded);
88
89 let key_bytes = self.get_key(key);
90 let reader = FromEmbed::new(self.storage.get::<S>(&key_bytes).ok()??);
91 let buffer = self.storage.buffer(&key_bytes);
92 let (elem, _buf) = postcard::from_eio::<V, _>((reader, buffer)).ok()?;
93 if let Some(mut loaded) = self.loaded.try_write() {
94 loaded.push_back((key.clone(), elem));
95 if loaded.len() > self.buf_size {
96 loaded.pop_front();
97 }
98 }
99 let read_guard = self.loaded.read();
100 let back = read_guard.back();
101 let guard = self.lock.read();
102 back.map(|(k, v)| unsafe { Ref::new(guard, k, v) })
103 }
104
105 pub fn get_elem_mut<'de, 's: 'de>(&'s mut self, index: usize) -> Option<RefMut<'s, K, V>>
106 where
107 K: Deserialize<'de>,
108 V: Deserialize<'de>,
109 S: 'de,
110 {
111 let loaded = self.loaded.get_mut();
112 let mut found = None;
113 for (i, (idx, _)) in loaded.iter().enumerate() {
114 if *idx == index {
115 found = Some(i);
116 }
117 }
118 if let Some(i) = found {
119 let guard = self.lock.write();
120 loaded
121 .get_mut(i)
122 .map(|(_, e)| unsafe { RefMut::new(guard, e) })
123 } else {
124 self.get_stored_mut(index)
125 }
126 }
127
128 fn get_stored_mut<'de, 's: 'de>(&'s mut self, index: usize) -> Option<RefMut<'s, K, V>>
129 where
130 K: Deserialize<'de>,
131 V: Deserialize<'de>,
132 S: 'de,
133 {
134 let key = self.get_key(index);
135 let reader = FromEmbed::new(self.storage.get::<S>(&key).ok()??);
136 let buffer = self.storage.buffer(&key);
137 let (elem, _buf) = postcard::from_eio::<T, _>((reader, buffer)).ok()?;
138 let loaded = self.loaded.get_mut();
139 loaded.push_back((index, elem));
140 if loaded.len() > self.buf_size {
141 loaded.pop_front();
142 }
143 let guard = self.lock.write();
144 let back = loaded.back_mut();
145 back.map(|(_, e)| unsafe { RefMut::new(guard, e) })
146 }
147
148 // pub fn get<I>(&self, index: I) -> Option<&<I as SliceIndex<[T]>>::Output>
149 // where
150 // I: SliceIndex<[T]>,
151 // {
152 // unimplemented!()
153 // }
154
155 // pub fn get_mut<I>(&mut self, index: I) -> Option<&mut <I as SliceIndex<[T]>>::Output>
156 // where
157 // I: SliceIndex<[T]>,
158 // {
159 // unimplemented!()
160 // }
161
162 pub fn insert(&mut self, index: K, elem: V) {
163 let key = self.get_key(index);
164 let mut writer = FromEmbed::new(self.storage.writer(&key).expect("failed to get writer"));
165 postcard::to_eio(&elem, &mut writer).expect("Failed to serialize element");
166 self.len += 1;
167 let loaded = self.loaded.get_mut();
168 loaded.push_back((index, elem));
169 self.storage.put::<_>(&key, writer.into_inner()).ok();
170 }
171
172 pub fn remove<'de, 's: 'de>(&'s mut self, index: &K) -> Option<V>
173 where
174 K: Deserialize<'de>,
175 V: Deserialize<'de>,
176 S: 'de,
177 {
178 if index >= self.len {
179 return None;
180 }
181 let mut found = None;
182 let mut loaded = self.loaded.write();
183 for (i, (idx, _)) in loaded.iter().enumerate() {
184 if *idx == index {
185 found = Some(i);
186 }
187 }
188 let key = self.get_key(index);
189 if let Some(i) = found {
190 self.storage.del(&key).ok();
191 loaded.remove(i).map(|(_, e)| e)
192 } else {
193 let reader = FromEmbed::new(self.storage.get::<S>(&key).ok()??);
194 let buffer = self.storage.buffer(&key);
195 let (elem, _buf) = postcard::from_eio::<T, _>((reader, buffer)).ok()?;
196 self.storage.del(&key).ok();
197 Some(elem)
198 }
199 }
200}
201
202pub struct Ref<'a, K, V> {
203 _guard: RwLockReadGuard<'a, ()>,
204 k: *const K,
205 v: *const V,
206}
207
208unsafe impl<'a, K: Ord + Sync, V: Sync> Send for Ref<'a, K, V> {}
209unsafe impl<'a, K: Ord + Sync, V: Sync> Sync for Ref<'a, K, V> {}
210
211impl<'a, K: Ord, V> Ref<'a, K, V> {
212 pub(crate) unsafe fn new(guard: RwLockReadGuard<'a, ()>, k: *const K, v: *const V) -> Self {
213 Self {
214 _guard: guard,
215 k,
216 v,
217 }
218 }
219
220 pub fn key(&self) -> &K {
221 self.pair().0
222 }
223
224 pub fn value(&self) -> &V {
225 self.pair().1
226 }
227
228 pub fn pair(&self) -> (&K, &V) {
229 unsafe { (&*self.k, &*self.v) }
230 }
231
232 pub fn map<F, T>(self, f: F) -> MappedRef<'a, K, V, T>
233 where
234 F: FnOnce(&V) -> &T,
235 {
236 MappedRef {
237 _guard: self._guard,
238 k: self.k,
239 v: f(unsafe { &*self.v }),
240 _marker: PhantomData,
241 }
242 }
243
244 pub fn try_map<F, T>(self, f: F) -> Result<MappedRef<'a, K, V, T>, Self>
245 where
246 F: FnOnce(&V) -> Option<&T>,
247 {
248 if let Some(v) = f(unsafe { &*self.v }) {
249 Ok(MappedRef {
250 _guard: self._guard,
251 k: self.k,
252 v,
253 _marker: PhantomData,
254 })
255 } else {
256 Err(self)
257 }
258 }
259}
260
261impl<'a, K: Ord + Debug, V: Debug> Debug for Ref<'a, K, V> {
262 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
263 f.debug_struct("Ref")
264 .field("k", &self.k)
265 .field("v", &self.v)
266 .finish()
267 }
268}
269
270impl<'a, K: Ord, V> Deref for Ref<'a, K, V> {
271 type Target = V;
272
273 fn deref(&self) -> &V {
274 self.value()
275 }
276}
277
278pub struct RefMut<'a, K, V> {
279 guard: RwLockWriteGuard<'a, ()>,
280 k: *const K,
281 v: *mut V,
282}
283
284unsafe impl<'a, K: Ord + Sync, V: Sync> Send for RefMut<'a, K, V> {}
285unsafe impl<'a, K: Ord + Sync, V: Sync> Sync for RefMut<'a, K, V> {}
286
287impl<'a, K: Ord, V> RefMut<'a, K, V> {
288 pub(crate) unsafe fn new(guard: RwLockWriteGuard<'a, ()>, k: *const K, v: *mut V) -> Self {
289 Self { guard, k, v }
290 }
291
292 pub fn key(&self) -> &K {
293 self.pair().0
294 }
295
296 pub fn value(&self) -> &V {
297 self.pair().1
298 }
299
300 pub fn value_mut(&mut self) -> &mut V {
301 self.pair_mut().1
302 }
303
304 pub fn pair(&self) -> (&K, &V) {
305 unsafe { (&*self.k, &*self.v) }
306 }
307
308 pub fn pair_mut(&mut self) -> (&K, &mut V) {
309 unsafe { (&*self.k, &mut *self.v) }
310 }
311
312 pub fn map<F, T>(self, f: F) -> MappedRefMut<'a, K, V, T>
313 where
314 F: FnOnce(&mut V) -> &mut T,
315 {
316 MappedRefMut {
317 _guard: self.guard,
318 k: self.k,
319 v: f(unsafe { &mut *self.v }),
320 _marker: PhantomData,
321 }
322 }
323
324 pub fn try_map<F, T>(self, f: F) -> Result<MappedRefMut<'a, K, V, T>, Self>
325 where
326 F: FnOnce(&mut V) -> Option<&mut T>,
327 {
328 let v = match f(unsafe { &mut *(self.v as *mut _) }) {
329 Some(v) => v,
330 None => return Err(self),
331 };
332 let guard = self.guard;
333 let k = self.k;
334 Ok(MappedRefMut {
335 _guard: guard,
336 k,
337 v,
338 _marker: PhantomData,
339 })
340 }
341}
342
343impl<'a, K: Ord + Debug, V: Debug> Debug for RefMut<'a, K, V> {
344 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
345 f.debug_struct("RefMut")
346 .field("k", &self.k)
347 .field("v", &self.v)
348 .finish()
349 }
350}
351
352impl<'a, K: Ord, V> Deref for RefMut<'a, K, V> {
353 type Target = V;
354
355 fn deref(&self) -> &V {
356 self.value()
357 }
358}
359
360impl<'a, K: Ord, V> DerefMut for RefMut<'a, K, V> {
361 fn deref_mut(&mut self) -> &mut V {
362 self.value_mut()
363 }
364}
365
366pub struct MappedRef<'a, K, V, T> {
367 _guard: RwLockReadGuard<'a, ()>,
368 k: *const K,
369 v: *const T,
370 _marker: PhantomData<V>,
371}
372
373impl<'a, K: Ord, V, T> MappedRef<'a, K, V, T> {
374 pub fn key(&self) -> &K {
375 self.pair().0
376 }
377
378 pub fn value(&self) -> &T {
379 self.pair().1
380 }
381
382 pub fn pair(&self) -> (&K, &T) {
383 unsafe { (&*self.k, &*self.v) }
384 }
385
386 pub fn map<F, T2>(self, f: F) -> MappedRef<'a, K, V, T2>
387 where
388 F: FnOnce(&T) -> &T2,
389 {
390 MappedRef {
391 _guard: self._guard,
392 k: self.k,
393 v: f(unsafe { &*self.v }),
394 _marker: PhantomData,
395 }
396 }
397
398 pub fn try_map<F, T2>(self, f: F) -> Result<MappedRef<'a, K, V, T2>, Self>
399 where
400 F: FnOnce(&T) -> Option<&T2>,
401 {
402 let v = match f(unsafe { &*self.v }) {
403 Some(v) => v,
404 None => return Err(self),
405 };
406 let guard = self._guard;
407 Ok(MappedRef {
408 _guard: guard,
409 k: self.k,
410 v,
411 _marker: PhantomData,
412 })
413 }
414}
415
416impl<'a, K: Ord + Debug, V, T: Debug> Debug for MappedRef<'a, K, V, T> {
417 fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
418 f.debug_struct("MappedRef")
419 .field("k", &self.k)
420 .field("v", &self.v)
421 .finish()
422 }
423}
424
425impl<'a, K: Ord, V, T> Deref for MappedRef<'a, K, V, T> {
426 type Target = T;
427
428 fn deref(&self) -> &T {
429 self.value()
430 }
431}
432
433impl<'a, K: Ord, V, T: core::fmt::Display> core::fmt::Display for MappedRef<'a, K, V, T> {
434 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
435 core::fmt::Display::fmt(self.value(), f)
436 }
437}
438
439impl<'a, K: Ord, V, T: AsRef<TDeref>, TDeref: ?Sized> AsRef<TDeref> for MappedRef<'a, K, V, T> {
440 fn as_ref(&self) -> &TDeref {
441 self.value().as_ref()
442 }
443}
444
445pub struct MappedRefMut<'a, K, V, T> {
446 _guard: RwLockWriteGuard<'a, ()>,
447 k: *const K,
448 v: *mut T,
449 _marker: PhantomData<V>,
450}
451
452impl<'a, K: Ord, V, T> MappedRefMut<'a, K, V, T> {
453 pub fn key(&self) -> &K {
454 self.pair().0
455 }
456
457 pub fn value(&self) -> &T {
458 self.pair().1
459 }
460
461 pub fn value_mut(&mut self) -> &mut T {
462 self.pair_mut().1
463 }
464
465 pub fn pair(&self) -> (&K, &T) {
466 unsafe { (&*self.k, &*self.v) }
467 }
468
469 pub fn pair_mut(&mut self) -> (&K, &mut T) {
470 unsafe { (&*self.k, &mut *self.v) }
471 }
472
473 pub fn map<F, T2>(self, f: F) -> MappedRefMut<'a, K, V, T2>
474 where
475 F: FnOnce(&mut T) -> &mut T2,
476 {
477 MappedRefMut {
478 _guard: self._guard,
479 k: self.k,
480 v: f(unsafe { &mut *self.v }),
481 _marker: PhantomData,
482 }
483 }
484
485 pub fn try_map<F, T2>(self, f: F) -> Result<MappedRefMut<'a, K, V, T2>, Self>
486 where
487 F: FnOnce(&mut T) -> Option<&mut T2>,
488 {
489 let v = match f(unsafe { &mut *(self.v as *mut _) }) {
490 Some(v) => v,
491 None => return Err(self),
492 };
493 let guard = self._guard;
494 let k = self.k;
495 Ok(MappedRefMut {
496 _guard: guard,
497 k,
498 v,
499 _marker: PhantomData,
500 })
501 }
502}
503
504impl<'a, K: Ord + Debug, V, T: Debug> Debug for MappedRefMut<'a, K, V, T> {
505 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
506 f.debug_struct("MappedRefMut")
507 .field("k", &self.k)
508 .field("v", &self.v)
509 .finish()
510 }
511}
512
513impl<'a, K: Ord, V, T> Deref for MappedRefMut<'a, K, V, T> {
514 type Target = T;
515
516 fn deref(&self) -> &T {
517 self.value()
518 }
519}
520
521impl<'a, K: Ord, V, T> DerefMut for MappedRefMut<'a, K, V, T> {
522 fn deref_mut(&mut self) -> &mut T {
523 self.value_mut()
524 }
525}