1 | pub struct Entry { |
2 | /// The pointer to the object stored in the registry. This is a type-erased |
3 | /// `Box<T>`. |
4 | pub ptr: *mut (), |
5 | /// The function that can be called on the above pointer to drop the object |
6 | /// and free its allocation. |
7 | pub drop: unsafe fn(*mut ()), |
8 | } |
9 | |
10 | #[cfg (feature = "slab" )] |
11 | mod slab_impl { |
12 | use std::cell::UnsafeCell; |
13 | use std::num::NonZeroUsize; |
14 | |
15 | use super::Entry; |
16 | |
17 | pub struct Registry(pub slab::Slab<Entry>); |
18 | |
19 | thread_local!(static REGISTRY: UnsafeCell<Registry> = UnsafeCell::new(Registry(slab::Slab::new()))); |
20 | |
21 | pub use usize as ItemId; |
22 | |
23 | pub fn insert(thread_id: NonZeroUsize, entry: Entry) -> ItemId { |
24 | let _ = thread_id; |
25 | REGISTRY.with(|registry| unsafe { (*registry.get()).0.insert(entry) }) |
26 | } |
27 | |
28 | pub fn with<R, F: FnOnce(&Entry) -> R>(item_id: ItemId, thread_id: NonZeroUsize, f: F) -> R { |
29 | let _ = thread_id; |
30 | REGISTRY.with(|registry| f(unsafe { &*registry.get() }.0.get(item_id).unwrap())) |
31 | } |
32 | |
33 | pub fn remove(item_id: ItemId, thread_id: NonZeroUsize) -> Entry { |
34 | let _ = thread_id; |
35 | REGISTRY.with(|registry| unsafe { (*registry.get()).0.remove(item_id) }) |
36 | } |
37 | |
38 | pub fn try_remove(item_id: ItemId, thread_id: NonZeroUsize) -> Option<Entry> { |
39 | let _ = thread_id; |
40 | REGISTRY.with(|registry| unsafe { (*registry.get()).0.try_remove(item_id) }) |
41 | } |
42 | } |
43 | |
44 | #[cfg (not(feature = "slab" ))] |
45 | mod map_impl { |
46 | use std::cell::UnsafeCell; |
47 | use std::num::NonZeroUsize; |
48 | use std::sync::atomic::{AtomicUsize, Ordering}; |
49 | |
50 | use super::Entry; |
51 | |
52 | pub struct Registry(pub std::collections::HashMap<(NonZeroUsize, NonZeroUsize), Entry>); |
53 | |
54 | thread_local!(static REGISTRY: UnsafeCell<Registry> = UnsafeCell::new(Registry(Default::default()))); |
55 | |
56 | pub type ItemId = NonZeroUsize; |
57 | |
58 | fn next_item_id() -> NonZeroUsize { |
59 | static COUNTER: AtomicUsize = AtomicUsize::new(1); |
60 | NonZeroUsize::new(COUNTER.fetch_add(1, Ordering::SeqCst)) |
61 | .expect("more than usize::MAX items" ) |
62 | } |
63 | |
64 | pub fn insert(thread_id: NonZeroUsize, entry: Entry) -> ItemId { |
65 | let item_id = next_item_id(); |
66 | REGISTRY |
67 | .with(|registry| unsafe { (*registry.get()).0.insert((thread_id, item_id), entry) }); |
68 | item_id |
69 | } |
70 | |
71 | pub fn with<R, F: FnOnce(&Entry) -> R>(item_id: ItemId, thread_id: NonZeroUsize, f: F) -> R { |
72 | REGISTRY.with(|registry| { |
73 | f(unsafe { &*registry.get() } |
74 | .0 |
75 | .get(&(thread_id, item_id)) |
76 | .unwrap()) |
77 | }) |
78 | } |
79 | |
80 | pub fn remove(item_id: ItemId, thread_id: NonZeroUsize) -> Entry { |
81 | REGISTRY |
82 | .with(|registry| unsafe { (*registry.get()).0.remove(&(thread_id, item_id)).unwrap() }) |
83 | } |
84 | |
85 | pub fn try_remove(item_id: ItemId, thread_id: NonZeroUsize) -> Option<Entry> { |
86 | REGISTRY.with(|registry| unsafe { (*registry.get()).0.remove(&(thread_id, item_id)) }) |
87 | } |
88 | } |
89 | |
90 | #[cfg (feature = "slab" )] |
91 | pub use self::slab_impl::*; |
92 | |
93 | #[cfg (not(feature = "slab" ))] |
94 | pub use self::map_impl::*; |
95 | |
96 | impl Drop for Registry { |
97 | fn drop(&mut self) { |
98 | for (_, value) in self.0.iter() { |
99 | // SAFETY: This function is only called once, and is called with the |
100 | // pointer it was created with. |
101 | unsafe { (value.drop)(value.ptr) }; |
102 | } |
103 | } |
104 | } |
105 | |