fragile/
registry.rs

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}