1 | use crate::loom::sync::atomic::AtomicBool; |
2 | |
3 | use std::cell::UnsafeCell; |
4 | use std::marker::PhantomData; |
5 | use std::ops::{Deref, DerefMut}; |
6 | use std::sync::atomic::Ordering::SeqCst; |
7 | |
8 | pub(crate) struct TryLock<T> { |
9 | locked: AtomicBool, |
10 | data: UnsafeCell<T>, |
11 | } |
12 | |
13 | pub(crate) struct LockGuard<'a, T> { |
14 | lock: &'a TryLock<T>, |
15 | _p: PhantomData<std::rc::Rc<()>>, |
16 | } |
17 | |
18 | unsafe impl<T: Send> Send for TryLock<T> {} |
19 | unsafe impl<T: Send> Sync for TryLock<T> {} |
20 | |
21 | unsafe impl<T: Sync> Sync for LockGuard<'_, T> {} |
22 | |
23 | macro_rules! new { |
24 | ($data:ident) => { |
25 | TryLock { |
26 | locked: AtomicBool::new(false), |
27 | data: UnsafeCell::new($data), |
28 | } |
29 | }; |
30 | } |
31 | |
32 | impl<T> TryLock<T> { |
33 | #[cfg (not(loom))] |
34 | /// Create a new `TryLock` |
35 | pub(crate) const fn new(data: T) -> TryLock<T> { |
36 | new!(data) |
37 | } |
38 | |
39 | #[cfg (loom)] |
40 | /// Create a new `TryLock` |
41 | pub(crate) fn new(data: T) -> TryLock<T> { |
42 | new!(data) |
43 | } |
44 | |
45 | /// Attempt to acquire lock |
46 | pub(crate) fn try_lock(&self) -> Option<LockGuard<'_, T>> { |
47 | if self |
48 | .locked |
49 | .compare_exchange(false, true, SeqCst, SeqCst) |
50 | .is_err() |
51 | { |
52 | return None; |
53 | } |
54 | |
55 | Some(LockGuard { |
56 | lock: self, |
57 | _p: PhantomData, |
58 | }) |
59 | } |
60 | } |
61 | |
62 | impl<T> Deref for LockGuard<'_, T> { |
63 | type Target = T; |
64 | |
65 | fn deref(&self) -> &T { |
66 | unsafe { &*self.lock.data.get() } |
67 | } |
68 | } |
69 | |
70 | impl<T> DerefMut for LockGuard<'_, T> { |
71 | fn deref_mut(&mut self) -> &mut T { |
72 | unsafe { &mut *self.lock.data.get() } |
73 | } |
74 | } |
75 | |
76 | impl<T> Drop for LockGuard<'_, T> { |
77 | fn drop(&mut self) { |
78 | self.lock.locked.store(false, SeqCst); |
79 | } |
80 | } |
81 | |