chore: refresco desde el monorepo (pin tawasuyu.git v0.2.0)
1//! Various utilities used for user data implementations
2
3use std::any::Any;
4use std::mem::ManuallyDrop;
5use std::sync::OnceLock;
6use std::thread::{self, ThreadId};
7
8use self::list::AppendList;
9
10// `UserData.get()` is called frequently, and unfortunately
11// `thread::current().id()` is not very efficient to be calling every time.
12#[inline]
13fn current_thread_id() -> ThreadId {
14 thread_local! {
15 static ID: ThreadId = thread::current().id();
16 }
17 ID.with(|id| *id)
18}
19
20/// A wrapper for user data, able to store any type, and correctly
21/// handling access from a wrong thread
22#[derive(Debug)]
23pub struct UserData {
24 inner: OnceLock<UserDataInner>,
25}
26
27#[derive(Debug)]
28enum UserDataInner {
29 ThreadSafe(Box<dyn Any + Send + Sync + 'static>),
30 NonThreadSafe(Box<ManuallyDrop<dyn Any + 'static>>, ThreadId),
31}
32
33// UserData itself is always threadsafe, as it only gives access to its
34// content if it is Send + Sync, or we are on the right thread
35unsafe impl Send for UserData {}
36unsafe impl Sync for UserData {}
37
38impl Default for UserData {
39 fn default() -> Self {
40 Self::new()
41 }
42}
43
44impl UserData {
45 /// Create a new UserData instance
46 pub const fn new() -> UserData {
47 UserData {
48 inner: OnceLock::new(),
49 }
50 }
51
52 /// Sets the UserData to a given value
53 ///
54 /// The provided closure is called to init the UserData,
55 /// does nothing is the UserData had already been set.
56 pub fn set<T: Any + 'static, F: FnOnce() -> T>(&self, f: F) {
57 self.inner.get_or_init(|| {
58 UserDataInner::NonThreadSafe(Box::new(ManuallyDrop::new(f())), current_thread_id())
59 });
60 }
61
62 /// Sets the UserData to a given threadsafe value
63 ///
64 /// The provided closure is called to init the UserData,
65 /// does nothing is the UserData had already been set.
66 pub fn set_threadsafe<T: Any + Send + Sync + 'static, F: FnOnce() -> T>(&self, f: F) {
67 self.inner
68 .get_or_init(|| UserDataInner::ThreadSafe(Box::new(f())));
69 }
70
71 /// Attempt to access the wrapped user data
72 ///
73 /// Will return `None` if either:
74 ///
75 /// - The requested type `T` does not match the type used for construction
76 /// - This `UserData` has been created using the non-threadsafe variant and access
77 /// is attempted from another thread than the one it was created on
78 pub fn get<T: 'static>(&self) -> Option<&T> {
79 match self.inner.get() {
80 Some(UserDataInner::ThreadSafe(val)) => <dyn Any>::downcast_ref::<T>(&**val),
81 Some(&UserDataInner::NonThreadSafe(ref val, threadid)) => {
82 // only give access if we are on the right thread
83 if threadid == current_thread_id() {
84 <dyn Any>::downcast_ref::<T>(&***val)
85 } else {
86 None
87 }
88 }
89 None => None,
90 }
91 }
92}
93
94impl Drop for UserData {
95 fn drop(&mut self) {
96 // only drop non-Send user data if we are on the right thread, leak it otherwise
97 if let Some(&mut UserDataInner::NonThreadSafe(ref mut val, threadid)) = self.inner.get_mut() {
98 if threadid == current_thread_id() {
99 unsafe {
100 ManuallyDrop::drop(&mut **val);
101 }
102 }
103 }
104 }
105}
106
107/// A storage able to store several values of `UserData`
108/// of different types. It behaves similarly to a `TypeMap`.
109#[derive(Debug)]
110pub struct UserDataMap {
111 list: AppendList<UserData>,
112}
113
114impl UserDataMap {
115 /// Create a new map
116 pub fn new() -> UserDataMap {
117 UserDataMap {
118 list: AppendList::new(),
119 }
120 }
121
122 /// Attempt to access the wrapped user data of a given type
123 ///
124 /// Will return `None` if no value of type `T` is stored in this `UserDataMap`
125 /// and accessible from this thread
126 pub fn get<T: 'static>(&self) -> Option<&T> {
127 for user_data in &self.list {
128 if let Some(val) = user_data.get::<T>() {
129 return Some(val);
130 }
131 }
132 None
133 }
134
135 /// Access the user data of a given type, initializing it if required.
136 pub fn get_or_insert<T: 'static, F: FnOnce() -> T>(&self, init: F) -> &T {
137 match self.get() {
138 Some(data) => data,
139 None => {
140 // Insert the new node.
141 self.insert(init);
142
143 // Return it again immediately.
144 self.list.iter().last().and_then(|data| data.get::<T>()).unwrap()
145 }
146 }
147 }
148
149 /// Access the user data of a given type, initializing it if required.
150 pub fn get_or_insert_threadsafe<T: Send + Sync + 'static, F: FnOnce() -> T>(&self, init: F) -> &T {
151 match self.get() {
152 Some(data) => data,
153 None => {
154 // Insert the new node.
155 self.insert_threadsafe(init);
156
157 // Return it again immediately.
158 self.list.iter().last().and_then(|data| data.get::<T>()).unwrap()
159 }
160 }
161 }
162
163 /// Insert a value in the map if it is not already there
164 ///
165 /// This is the non-threadsafe variant, the type you insert don't have to be
166 /// threadsafe, but they will not be visible from other threads (even if they are
167 /// actually threadsafe).
168 ///
169 /// If the value does not already exist, the closure is called to create it and
170 /// this function returns `true`. If the value already exists, the closure is not
171 /// called, and this function returns `false`.
172 pub fn insert_if_missing<T: 'static, F: FnOnce() -> T>(&self, init: F) -> bool {
173 if self.get::<T>().is_some() {
174 return false;
175 }
176
177 self.insert(init);
178
179 true
180 }
181
182 /// Insert a value in the map if it is not already there
183 ///
184 /// This is the threadsafe variant, the type you insert must be threadsafe and will
185 /// be visible from all threads.
186 ///
187 /// If the value does not already exist, the closure is called to create it and
188 /// this function returns `true`. If the value already exists, the closure is not
189 /// called, and this function returns `false`.
190 pub fn insert_if_missing_threadsafe<T: Send + Sync + 'static, F: FnOnce() -> T>(&self, init: F) -> bool {
191 if self.get::<T>().is_some() {
192 return false;
193 }
194
195 self.insert_threadsafe(init);
196
197 true
198 }
199
200 /// Insert a value into the user data.
201 fn insert<T: 'static, F: FnOnce() -> T>(&self, init: F) {
202 let data = UserData::new();
203 data.set(init);
204 self.list.append(data);
205 }
206
207 /// Insert a value into the user data.
208 fn insert_threadsafe<T: Send + Sync + 'static, F: FnOnce() -> T>(&self, init: F) {
209 let data = UserData::new();
210 data.set_threadsafe(init);
211 self.list.append(data);
212 }
213}
214
215impl Default for UserDataMap {
216 fn default() -> UserDataMap {
217 UserDataMap::new()
218 }
219}
220
221mod list {
222 /*
223 * This is a lock-free append-only list, it is used as an implementation
224 * detail of the UserDataMap.
225 *
226 * It was extracted from https://github.com/Diggsey/lockless under MIT license
227 * Copyright © Diggory Blake <diggsey@googlemail.com>
228 */
229
230 use std::sync::atomic::{AtomicPtr, Ordering};
231 use std::{mem, ptr};
232
233 type NodePtr<T> = Option<Box<Node<T>>>;
234
235 #[derive(Debug)]
236 struct Node<T> {
237 value: T,
238 next: AppendList<T>,
239 }
240
241 #[derive(Debug)]
242 pub struct AppendList<T>(AtomicPtr<Node<T>>);
243
244 impl<T> AppendList<T> {
245 fn node_into_raw(ptr: NodePtr<T>) -> *mut Node<T> {
246 match ptr {
247 Some(b) => Box::into_raw(b),
248 None => ptr::null_mut(),
249 }
250 }
251 unsafe fn node_from_raw(ptr: *mut Node<T>) -> NodePtr<T> {
252 if ptr.is_null() {
253 None
254 } else {
255 Some(Box::from_raw(ptr))
256 }
257 }
258
259 fn new_internal(ptr: NodePtr<T>) -> Self {
260 AppendList(AtomicPtr::new(Self::node_into_raw(ptr)))
261 }
262
263 pub fn new() -> Self {
264 Self::new_internal(None)
265 }
266
267 pub fn append(&self, value: T) {
268 self.append_list(AppendList::new_internal(Some(Box::new(Node {
269 value,
270 next: AppendList::new(),
271 }))));
272 }
273
274 unsafe fn append_ptr(&self, p: *mut Node<T>) {
275 loop {
276 match self
277 .0
278 .compare_exchange_weak(ptr::null_mut(), p, Ordering::AcqRel, Ordering::Acquire)
279 {
280 Ok(_) => return,
281 Err(head) => {
282 if !head.is_null() {
283 return (*head).next.append_ptr(p);
284 }
285 }
286 }
287 }
288 }
289
290 pub fn append_list(&self, other: AppendList<T>) {
291 let p = other.0.load(Ordering::Acquire);
292 mem::forget(other);
293 unsafe { self.append_ptr(p) };
294 }
295
296 pub fn iter(&self) -> AppendListIterator<'_, T> {
297 AppendListIterator(&self.0)
298 }
299
300 pub fn iter_mut(&mut self) -> AppendListMutIterator<'_, T> {
301 AppendListMutIterator(&mut self.0)
302 }
303 }
304
305 impl<'a, T> IntoIterator for &'a AppendList<T> {
306 type Item = &'a T;
307 type IntoIter = AppendListIterator<'a, T>;
308
309 fn into_iter(self) -> AppendListIterator<'a, T> {
310 self.iter()
311 }
312 }
313
314 impl<'a, T> IntoIterator for &'a mut AppendList<T> {
315 type Item = &'a mut T;
316 type IntoIter = AppendListMutIterator<'a, T>;
317
318 fn into_iter(self) -> AppendListMutIterator<'a, T> {
319 self.iter_mut()
320 }
321 }
322
323 impl<T> Drop for AppendList<T> {
324 fn drop(&mut self) {
325 unsafe { Self::node_from_raw(mem::replace(self.0.get_mut(), ptr::null_mut())) };
326 }
327 }
328
329 #[derive(Debug)]
330 pub struct AppendListIterator<'a, T>(&'a AtomicPtr<Node<T>>);
331
332 impl<'a, T: 'a> Iterator for AppendListIterator<'a, T> {
333 type Item = &'a T;
334
335 fn next(&mut self) -> Option<&'a T> {
336 let p = self.0.load(Ordering::Acquire);
337 if p.is_null() {
338 None
339 } else {
340 unsafe {
341 self.0 = &(*p).next.0;
342 Some(&(*p).value)
343 }
344 }
345 }
346 }
347
348 #[derive(Debug)]
349 pub struct AppendListMutIterator<'a, T>(&'a mut AtomicPtr<Node<T>>);
350
351 impl<'a, T: 'a> Iterator for AppendListMutIterator<'a, T> {
352 type Item = &'a mut T;
353
354 fn next(&mut self) -> Option<&'a mut T> {
355 let p = self.0.load(Ordering::Acquire);
356 if p.is_null() {
357 None
358 } else {
359 unsafe {
360 self.0 = &mut (*p).next.0;
361 Some(&mut (*p).value)
362 }
363 }
364 }
365 }
366}
367
368#[cfg(test)]
369mod tests {
370 use super::UserDataMap;
371
372 #[test]
373 fn insert_twice() {
374 let map = UserDataMap::new();
375
376 assert_eq!(map.get::<usize>(), None);
377 assert!(map.insert_if_missing(|| 42usize));
378 assert!(!map.insert_if_missing(|| 43usize));
379 assert_eq!(map.get::<usize>(), Some(&42));
380 }
381}