mtb_entity_slab/
id.rs

1use std::{
2    fmt::{Debug, Formatter, LowerHex, Pointer, UpperHex},
3    hash::{Hash, Hasher},
4    marker::PhantomData,
5    ptr::NonNull,
6};
7
8use crate::{
9    EntityAlloc, NULL_INDEXED_ID, SlicePtrError,
10    chunk::{Chunk, Unit},
11    policy::IAllocPolicy,
12};
13
14pub trait IEntityAllocID<E, P: IAllocPolicy>: Sized + Copy {
15    fn from_ptr(alloc: &EntityAlloc<E, P>, ptr: PtrID<E, P>) -> Option<Self>;
16    fn from_index(alloc: &EntityAlloc<E, P>, indexed: usize) -> Option<Self>;
17    fn try_deref(self, alloc: &EntityAlloc<E, P>) -> Option<&E>;
18    fn try_deref_mut(self, alloc: &mut EntityAlloc<E, P>) -> Option<&mut E>;
19
20    #[inline]
21    fn deref(self, alloc: &EntityAlloc<E, P>) -> &E {
22        self.try_deref(alloc).expect("UAF detected!")
23    }
24    #[inline]
25    fn deref_mut(self, alloc: &mut EntityAlloc<E, P>) -> &mut E {
26        self.try_deref_mut(alloc).expect("UAF detected!")
27    }
28
29    fn free(self, alloc: &mut EntityAlloc<E, P>) -> Option<E>;
30}
31
32#[repr(transparent)]
33pub struct PtrID<E, P>(NonNull<Unit<E>>, PhantomData<P>);
34
35impl<E, P> Debug for PtrID<E, P> {
36    #[inline]
37    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
38        Debug::fmt(&self.0, f)
39    }
40}
41impl<E, P> Pointer for PtrID<E, P> {
42    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
43        Pointer::fmt(&self.0, f)
44    }
45}
46impl<E, P> Clone for PtrID<E, P> {
47    #[inline]
48    fn clone(&self) -> Self {
49        Self(self.0, PhantomData)
50    }
51}
52impl<E, P> Copy for PtrID<E, P> {}
53impl<E, P> PartialEq for PtrID<E, P> {
54    #[inline]
55    fn eq(&self, other: &Self) -> bool {
56        self.0 == other.0
57    }
58}
59impl<E, P> Eq for PtrID<E, P> {}
60impl<E, P> PartialOrd for PtrID<E, P> {
61    #[inline]
62    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
63        self.0.partial_cmp(&other.0)
64    }
65}
66impl<E, P> Ord for PtrID<E, P> {
67    #[inline]
68    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
69        self.0.cmp(&other.0)
70    }
71}
72impl<E, P> Hash for PtrID<E, P> {
73    #[inline]
74    fn hash<H: Hasher>(&self, state: &mut H) {
75        self.0.hash(state);
76    }
77}
78unsafe impl<E: Send, P> Send for PtrID<E, P> {}
79unsafe impl<E: Sync, P> Sync for PtrID<E, P> {}
80
81impl<E, P: IAllocPolicy> IEntityAllocID<E, P> for PtrID<E, P> {
82    fn from_ptr(_: &EntityAlloc<E, P>, ptr: PtrID<E, P>) -> Option<Self> {
83        Some(ptr)
84    }
85    fn from_index(alloc: &EntityAlloc<E, P>, indexed: usize) -> Option<Self> {
86        let chunk_id = P::chunk_of_indexed_id(indexed);
87        let unit_id = P::unit_of_indexed_id(indexed) as u16;
88        let chunks = alloc.chunks.borrow();
89        let chunk = chunks.get(chunk_id as usize)?;
90        let unit_ptr = chunk.unit(unit_id) as *const Unit<E> as *mut Unit<E>;
91        unsafe { Some(PtrID::from_raw_ptr(unit_ptr)) }
92    }
93
94    fn try_deref(self, alloc: &EntityAlloc<E, P>) -> Option<&E> {
95        if cfg!(debug_assertions) {
96            self.validate(alloc).expect("dangling PtrID detected");
97        }
98        if unsafe { self.direct_get_indexed() } == NULL_INDEXED_ID {
99            None
100        } else {
101            unsafe { Some(self.direct_deref()) }
102        }
103    }
104    fn try_deref_mut(self, alloc: &mut EntityAlloc<E, P>) -> Option<&mut E> {
105        if cfg!(debug_assertions) {
106            self.validate_mut(alloc).expect("dangling PtrID detected");
107        }
108        if unsafe { self.direct_get_indexed() } == NULL_INDEXED_ID {
109            None
110        } else {
111            unsafe { Some(self.direct_deref_mut()) }
112        }
113    }
114    fn free(self, alloc: &mut EntityAlloc<E, P>) -> Option<E> {
115        alloc.free_ptr(self)
116    }
117}
118
119impl<E, P: IAllocPolicy> PtrID<E, P> {
120    pub unsafe fn from_nonnull(ptr: NonNull<Unit<E>>) -> Self {
121        Self(ptr, PhantomData)
122    }
123    pub unsafe fn into_nonnull(self) -> NonNull<Unit<E>> {
124        self.0
125    }
126
127    pub unsafe fn from_raw_ptr(ptr: *mut Unit<E>) -> Self {
128        Self(NonNull::new(ptr).unwrap(), PhantomData)
129    }
130    pub unsafe fn into_raw_ptr(self) -> *mut Unit<E> {
131        self.0.as_ptr()
132    }
133
134    pub fn from_ref(unit: &Unit<E>) -> Self {
135        Self(NonNull::from_ref(unit), PhantomData)
136    }
137
138    pub fn get_index(self, alloc: &EntityAlloc<E, P>) -> Option<usize> {
139        if cfg!(debug_assertions) {
140            self.internal_validate(&alloc.chunks.borrow()).ok()
141        } else {
142            let id = unsafe { self.direct_get_indexed() };
143            if id == NULL_INDEXED_ID {
144                None
145            } else {
146                Some(id)
147            }
148        }
149    }
150
151    #[inline]
152    pub unsafe fn direct_get_indexed(self) -> usize {
153        unsafe { self.0.as_ref().indexed_id }
154    }
155    #[inline]
156    pub unsafe fn direct_deref<'a>(self) -> &'a E {
157        unsafe { self.0.as_ref() }.as_init_ref().unwrap()
158    }
159    #[inline]
160    pub unsafe fn direct_deref_mut<'a>(mut self) -> &'a mut E {
161        unsafe { self.0.as_mut() }.as_init_mut().unwrap()
162    }
163
164    pub fn validate(&self, alloc: &EntityAlloc<E, P>) -> Result<usize, SlicePtrError> {
165        self.internal_validate(&alloc.chunks.borrow())
166    }
167    pub fn validate_mut(&self, alloc: &mut EntityAlloc<E, P>) -> Result<usize, SlicePtrError> {
168        self.internal_validate(alloc.chunks.get_mut())
169    }
170    pub(crate) fn internal_validate(&self, chunks: &[Chunk<E, P>]) -> Result<usize, SlicePtrError> {
171        let ptr = self.0.as_ptr() as *const Unit<E>;
172        for chunk in chunks {
173            match chunk.unit_of_ptr(ptr) {
174                Ok(unit_id) => return Ok(chunk.indexed_id_of_unit(unit_id)),
175                Err(SlicePtrError::OutOfRange) => continue,
176                Err(SlicePtrError::NotAlignedWith) => return Err(SlicePtrError::NotAlignedWith),
177            }
178        }
179        Err(SlicePtrError::OutOfRange)
180    }
181}
182
183#[repr(transparent)]
184pub struct IndexedID<E, P>(pub usize, pub PhantomData<(E, P)>);
185
186impl<E, P> LowerHex for IndexedID<E, P> {
187    #[inline]
188    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
189        LowerHex::fmt(&self.0, f)
190    }
191}
192impl<E, P> UpperHex for IndexedID<E, P> {
193    #[inline]
194    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
195        UpperHex::fmt(&self.0, f)
196    }
197}
198impl<E, P> Clone for IndexedID<E, P> {
199    #[inline]
200    fn clone(&self) -> Self {
201        Self(self.0, PhantomData)
202    }
203}
204impl<E, P> Copy for IndexedID<E, P> {}
205impl<E, P> PartialEq for IndexedID<E, P> {
206    #[inline]
207    fn eq(&self, other: &Self) -> bool {
208        self.0 == other.0
209    }
210}
211impl<E, P> Eq for IndexedID<E, P> {}
212impl<E, P> PartialOrd for IndexedID<E, P> {
213    #[inline]
214    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
215        self.0.partial_cmp(&other.0)
216    }
217}
218impl<E, P> Ord for IndexedID<E, P> {
219    #[inline]
220    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
221        self.0.cmp(&other.0)
222    }
223}
224impl<E, P> Hash for IndexedID<E, P> {
225    #[inline]
226    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
227        self.0.hash(state);
228    }
229}
230impl<E, P> From<usize> for IndexedID<E, P> {
231    fn from(indexed: usize) -> Self {
232        Self(indexed, PhantomData)
233    }
234}
235impl<E, P> From<IndexedID<E, P>> for usize {
236    fn from(indexed: IndexedID<E, P>) -> Self {
237        indexed.0
238    }
239}
240impl<E, P: IAllocPolicy> IEntityAllocID<E, P> for IndexedID<E, P> {
241    fn from_ptr(alloc: &EntityAlloc<E, P>, ptr: PtrID<E, P>) -> Option<Self> {
242        ptr.get_index(alloc).map(IndexedID::from)
243    }
244    fn from_index(_: &EntityAlloc<E, P>, indexed: usize) -> Option<Self> {
245        Some(IndexedID::from(indexed))
246    }
247
248    fn try_deref(self, alloc: &EntityAlloc<E, P>) -> Option<&E> {
249        PtrID::from_index(alloc, self.0)?.try_deref(alloc)
250    }
251    fn try_deref_mut(self, alloc: &mut EntityAlloc<E, P>) -> Option<&mut E> {
252        PtrID::from_index(alloc, self.0)?.try_deref_mut(alloc)
253    }
254
255    fn free(self, alloc: &mut EntityAlloc<E, P>) -> Option<E> {
256        alloc.free_indexed(self.into())
257    }
258}
259
260/// Pointer ID wrapper of allocated entities with a specialized policy
261pub trait IPolicyPtrID: Copy + Eq + Debug {
262    type ObjectT: Sized;
263    type PolicyT: IAllocPolicy;
264
265    fn from_raw_ptrid(ptr: PtrID<Self::ObjectT, Self::PolicyT>) -> Self;
266    fn into_raw_ptrid(self) -> PtrID<Self::ObjectT, Self::PolicyT>;
267
268    fn deref_alloc(self, alloc: &EntityAlloc<Self::ObjectT, Self::PolicyT>) -> &Self::ObjectT {
269        self.into_raw_ptrid().deref(alloc.as_ref())
270    }
271    fn deref_from_alloc_mut(
272        self,
273        alloc: &mut EntityAlloc<Self::ObjectT, Self::PolicyT>,
274    ) -> &mut Self::ObjectT {
275        self.into_raw_ptrid().deref_mut(alloc.as_mut())
276    }
277}
278impl<T, P: IAllocPolicy> IPolicyPtrID for PtrID<T, P> {
279    type ObjectT = T;
280    type PolicyT = P;
281
282    #[inline]
283    fn from_raw_ptrid(ptr: PtrID<Self::ObjectT, Self::PolicyT>) -> Self {
284        ptr
285    }
286    #[inline]
287    fn into_raw_ptrid(self) -> PtrID<Self::ObjectT, Self::PolicyT> {
288        self
289    }
290}