1use std::{
2 hash::Hash,
3 ops::{Index, IndexMut},
4};
5
6use rustc_hash::FxHashMap;
7use stable_id_traits::{CastUsize, Maximum, Successor};
8
9use crate::{Sequence, Tec};
10
11use super::Entities;
12
13impl<IndexT, DataT> Entities<IndexT, DataT>
14where
15 IndexT: Default + Successor + Clone + Copy + Hash + Eq + CastUsize + Ord + Maximum,
16{
17 pub fn with_capacity(capacity: usize) -> Self {
19 Self {
20 vtable: Default::default(),
21 data: Tec::with_capacity(capacity),
22 seq: Default::default(),
23 }
24 }
25
26 pub fn len(&self) -> usize {
28 self.data.len()
29 }
30
31 pub fn is_empty(&self) -> bool {
33 self.data.is_empty()
34 }
35
36 pub fn get(&self, index: IndexT) -> Option<&DataT> {
38 self.vtable
39 .get(&index)
40 .and_then(|physical_id| self.data.get(*physical_id).map(|data| data))
41 }
42
43 pub fn get_mut(&mut self, index: IndexT) -> Option<&mut DataT> {
45 self.vtable
46 .get(&index)
47 .and_then(|physical_id| self.data.get_mut(*physical_id).map(|data| data))
48 }
49
50 pub fn remove(&mut self, index: IndexT) -> Option<DataT> {
54 let virtual_id = index;
55 let physical_id = self.vtable.get(&virtual_id);
56
57 if let Some(&physical_id) = physical_id {
58 let data = self.data.remove(physical_id);
59
60 self.vtable.remove(&virtual_id).expect("cannot remove item"); assert_eq!(self.vtable.len(), self.data.len());
63
64 let len = self.len();
65 let capacity = self.data.capacity();
66 let num_dead_slots = capacity - len;
67 let logn = len.checked_ilog2();
68
69 if let Some(logn) = logn {
70 if num_dead_slots >= logn.cast_to() {
72 self.coalesce();
73 }
74 } else {
75 debug_assert!(len == 0);
76 }
77
78 Some(data)
79 } else {
80 None
81 }
82 }
83
84 pub fn alloc(&mut self, data: DataT) -> IndexT {
88 let virtual_id = self.seq.next_value();
89 let phyiscal_id = self.data.alloc(data);
90
91 self.vtable.insert(virtual_id, phyiscal_id);
92
93 virtual_id
94 }
95
96 pub fn iter(&self) -> impl Iterator<Item = &DataT> {
98 self.data.iter()
99 }
100
101 pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut DataT> {
103 self.data.iter_mut()
104 }
105
106 pub fn iter_with_id(&self) -> impl Iterator<Item = (IndexT, &DataT)> {
110 self.vtable.iter().map(|(virtual_id, physical_id)| {
111 let data = &self.data[*physical_id];
112
113 (*virtual_id, data)
114 })
115 }
116
117 fn coalesce(&mut self) {
121 let reverse_mapping: FxHashMap<_, _> = self.vtable.iter().map(|(a, b)| (*b, *a)).collect();
122
123 self.data.coalesce(|old_physical_id, new_physical_id| {
124 let virtual_id = reverse_mapping
125 .get(&old_physical_id)
126 .cloned()
127 .expect("inconsistent index");
128
129 self.vtable.entry(virtual_id).and_modify(|c| {
130 *c = new_physical_id;
131 });
132 })
133 }
134}
135
136impl<IndexT, DataT> Default for Entities<IndexT, DataT>
137where
138 IndexT: Default + Maximum,
139{
140 fn default() -> Self {
141 Self {
142 vtable: Default::default(),
143 data: Default::default(),
144 seq: Default::default(),
145 }
146 }
147}
148
149impl<IndexT, DataT> Entities<IndexT, DataT>
150where
151 IndexT: Successor + CastUsize + Ord + Copy + Maximum + Hash,
152 DataT: Clone,
153{
154 pub fn populate(data: DataT, count: usize) -> Self {
158 let data = Tec::populate(data, count);
159 let seq = Sequence::continue_from(CastUsize::cast_from(count));
160
161 let vtable = (0..count)
162 .map(|i| {
163 let i = CastUsize::cast_from(i);
164 (i, i)
165 })
166 .collect();
167
168 Self { vtable, data, seq }
169 }
170}
171
172impl<IndexT, DataT> Entities<IndexT, DataT>
173where
174 IndexT: CastUsize + Ord + Copy + Maximum + Successor + Hash,
175 DataT: Clone + Default,
176{
177 pub fn populate_defaults(count: usize) -> Self {
181 Self::populate(Default::default(), count)
182 }
183}
184
185impl<IndexT, DataT> Index<IndexT> for Entities<IndexT, DataT>
186where
187 IndexT: Successor + Clone + Copy + Hash + Eq + Default + CastUsize + Ord + Maximum,
188{
189 type Output = DataT;
190
191 fn index(&self, index: IndexT) -> &Self::Output {
192 self.get(index).expect("element not exist")
193 }
194}
195
196impl<IndexT, DataT> IndexMut<IndexT> for Entities<IndexT, DataT>
197where
198 IndexT: Successor + Clone + Copy + Hash + Eq + Default + CastUsize + Ord + Maximum,
199{
200 fn index_mut(&mut self, index: IndexT) -> &mut Self::Output {
201 self.get_mut(index).expect("element not exist")
202 }
203}
204
205#[cfg(test)]
206mod tests {
207 use std::collections::{HashMap, HashSet};
208
209 use crate::Entities;
210
211 #[test]
212 fn access_out_of_bound() {
213 let mut entities = Entities::default();
214 entities.alloc(1232);
215 assert_eq!(entities.get(312u16), None);
216 }
217
218 #[test]
219 #[should_panic(expected = "element not exist")]
220 fn access_out_of_bound_mut() {
221 let mut entities = Entities::default();
222 entities.alloc(1232);
223 entities[312u16] = 3333;
224 }
225
226 #[test]
227 fn populate() {
228 let count = 50;
229 let mut entities = Entities::<u8, usize>::populate_defaults(count);
230
231 assert_eq!(entities.len(), count);
232 assert_eq!(entities.alloc(54354534), count as u8);
233 assert_eq!(entities.len(), count + 1);
234 }
235
236 #[test]
237 fn normal() {
238 let mut entities = Entities::default();
239
240 fn check_all(entities: &Entities<usize, &str>) {
241 entities
242 .iter_with_id()
243 .for_each(|(id, data)| assert_eq!(entities[id], *data));
244 }
245
246 vec!["0", "1", "2", "3", "4", "5"]
247 .into_iter()
248 .fold(HashMap::new(), |mut acc, data| {
249 acc.insert(entities.alloc(data), data);
250 acc
251 })
252 .into_iter()
253 .for_each(|(id, data)| assert_eq!(entities[id], data));
254
255 assert_eq!(entities.remove(1), Some("1"));
256 check_all(&entities);
257
258 assert_eq!(entities.remove(4), Some("4"));
259 check_all(&entities);
260
261 assert_eq!(entities.remove(5), Some("5"));
262 check_all(&entities);
263
264 assert_eq!(entities.remove(3), Some("3"));
265 check_all(&entities);
266
267 assert_eq!(entities.remove(2), Some("2"));
268 assert_eq!(entities.len(), 1);
269 check_all(&entities);
270
271 assert_eq!(entities.remove(0), Some("0"));
272 assert!(entities.is_empty());
273 check_all(&entities);
274 }
275
276 #[test]
277 fn iter() {
278 let mut entities = Entities::default();
279
280 fn check_all(entities: &Entities<usize, String>) {
281 entities
282 .iter_with_id()
283 .for_each(|(id, data)| assert_eq!(entities[id], *data));
284 }
285
286 vec![
287 "0".to_owned(),
288 "1".to_owned(),
289 "2".to_owned(),
290 "3".to_owned(),
291 "4".to_owned(),
292 "5".to_owned(),
293 ]
294 .into_iter()
295 .fold(HashMap::new(), |mut acc, data| {
296 acc.insert(entities.alloc(data.clone()), data);
297 acc
298 })
299 .into_iter()
300 .for_each(|(id, data)| assert_eq!(entities[id], data));
301
302 assert_eq!(entities.remove(1), Some("1".to_owned()));
303 check_all(&entities);
304
305 assert_eq!(entities.remove(4), Some("4".to_owned()));
306 check_all(&entities);
307
308 assert_eq!(entities.remove(5), Some("5".to_owned()));
309 check_all(&entities);
310
311 assert_eq!(entities.remove(2), Some("2".to_owned()));
312 check_all(&entities);
313
314 let data_with_id = HashSet::from([(3, "3".to_owned()), (0, "0".to_owned())]);
315
316 assert_eq!(
317 HashSet::from(["3".to_owned(), "0".to_owned()]),
318 entities.iter().cloned().collect(),
319 );
320
321 assert_eq!(
322 data_with_id,
323 entities
324 .iter_with_id()
325 .map(|(id, value)| (id, value.to_owned()))
326 .collect(),
327 );
328
329 entities
330 .iter_mut()
331 .for_each(|value| *value = format!("1{value}"));
332
333 assert_eq!(
334 HashSet::from([(3, "13".to_owned()), (0, "10".to_owned())]),
335 entities
336 .iter_with_id()
337 .map(|(id, value)| (id, value.to_owned()))
338 .collect(),
339 );
340 }
341
342 #[test]
343 fn coalesce_1() {
344 let mut entities: Entities<u8, u8> = Default::default();
345 (0..255).for_each(|i| {
346 assert_eq!(entities.alloc(i), i);
347 });
348
349 entities.remove(27);
350 entities.remove(254);
351 entities.remove(15);
352 entities.remove(252);
353 entities.remove(251);
354 entities.remove(253);
355
356 entities.coalesce();
357
358 let unique_values: HashSet<_> = entities.iter_with_id().map(|(_, data)| *data).collect();
359
360 assert_eq!(unique_values.len(), 249);
361 }
362
363 #[test]
364 fn coalesce_2() {
365 let mut entities: Entities<u8, u8> = Default::default();
366 (0..255).for_each(|i| {
367 assert_eq!(entities.alloc(i), i);
368 });
369
370 entities.remove(27);
371 entities.remove(15);
372
373 entities.remove(250);
374 entities.remove(232);
375 entities.remove(231);
376 entities.remove(254);
377 entities.remove(252);
378 entities.remove(251);
379 entities.remove(25);
380 entities.remove(253);
381 entities.remove(229);
382 entities.remove(233);
383 entities.remove(234);
384 entities.remove(235);
385 entities.remove(236);
386 entities.remove(237);
387 entities.remove(238);
388 entities.remove(239);
389 entities.remove(240);
390 entities.remove(35);
391 entities.remove(241);
392 entities.remove(242);
393 entities.remove(243);
394 entities.remove(245);
395 entities.remove(244);
396 entities.remove(246);
397 entities.remove(247);
398 entities.remove(248);
399 entities.remove(34);
400 entities.remove(249);
401 entities.remove(30);
402
403 entities.coalesce();
404
405 let unique_values: HashSet<_> = entities.iter_with_id().map(|(_, data)| *data).collect();
406 assert_eq!(unique_values.len(), 224);
407 }
408
409 #[test]
410 fn coalesce_from_remove() {
411 let mut entities: Entities<usize, char> = Default::default();
412
413 ['a', 'b', 'c', 'd', 'e'].into_iter().for_each(|c| {
414 entities.alloc(c);
415 });
416
417 entities.remove(2);
418 entities.remove(3);
419 entities.remove(1);
420
421 assert_eq!(entities.len(), 2);
422 assert_eq!(
423 HashSet::from(['a', 'e']),
424 entities.iter().cloned().collect()
425 );
426 assert_eq!(entities.data.capacity(), 2); }
428}