fn_map/
raw.rs

1use core::{mem::ManuallyDrop, ptr, ptr::NonNull};
2
3use bumpalo::Bump;
4use hashbrown::HashMap;
5use nohash_hasher::BuildNoHashHasher;
6use type_key::TypeKey;
7
8#[derive(Debug)]
9/// raw FnMap
10pub struct RawFnMap {
11    // [`TypeId`] only hashes lower 64 bits
12    map: HashMap<TypeKey, Val, BuildNoHashHasher<u64>>,
13
14    bump: ManuallyDrop<Bump>,
15}
16
17impl RawFnMap {
18    pub fn new() -> Self {
19        Self {
20            map: HashMap::default(),
21
22            bump: ManuallyDrop::new(Bump::new()),
23        }
24    }
25
26    pub fn get<T: 'static>(&self, key: &TypeKey) -> Option<NonNull<T>> {
27        Some(self.map.get(key)?.inner().cast::<T>())
28    }
29
30    /// insert value
31    ///
32    /// Returned pointer cannot outlive Self
33    pub fn insert<T: 'static>(&mut self, key: TypeKey, value: T) -> NonNull<T> {
34        let value = Val(NonNull::from(self.bump.alloc(value)) as NonNull<dyn Erased>);
35        let ptr = value.inner();
36
37        self.map.insert(key, value);
38
39        ptr.cast::<T>()
40    }
41
42    pub fn reset(&mut self) {
43        self.map.clear();
44        self.bump.reset();
45    }
46}
47
48impl Default for RawFnMap {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl Drop for RawFnMap {
55    fn drop(&mut self) {
56        self.map.clear();
57
58        // SAFETY: Manually dropped to ensure allocated objects to drop first
59        unsafe { ManuallyDrop::drop(&mut self.bump) }
60    }
61}
62
63trait Erased {}
64impl<T: ?Sized> Erased for T {}
65
66#[derive(Debug)]
67#[repr(transparent)]
68struct Val(NonNull<dyn Erased>);
69
70impl Val {
71    pub const fn inner(&self) -> NonNull<()> {
72        self.0.cast()
73    }
74}
75
76impl Drop for Val {
77    fn drop(&mut self) {
78        // SAFETY: Safe to drop since it is the only unique pointer
79        unsafe { ptr::drop_in_place(self.0.as_ptr()) }
80    }
81}