mtb_entity_slab/
id.rs

1use crate::{
2    EntityAlloc,
3    alloc::IEntityAllocatable,
4    chunk::{Chunk, NULL_INDEXED_ID, Unit, UnitPtrError},
5};
6use std::{
7    fmt::Debug,
8    hash::Hash,
9    marker::PhantomData,
10    ops::{Deref, DerefMut},
11    ptr::NonNull,
12};
13
14pub trait IEntityAllocID<E: IEntityAllocatable>: Sized + Copy {
15    fn from_ptr(alloc: &EntityAlloc<E>, ptr: PtrID<E>) -> Option<Self>;
16    fn from_index(alloc: &EntityAlloc<E>, indexed: usize) -> Option<Self>;
17    fn try_deref(self, alloc: &EntityAlloc<E>) -> Option<&E>;
18    fn try_deref_mut(self, alloc: &mut EntityAlloc<E>) -> Option<&mut E>;
19
20    #[inline]
21    fn deref(self, alloc: &EntityAlloc<E>) -> &E {
22        self.try_deref(alloc).expect("UAF detected!")
23    }
24    #[inline]
25    fn deref_mut(self, alloc: &mut EntityAlloc<E>) -> &mut E {
26        self.try_deref_mut(alloc).expect("UAF detected!")
27    }
28
29    fn free(self, alloc: &mut EntityAlloc<E>) -> Option<E>;
30}
31
32#[repr(transparent)]
33pub struct PtrID<E: IEntityAllocatable>(pub(crate) NonNull<Unit<E>>);
34
35impl<E: IEntityAllocatable> Debug for PtrID<E> {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        let ty_name = std::any::type_name::<E>();
38        let addr = self.0;
39        write!(f, "Ptr:{ty_name}({addr:p})")
40    }
41}
42impl<E: IEntityAllocatable> std::fmt::Pointer for PtrID<E> {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        write!(f, "{:p}", self.0)
45    }
46}
47impl<E: IEntityAllocatable> Clone for PtrID<E> {
48    fn clone(&self) -> Self {
49        Self(self.0)
50    }
51}
52impl<E: IEntityAllocatable> Copy for PtrID<E> {}
53impl<E: IEntityAllocatable> PartialEq for PtrID<E> {
54    fn eq(&self, other: &Self) -> bool {
55        self.0 == other.0
56    }
57}
58impl<E: IEntityAllocatable> Eq for PtrID<E> {}
59impl<E: IEntityAllocatable> PartialOrd for PtrID<E> {
60    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
61        self.0.partial_cmp(&other.0)
62    }
63}
64impl<E: IEntityAllocatable> Ord for PtrID<E> {
65    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
66        self.0.cmp(&other.0)
67    }
68}
69impl<E: IEntityAllocatable> Hash for PtrID<E> {
70    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
71        self.0.hash(state);
72    }
73}
74unsafe impl<E: Send + IEntityAllocatable> Send for PtrID<E> {}
75unsafe impl<E: Sync + IEntityAllocatable> Sync for PtrID<E> {}
76
77impl<E: IEntityAllocatable> IEntityAllocID<E> for PtrID<E> {
78    #[inline]
79    fn from_ptr(_: &EntityAlloc<E>, ptr: PtrID<E>) -> Option<Self> {
80        Some(ptr)
81    }
82
83    fn from_index(alloc: &EntityAlloc<E>, indexed: usize) -> Option<Self> {
84        let chunk_id = E::chunk_of_indexed_id(indexed) as usize;
85        let unit_id = E::unit_of_indexed_id(indexed);
86        let chunks = alloc.chunks.borrow();
87        let chunk = chunks.get(chunk_id)?;
88        Some(Self(NonNull::from_ref(chunk.unit(unit_id))))
89    }
90
91    #[inline]
92    fn try_deref(self, alloc: &EntityAlloc<E>) -> Option<&E> {
93        if cfg!(debug_assertions) {
94            if !self.check_validity(alloc.chunks.borrow().as_slice()) {
95                return None;
96            }
97        }
98        unsafe { self.0.as_ref() }.as_init_ref()
99    }
100
101    #[inline]
102    fn try_deref_mut(self, alloc: &mut EntityAlloc<E>) -> Option<&mut E> {
103        if cfg!(debug_assertions) {
104            if !self.check_validity(alloc.chunks.get_mut()) {
105                return None;
106            }
107        }
108        let mut ptr = self.0;
109        unsafe { ptr.as_mut() }.as_init_mut()
110    }
111
112    fn free(self, alloc: &mut EntityAlloc<E>) -> Option<E> {
113        alloc.free_ptr(self)
114    }
115}
116
117impl<E: IEntityAllocatable> PtrID<E> {
118    pub(crate) fn check_validity(self, chunks: &[Chunk<E, E::AllocatePolicyT>]) -> bool {
119        for chunk in chunks.iter() {
120            match chunk.unit_of_ptr(self.0.as_ptr()) {
121                Ok(_) => return true,
122                Err(UnitPtrError::OutOfRange) => continue,
123                Err(UnitPtrError::NotAlignedWith) => return false,
124            }
125        }
126        false
127    }
128
129    #[inline]
130    pub unsafe fn direct_get_indexed(self) -> usize {
131        unsafe { self.0.as_ref().indexed_id }
132    }
133    #[inline]
134    pub unsafe fn direct_deref<'a>(self) -> &'a E {
135        unsafe { self.0.as_ref() }.as_init_ref().unwrap()
136    }
137    #[inline]
138    pub unsafe fn direct_deref_mut<'a>(mut self) -> &'a mut E {
139        unsafe { self.0.as_mut() }.as_init_mut().unwrap()
140    }
141
142    #[inline]
143    pub fn as_indexed(self, alloc: &EntityAlloc<E>) -> Option<IndexedID<E>> {
144        IndexedID::from_ptr(alloc, self)
145    }
146
147    /// Get a raw pointer to the Unit<E>.
148    ///
149    /// Returning a raw pointer is safe; dereferencing it remains unsafe and
150    /// the caller must ensure the pointed unit is still allocated and valid.
151    #[inline]
152    pub fn as_unit_pointer(self) -> *const Unit<E> {
153        self.0.as_ptr()
154    }
155}
156
157#[repr(transparent)]
158pub struct IndexedID<E: IEntityAllocatable>(pub usize, pub PhantomData<E>);
159
160impl<E: IEntityAllocatable> Debug for IndexedID<E> {
161    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
162        let ty_name = std::any::type_name::<E>();
163        let indexed = self.0;
164        write!(f, "Idx:{ty_name}({indexed})")
165    }
166}
167impl<E: IEntityAllocatable> std::fmt::LowerHex for IndexedID<E> {
168    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
169        write!(f, "{:x}", self.0)
170    }
171}
172impl<E: IEntityAllocatable> std::fmt::UpperHex for IndexedID<E> {
173    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
174        write!(f, "{:X}", self.0)
175    }
176}
177impl<E: IEntityAllocatable> Clone for IndexedID<E> {
178    fn clone(&self) -> Self {
179        Self(self.0, PhantomData)
180    }
181}
182impl<E: IEntityAllocatable> Copy for IndexedID<E> {}
183impl<E: IEntityAllocatable> PartialEq for IndexedID<E> {
184    fn eq(&self, other: &Self) -> bool {
185        self.0 == other.0
186    }
187}
188impl<E: IEntityAllocatable> Eq for IndexedID<E> {}
189impl<E: IEntityAllocatable> PartialOrd for IndexedID<E> {
190    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
191        self.0.partial_cmp(&other.0)
192    }
193}
194impl<E: IEntityAllocatable> Ord for IndexedID<E> {
195    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
196        self.0.cmp(&other.0)
197    }
198}
199impl<E: IEntityAllocatable> Hash for IndexedID<E> {
200    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
201        self.0.hash(state);
202    }
203}
204
205impl<E: IEntityAllocatable> IEntityAllocID<E> for IndexedID<E> {
206    fn from_ptr(_: &EntityAlloc<E>, ptr: PtrID<E>) -> Option<Self> {
207        let indexed = unsafe { ptr.direct_get_indexed() };
208        if indexed == NULL_INDEXED_ID {
209            None
210        } else {
211            Some(Self(indexed, PhantomData))
212        }
213    }
214
215    fn from_index(_: &EntityAlloc<E>, indexed: usize) -> Option<Self> {
216        Some(Self(indexed, PhantomData))
217    }
218
219    fn try_deref(self, alloc: &EntityAlloc<E>) -> Option<&E> {
220        let chunk_id = E::chunk_of_indexed_id(self.0) as usize;
221        let unit_id = E::unit_of_indexed_id(self.0) as usize;
222        let chunks_ref = alloc.chunks.borrow();
223        let Some(chunk) = chunks_ref.get(chunk_id) else {
224            return None;
225        };
226        // Capture a stable raw pointer to the Units buffer (Box allocation is stable even if the
227        // chunks Vec reallocates). Then release the RefCell borrow before creating &E.
228        let units_ptr = chunk.units().as_ptr();
229        drop(chunks_ref);
230        unsafe { units_ptr.add(unit_id).as_ref()?.as_init_ref() }
231    }
232    fn try_deref_mut(self, alloc: &mut EntityAlloc<E>) -> Option<&mut E> {
233        let chunk_id = E::chunk_of_indexed_id(self.0);
234        let unit_id = E::unit_of_indexed_id(self.0);
235        let chunks = alloc.chunks.get_mut();
236        let chunk = &mut chunks[chunk_id as usize];
237        chunk.unit_mut(unit_id).as_init_mut()
238    }
239
240    fn free(self, alloc: &mut EntityAlloc<E>) -> Option<E> {
241        alloc.free_indexed(self.0)
242    }
243}
244
245impl<E: IEntityAllocatable> IndexedID<E> {
246    pub fn as_ptr(self, alloc: &EntityAlloc<E>) -> Option<PtrID<E>> {
247        PtrID::from_index(alloc, self.0)
248    }
249}
250
251/// Maybe unsafe!
252pub struct IDProxy<'alloc, E: IEntityAllocatable> {
253    pub(crate) unit: &'alloc mut Unit<E>,
254}
255impl<'alloc, E: IEntityAllocatable> Deref for IDProxy<'alloc, E> {
256    type Target = E;
257
258    fn deref(&self) -> &Self::Target {
259        self.unit.as_init_ref().unwrap()
260    }
261}
262impl<'alloc, E: IEntityAllocatable> DerefMut for IDProxy<'alloc, E> {
263    fn deref_mut(&mut self) -> &mut Self::Target {
264        self.unit.as_init_mut().unwrap()
265    }
266}
267impl<'alloc, E: IEntityAllocatable> IDProxy<'alloc, E> {
268    pub fn new(alloc: &'alloc EntityAlloc<E>, val: E) -> Self {
269        let mut ptr = alloc.allocate(val);
270        let unit = unsafe { ptr.0.as_mut() };
271        Self { unit }
272    }
273
274    pub fn get(&self) -> Option<&E> {
275        self.unit.as_init_ref()
276    }
277    pub fn get_mut(&mut self) -> Option<&mut E> {
278        self.unit.as_init_mut()
279    }
280    pub(super) unsafe fn raw_free(self) -> Option<E> {
281        if self.unit.indexed_id == NULL_INDEXED_ID {
282            return None;
283        }
284        self.unit.indexed_id = NULL_INDEXED_ID;
285        unsafe { Some(self.unit.data.assume_init_read()) }
286    }
287
288    pub fn release_ptr(self) -> PtrID<E> {
289        PtrID(NonNull::from(self.unit))
290    }
291    pub fn release_indexed(self) -> IndexedID<E> {
292        IndexedID(self.unit.indexed_id, PhantomData)
293    }
294    pub fn release_all(self) -> (PtrID<E>, IndexedID<E>) {
295        (
296            PtrID(NonNull::from_ref(self.unit)),
297            IndexedID(self.unit.indexed_id, PhantomData),
298        )
299    }
300    pub(super) unsafe fn copy_release_all(&self) -> (PtrID<E>, IndexedID<E>) {
301        (
302            PtrID(NonNull::from_ref(self.unit)),
303            IndexedID(self.unit.indexed_id, PhantomData),
304        )
305    }
306}
mtb_entity_slab/id.rs

mtb_entity_slab/
id.rs
