| 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 | |