1pub 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")]
11mod 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"))]
45mod 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")]
91pub use self::slab_impl::*;
92
93#[cfg(not(feature = "slab"))]
94pub use self::map_impl::*;
95
96impl 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