1use crate::{
2 EntityAllocConsumeIter, EntityAllocEditIter, EntityAllocReadIter, IEntityAllocPolicy,
3 IndexedID, PtrID,
4 chunk::{Chunk, NULL_INDEXED_ID},
5 iter::EntitySlotIter,
6};
7use std::{
8 cell::{BorrowMutError, Cell, RefCell},
9 marker::PhantomData,
10 ptr::NonNull,
11};
12
13pub trait IEntityAllocatable: Sized {
14 type AllocatePolicyT: IEntityAllocPolicy<Self>;
15 const CHUNK_SIZE: usize = Self::AllocatePolicyT::CHUNK_SIZE;
16 const BITSET_LEN: usize = Self::AllocatePolicyT::BITSET_LEN;
17
18 fn compose_indexed_id(chunk_id: u32, unit_id: u16) -> usize {
19 Self::AllocatePolicyT::compose_indexed_id(chunk_id, unit_id)
20 }
21 fn split_indexed_id(indexed_id: usize) -> (u32, u16) {
22 Self::AllocatePolicyT::split_indexed_id(indexed_id)
23 }
24 fn unit_of_indexed_id(indexed_id: usize) -> u16 {
25 Self::AllocatePolicyT::unit_of_indexed_id(indexed_id)
26 }
27 fn chunk_of_indexed_id(indexed_id: usize) -> u32 {
28 Self::AllocatePolicyT::chunk_of_indexed_id(indexed_id)
29 }
30}
31
32pub struct EntityAlloc<E: IEntityAllocatable> {
33 pub(crate) chunks: RefCell<Vec<Chunk<E, E::AllocatePolicyT>>>,
34 pub(crate) free_head: Cell<u32>,
35 pub(crate) num_allocated: Cell<usize>,
36}
37
38impl<'alloc, E: IEntityAllocatable> IntoIterator for &'alloc EntityAlloc<E> {
39 type Item = (usize, PtrID<E>, &'alloc E);
40 type IntoIter = EntityAllocReadIter<'alloc, E>;
41
42 #[inline]
43 fn into_iter(self) -> Self::IntoIter {
44 EntityAllocReadIter::new(self)
45 }
46}
47impl<'alloc, E: IEntityAllocatable> IntoIterator for &'alloc mut EntityAlloc<E> {
48 type Item = (usize, PtrID<E>, &'alloc mut E);
49 type IntoIter = EntityAllocEditIter<'alloc, E>;
50
51 #[inline]
52 fn into_iter(self) -> Self::IntoIter {
53 EntityAllocEditIter::new(self)
54 }
55}
56impl<E: IEntityAllocatable> IntoIterator for EntityAlloc<E> {
57 type Item = (usize, E);
58 type IntoIter = EntityAllocConsumeIter<E>;
59
60 #[inline]
61 fn into_iter(self) -> Self::IntoIter {
62 EntityAllocConsumeIter::new(self)
63 }
64}
65
66impl<E: IEntityAllocatable> EntityAlloc<E> {
67 pub fn new() -> Self {
68 Self {
69 chunks: RefCell::new(Vec::new()),
70 free_head: Cell::new(0),
71 num_allocated: Cell::new(0),
72 }
73 }
74
75 pub fn with_capacity(cap: usize) -> Self {
76 let cap = cap.next_multiple_of(E::AllocatePolicyT::CHUNK_SIZE);
77 let chunks = {
78 let nchunks = cap / E::AllocatePolicyT::CHUNK_SIZE;
79 let mut chunks = Vec::with_capacity(nchunks);
80 for chunk_id in 0..(nchunks as u32) {
81 let chunk = Chunk::new(chunk_id, chunk_id + 1);
82 chunks.push(chunk);
83 }
84 RefCell::new(chunks)
85 };
86 Self {
87 chunks,
88 free_head: Cell::new(0),
89 num_allocated: Cell::new(0),
90 }
91 }
92
93 pub fn allocate(&self, val: E) -> PtrID<E> {
94 let Ok(ptr) = self.try_allocate(val) else {
95 panic!(
96 "Cannot borrow EntityAlloc {self:p} mutably for allocation. Are you traversing it now?"
97 );
98 };
99 ptr
100 }
101 pub fn try_allocate(&self, val: E) -> Result<PtrID<E>, (BorrowMutError, E)> {
102 let mut chunks = match self.chunks.try_borrow_mut() {
103 Ok(c) => c,
104 Err(e) => return Err((e, val)),
105 };
106 let chunk = Self::find_chunk(self.free_head.get(), &mut chunks);
107 let Ok((ptr, _)) = chunk.allocate(val) else {
108 panic!("Allocation error!");
109 };
110 self.num_allocated.set(self.num_allocated.get() + 1);
111 if chunk.is_full() {
112 self.free_head.set(chunk.free_next);
113 }
114 Ok(PtrID(NonNull::new(ptr as *mut _).unwrap()))
115 }
116 pub fn next_indexed(&self) -> IndexedID<E> {
118 let chunks = self.chunks.borrow();
119 let free_head = self.free_head.get() as usize;
120 let Some(chunk) = chunks.get(free_head) else {
121 return IndexedID(
122 free_head << E::AllocatePolicyT::CHUNK_SIZE_LOG2,
123 PhantomData,
124 );
125 };
126 let unit_id = chunk.next_available().unwrap();
127 IndexedID(chunk.indexed_id_of_unit(unit_id), PhantomData)
128 }
129 fn find_chunk(
130 free: u32,
131 chunks: &mut Vec<Chunk<E, E::AllocatePolicyT>>,
132 ) -> &mut Chunk<E, E::AllocatePolicyT> {
133 let free = free as usize;
134 while chunks.len() <= free {
135 let id = chunks.len() as u32;
136 chunks.push(Chunk::new(id, id + 1));
137 }
138 &mut chunks[free]
139 }
140
141 pub fn free_indexed(&mut self, indexed: usize) -> Option<E> {
142 let chunks = self.chunks.get_mut();
143 let chunk_id = E::chunk_of_indexed_id(indexed);
144 let unit_id = E::unit_of_indexed_id(indexed);
145 let chunk = chunks.get_mut(chunk_id as usize)?;
146 let chunk_full = chunk.is_full();
147 let e = chunk.deallocate_unit_id(unit_id)?;
148 if chunk_full {
150 chunk.free_next = self.free_head.get();
151 self.free_head.set(chunk_id);
152 }
153 self.num_allocated.set(self.num_allocated.get() - 1);
154 Some(e)
155 }
156 pub fn free_ptr(&mut self, ptr: PtrID<E>) -> Option<E> {
157 if cfg!(debug_assertions) && !ptr.check_validity(self.chunks.get_mut()) {
158 return None;
159 }
160 let indexed = unsafe { ptr.direct_get_indexed() };
161 if indexed == NULL_INDEXED_ID {
162 return None;
163 }
164 self.free_indexed(indexed)
165 }
166 pub fn free_if(&mut self, pred: impl Fn(&E, PtrID<E>, IndexedID<E>) -> bool) {
167 self.fully_free_if(pred, drop);
168 }
169 pub fn fully_free_if(
170 &mut self,
171 pred: impl Fn(&E, PtrID<E>, IndexedID<E>) -> bool,
172 mut consume: impl FnMut(E),
173 ) {
174 let mut num_allocated_dec = self.num_allocated.get();
175 let mut num_allocated_inc = 0;
176 for slot in EntitySlotIter::new(self) {
177 let (ptr, indexed_id) = unsafe { slot.copy_release_all() };
178 if pred(slot.get().unwrap(), ptr, indexed_id) {
179 let elem = unsafe { slot.raw_free() };
180 let Some(elem) = elem else {
181 continue;
182 };
183 consume(elem);
184 num_allocated_dec -= 1;
185 } else {
186 num_allocated_inc += 1;
187 }
188 }
189 assert_eq!(
190 num_allocated_dec, num_allocated_inc,
191 "Length maintenance error in fully_free_if"
192 );
193 self.num_allocated.set(num_allocated_dec);
194 }
195
196 #[inline]
197 pub fn len(&self) -> usize {
198 self.num_allocated.get()
199 }
200 #[inline]
201 pub fn capacity(&self) -> usize {
202 self.chunks.borrow().len() * E::AllocatePolicyT::CHUNK_SIZE
203 }
204 #[inline]
205 pub fn is_empty(&self) -> bool {
206 self.len() == 0
207 }
208
209 pub fn clear(&mut self) {
210 self.chunks.get_mut().clear();
211 self.free_head.set(0);
212 self.num_allocated.set(0);
213 }
214
215 #[inline]
216 pub fn iter(&self) -> EntityAllocReadIter<'_, E> {
217 EntityAllocReadIter::new(self)
218 }
219 #[inline]
220 pub fn iter_mut(&mut self) -> EntityAllocEditIter<'_, E> {
221 EntityAllocEditIter::new(self)
222 }
223}