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 
insert(thread_id: NonZeroUsize, entry: Entry) -> ItemId23     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 
with<R, F: FnOnce(&Entry) -> R>(item_id: ItemId, thread_id: NonZeroUsize, f: F) -> R28     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 
remove(item_id: ItemId, thread_id: NonZeroUsize) -> Entry33     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 
try_remove(item_id: ItemId, thread_id: NonZeroUsize) -> Option<Entry>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 
next_item_id() -> NonZeroUsize58     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 
insert(thread_id: NonZeroUsize, entry: Entry) -> ItemId64     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 
with<R, F: FnOnce(&Entry) -> R>(item_id: ItemId, thread_id: NonZeroUsize, f: F) -> R71     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 
remove(item_id: ItemId, thread_id: NonZeroUsize) -> Entry80     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 
try_remove(item_id: ItemId, thread_id: NonZeroUsize) -> Option<Entry>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 {
drop(&mut self)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