stable_id/
entities.rs

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    /** Reserves spaces similar to [`Vec::with_capacity()`]. */
18    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    /** Returns the number of items in this data structure. */
27    pub fn len(&self) -> usize {
28        self.data.len()
29    }
30
31    /** Tells you if the collection is empty. */
32    pub fn is_empty(&self) -> bool {
33        self.data.is_empty()
34    }
35
36    /** Try getting the item with the given id. */
37    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    /** Mutable version of get. */
44    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    /**
51    Removes an element for the given id.
52    */
53    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"); // contradiction: we just found the physical id
61
62            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                // we can perform the cast because log(MAX) is always smaller than MAX
71                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    /**
85    Allocate an entity with monotonically increase ids, just like [`crate::SparseEntities`].
86    */
87    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    /// Return all data's references.
97    pub fn iter(&self) -> impl Iterator<Item = &DataT> {
98        self.data.iter()
99    }
100
101    /// Return all data in mutable references.
102    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut DataT> {
103        self.data.iter_mut()
104    }
105
106    /**
107    Iterate every entries. This takes O(`HashMap::iter()`) to iterate the entire collection.
108    */
109    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    /**
118    Compact spaces internally.
119    */
120    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    /**
155    Populate `count` number of items by cloning the given `data`.
156    */
157    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    /**
178    Populate `count` number of items with the default value.
179    */
180    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); // coalesce() was called since we removed a majority of items.
427    }
428}