1use alloc::sync::{Arc, Weak};
2use core::{
3 cell::UnsafeCell,
4 ops::{Deref, DerefMut},
5 sync::atomic::{AtomicBool, Ordering},
6};
7
8use crate::custom_type;
9
10custom_type!(PId, usize, "{:?}");
11
12pub enum LockError {
13 UsedByOthers(PId),
14}
15
16pub struct Lock<T> {
17 data: Arc<LockData<T>>,
18}
19
20impl<T> Lock<T> {
21 pub fn new(data: T) -> Self {
22 Lock {
23 data: Arc::new(LockData::new(data)),
24 }
25 }
26
27 pub fn try_borrow(&self, pid: PId) -> Result<LockGuard<T>, LockError> {
28 match self
29 .data
30 .borrowed
31 .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
32 {
33 Ok(_) => {
34 unsafe {
35 (*self.data.used.get()).replace(pid);
36 }
37 Ok(LockGuard {
38 data: self.data.clone(),
39 })
40 }
41 Err(_) => {
42 let pid = unsafe { *self.data.used.get() };
43 Err(LockError::UsedByOthers(pid.unwrap()))
44 }
45 }
46 }
47
48 pub fn weak(&self) -> LockWeak<T> {
49 LockWeak {
50 data: Arc::downgrade(&self.data),
51 }
52 }
53
54 pub unsafe fn force_use(&self) -> *mut T {
59 self.data.data.get()
60 }
61}
62
63impl<T: Sync + Send> Deref for Lock<T> {
64 type Target = T;
65
66 fn deref(&self) -> &Self::Target {
67 unsafe { &*self.data.data.get() }
68 }
69}
70
71pub struct LockWeak<T> {
72 data: Weak<LockData<T>>,
73}
74
75impl<T> LockWeak<T> {
76 pub fn upgrade(&self) -> Option<Lock<T>> {
77 self.data.upgrade().map(|data| Lock { data })
78 }
79}
80
81struct LockData<T> {
82 borrowed: AtomicBool,
83 used: UnsafeCell<Option<PId>>,
84 data: UnsafeCell<T>,
85}
86
87unsafe impl<T: Send> Send for LockData<T> {}
88unsafe impl<T: Send> Sync for LockData<T> {}
89
90impl<T> LockData<T> {
91 fn new(data: T) -> Self {
92 LockData {
93 borrowed: AtomicBool::new(false),
94 used: UnsafeCell::new(None),
95 data: UnsafeCell::new(data),
96 }
97 }
98}
99
100pub struct LockGuard<T> {
101 data: Arc<LockData<T>>,
102}
103
104impl<T> Drop for LockGuard<T> {
105 fn drop(&mut self) {
106 unsafe {
107 (*self.data.used.get()).take();
108 }
109 self.data.borrowed.store(false, Ordering::Release);
110 }
111}
112
113impl<T> Deref for LockGuard<T> {
114 type Target = T;
115
116 fn deref(&self) -> &Self::Target {
117 unsafe { &*self.data.data.get() }
118 }
119}
120
121impl<T> DerefMut for LockGuard<T> {
122 fn deref_mut(&mut self) -> &mut Self::Target {
123 unsafe { &mut *self.data.data.get() }
124 }
125}