1use core::{cell::UnsafeCell, fmt, mem::ManuallyDrop};
2
3use crate::lock::{rank, RankData, RwLock, RwLockReadGuard, RwLockWriteGuard};
4
5pub struct SnatchGuard<'a>(RwLockReadGuard<'a, ()>);
7pub struct ExclusiveSnatchGuard<'a>(#[expect(dead_code)] RwLockWriteGuard<'a, ()>);
9
10pub struct Snatchable<T> {
17 value: UnsafeCell<Option<T>>,
18}
19
20impl<T> Snatchable<T> {
21 pub fn new(val: T) -> Self {
22 Snatchable {
23 value: UnsafeCell::new(Some(val)),
24 }
25 }
26
27 #[allow(dead_code)]
28 pub fn empty() -> Self {
29 Snatchable {
30 value: UnsafeCell::new(None),
31 }
32 }
33
34 pub fn get<'a>(&'a self, _guard: &'a SnatchGuard) -> Option<&'a T> {
36 unsafe { (*self.value.get()).as_ref() }
37 }
38
39 pub fn snatch(&self, _guard: &mut ExclusiveSnatchGuard) -> Option<T> {
41 unsafe { (*self.value.get()).take() }
42 }
43
44 pub fn take(&mut self) -> Option<T> {
49 self.value.get_mut().take()
50 }
51}
52
53impl<T> fmt::Debug for Snatchable<T> {
56 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
57 write!(f, "<snatchable>")
58 }
59}
60
61unsafe impl<T> Sync for Snatchable<T> {}
62
63use trace::LockTrace;
64#[cfg(all(debug_assertions, feature = "std"))]
65mod trace {
66 use core::{cell::Cell, fmt, panic::Location};
67 use std::{backtrace::Backtrace, thread};
68
69 pub(super) struct LockTrace {
70 purpose: &'static str,
71 caller: &'static Location<'static>,
72 backtrace: Backtrace,
73 }
74
75 impl fmt::Display for LockTrace {
76 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
77 write!(
78 f,
79 "a {} lock at {}\n{}",
80 self.purpose, self.caller, self.backtrace
81 )
82 }
83 }
84
85 impl LockTrace {
86 #[track_caller]
87 pub(super) fn enter(purpose: &'static str) {
88 let new = LockTrace {
89 purpose,
90 caller: Location::caller(),
91 backtrace: Backtrace::capture(),
92 };
93
94 if let Some(prev) = SNATCH_LOCK_TRACE.take() {
95 let current = thread::current();
96 let name = current.name().unwrap_or("<unnamed>");
97 panic!(
98 "thread '{name}' attempted to acquire a snatch lock recursively.\n\
99 - Currently trying to acquire {new}\n\
100 - Previously acquired {prev}",
101 );
102 } else {
103 SNATCH_LOCK_TRACE.set(Some(new));
104 }
105 }
106
107 pub(super) fn exit() {
108 SNATCH_LOCK_TRACE.take();
109 }
110 }
111
112 std::thread_local! {
113 static SNATCH_LOCK_TRACE: Cell<Option<LockTrace>> = const { Cell::new(None) };
114 }
115}
116#[cfg(not(all(debug_assertions, feature = "std")))]
117mod trace {
118 pub(super) struct LockTrace {
119 _private: (),
120 }
121
122 impl LockTrace {
123 pub(super) fn enter(_purpose: &'static str) {}
124 pub(super) fn exit() {}
125 }
126}
127
128pub struct SnatchLock {
130 lock: RwLock<()>,
131}
132
133impl SnatchLock {
134 pub unsafe fn new(rank: rank::LockRank) -> Self {
139 SnatchLock {
140 lock: RwLock::new(rank, ()),
141 }
142 }
143
144 #[track_caller]
146 pub fn read(&self) -> SnatchGuard<'_> {
147 LockTrace::enter("read");
148 SnatchGuard(self.lock.read())
149 }
150
151 #[track_caller]
157 pub fn write(&self) -> ExclusiveSnatchGuard<'_> {
158 LockTrace::enter("write");
159 ExclusiveSnatchGuard(self.lock.write())
160 }
161
162 #[track_caller]
163 pub unsafe fn force_unlock_read(&self, data: RankData) {
164 LockTrace::exit();
168 unsafe { self.lock.force_unlock_read(data) };
169 }
170}
171
172impl SnatchGuard<'_> {
173 pub fn forget(this: Self) -> RankData {
178 let manually_drop = ManuallyDrop::new(this);
180
181 let inner_guard = unsafe { core::ptr::read(&manually_drop.0) };
185
186 RwLockReadGuard::forget(inner_guard)
187 }
188}
189
190impl Drop for SnatchGuard<'_> {
191 fn drop(&mut self) {
192 LockTrace::exit();
193 }
194}
195
196impl Drop for ExclusiveSnatchGuard<'_> {
197 fn drop(&mut self) {
198 LockTrace::exit();
199 }
200}