ixa/data_structures/
entity_map.rs

1/*!
2
3An `EntityMap<E, V>` is a map from `EntityId<E>` to values of type `V` with a hash-map-like API
4optimized for densely populated maps.
5
6An `EntityMap` is "dense" in the sense that it uses a vector `Vec<Option<V>>` internally for
7storage, using the `EntityId<E>` to index into the vector. Use an `EntityMap` in any of the
8following cases:
9
10 - The number of entities is small
11 - Most entities are expected to be used as a key
12 - Access or creation speed is important
13 - You want access to the `EntityId<E>` key a value was stored with
14
15If you know beforehand how many entities you expect to store, use the `EntityMap::with_capacity`
16constructor or `EntityMap::reserve` to preallocate the map, as that is more efficient than letting
17it lazily reallocate as needed.
18
19Because every value added to an `EntityMap` is accompanied by a valid `EntityId<E>`, `EntityMap` is
20guaranteed to only "store" valid entity IDs. You can therefore use it as a replacement for
21`EntityVec<E, V>` (or just `Vec<V>`) for cases where you need to recover the original entity ID that
22a value was stored with, for example, by iterating over the (entity ID, value) pairs returned by
23`EntityMap::iter`. The only cost you pay for this is the extra memory needed to store `Option<V>`
24values instead of `V` values, which in some cases is nothing. The `EntityId<E>` itself is not
25stored.
26
27An `EntityMap<E, V>` can be cheaply converted to an `EntityVec<E, Option<V>>` using the
28`EntityMap::into_entity_vec` method.
29
30## Example
31
32Imagine you have `Person` and `Setting` entities, and you want an efficient way to store for each
33`SettingId` a `Vec<PersonId>` representing all the people that can be found in the setting. It is
34possible to use a `HashMap<SettingId, Vec<PersonId>>` to store this information, but an
35`EntityMap<SettingId, Vec<PersonId>>` is more efficient.
36
37```rust,ignore
38use ixa::data_structures::entity_map::EntityMap;
39
40let mut setting_membership = EntityMap::<SettingId, Vec<PersonId>>::new();
41
42// During population initialization you might initialize the map with data in, say, a `PersonRecord`
43// struct that has a `home_id` field of type `SettingId`.
44let person_id = context.add_entity(with!(Person, person_record.age));
45let setting_members = setting_membership.get_or_insert(person_record.home_id, Vec::new);
46setting_members.push(person_id);
47
48// Look-ups are extremely efficient.
49if let Some(setting_members) = setting_membership.get(setting_id){
50    // Do something with the setting members.
51}
52
53// You can also iterate over the (entity ID, value) pairs.
54for (setting_id, setting_members) in setting_membership.iter() {
55    // Do something with the setting and its members.
56}
57```
58
59*/
60
61use std::fmt::{self, Debug};
62use std::iter::FusedIterator;
63use std::marker::PhantomData;
64
65use crate::data_structures::entity_vec::EntityVec;
66use crate::entity::{Entity, EntityId};
67
68/// A `Vec`-backed map keyed by `EntityId<E>`.
69#[derive(Clone, PartialEq, Eq)]
70pub struct EntityMap<E: Entity, V> {
71    data: Vec<Option<V>>,
72    len: usize,
73    _phantom: PhantomData<E>,
74}
75
76impl<E: Entity, V> EntityMap<E, V> {
77    /// Creates an empty `EntityMap`.
78    #[must_use]
79    pub fn new() -> Self {
80        Self {
81            data: Vec::new(),
82            len: 0,
83            _phantom: PhantomData,
84        }
85    }
86
87    /// Creates an empty `EntityMap` with space for at least `capacity` values.
88    #[must_use]
89    pub fn with_capacity(capacity: usize) -> Self {
90        Self {
91            data: Vec::with_capacity(capacity),
92            len: 0,
93            _phantom: PhantomData,
94        }
95    }
96
97    /// Cheap conversion to an `EntityVec<E, Option<V>>`
98    #[must_use]
99    pub fn into_entity_vec(self) -> EntityVec<E, Option<V>> {
100        self.data.into()
101    }
102
103    /// Returns the number of stored key-value pairs.
104    #[inline]
105    #[must_use]
106    pub fn len(&self) -> usize {
107        self.len
108    }
109
110    /// Returns the capacity of the backing storage.
111    #[inline]
112    #[must_use]
113    pub fn capacity(&self) -> usize {
114        self.data.capacity()
115    }
116
117    /// Returns `true` if the map contains no values.
118    #[inline]
119    #[must_use]
120    pub fn is_empty(&self) -> bool {
121        self.len == 0
122    }
123
124    /// Reserves capacity for at least `additional` more values.
125    pub fn reserve(&mut self, additional: usize) {
126        self.data.reserve(additional);
127    }
128
129    /// Shrinks the backing vector to fit the highest occupied entity ID.
130    pub fn shrink_to_fit(&mut self) {
131        while self.data.last().is_some_and(Option::is_none) {
132            self.data.pop();
133        }
134        self.data.shrink_to_fit();
135    }
136
137    /// Returns `true` if `entity_id` is present in the map.
138    #[must_use]
139    pub fn contains_key(&self, entity_id: EntityId<E>) -> bool {
140        self.get(entity_id).is_some()
141    }
142
143    /// Returns the value for `entity_id`, or `None` if not present.
144    #[must_use]
145    pub fn get(&self, entity_id: EntityId<E>) -> Option<&V> {
146        self.data.get(entity_id.0).and_then(Option::as_ref)
147    }
148
149    /// Returns the value for `entity_id` mutably, or `None` if not present.
150    #[must_use]
151    pub fn get_mut(&mut self, entity_id: EntityId<E>) -> Option<&mut V> {
152        self.data.get_mut(entity_id.0).and_then(Option::as_mut)
153    }
154
155    /// Inserts `value` for `entity_id`, returning the previous value if one existed.
156    pub fn insert(&mut self, entity_id: EntityId<E>, value: V) -> Option<V> {
157        if entity_id.0 >= self.data.len() {
158            self.data.resize_with(entity_id.0 + 1, || None);
159        }
160
161        let slot = &mut self.data[entity_id.0];
162        let previous = slot.replace(value);
163        if previous.is_none() {
164            self.len += 1;
165        }
166        previous
167    }
168
169    /// Returns the value for `entity_id`, inserting `value` if it is not already present.
170    pub fn get_or_insert(&mut self, entity_id: EntityId<E>, value: V) -> &mut V {
171        self.get_or_insert_with(entity_id, || value)
172    }
173
174    /// Returns the value for `entity_id`, inserting a value from `f` if it is not already present.
175    pub fn get_or_insert_with<F>(&mut self, entity_id: EntityId<E>, f: F) -> &mut V
176    where
177        F: FnOnce() -> V,
178    {
179        if entity_id.0 >= self.data.len() {
180            self.data.resize_with(entity_id.0 + 1, || None);
181        }
182
183        let slot = &mut self.data[entity_id.0];
184        if slot.is_none() {
185            *slot = Some(f());
186            self.len += 1;
187        }
188
189        slot.as_mut().unwrap()
190    }
191
192    /// Removes and returns the value for `entity_id`, if present.
193    #[must_use]
194    pub fn remove(&mut self, entity_id: EntityId<E>) -> Option<V> {
195        let removed = self.data.get_mut(entity_id.0).and_then(Option::take);
196        if removed.is_some() {
197            self.len -= 1;
198        }
199        removed
200    }
201
202    /// Clears the map, removing all key-value pairs.
203    pub fn clear(&mut self) {
204        self.data.clear();
205        self.len = 0;
206    }
207
208    /// Returns an iterator over `(EntityId<E>, &V)` pairs.
209    #[must_use]
210    pub fn iter(&self) -> Iter<'_, E, V> {
211        Iter {
212            inner: self.data.iter().enumerate(),
213            remaining: self.len,
214            _phantom: PhantomData,
215        }
216    }
217}
218
219impl<E: Entity, V> Default for EntityMap<E, V> {
220    fn default() -> Self {
221        Self::new()
222    }
223}
224
225impl<E: Entity, V: Debug> Debug for EntityMap<E, V> {
226    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
227        f.debug_map().entries(self.iter()).finish()
228    }
229}
230
231impl<E: Entity, V> Extend<(EntityId<E>, V)> for EntityMap<E, V> {
232    fn extend<I: IntoIterator<Item = (EntityId<E>, V)>>(&mut self, iter: I) {
233        for (entity_id, value) in iter {
234            let _ = self.insert(entity_id, value);
235        }
236    }
237}
238
239impl<E: Entity, V> FromIterator<(EntityId<E>, V)> for EntityMap<E, V> {
240    fn from_iter<I: IntoIterator<Item = (EntityId<E>, V)>>(iter: I) -> Self {
241        let mut map = Self::new();
242        map.extend(iter);
243        map
244    }
245}
246
247impl<'a, E: Entity, V> IntoIterator for &'a EntityMap<E, V> {
248    type Item = (EntityId<E>, &'a V);
249    type IntoIter = Iter<'a, E, V>;
250
251    fn into_iter(self) -> Self::IntoIter {
252        self.iter()
253    }
254}
255
256/// Iterator over `(EntityId<E>, &V)` pairs from an `EntityMap<E, V>`.
257pub struct Iter<'a, E: Entity, V> {
258    inner: std::iter::Enumerate<std::slice::Iter<'a, Option<V>>>,
259    remaining: usize,
260    _phantom: PhantomData<E>,
261}
262
263impl<'a, E: Entity, V> Iterator for Iter<'a, E, V> {
264    type Item = (EntityId<E>, &'a V);
265
266    fn next(&mut self) -> Option<Self::Item> {
267        for (index, value) in self.inner.by_ref() {
268            if let Some(value) = value.as_ref() {
269                self.remaining -= 1;
270                return Some((EntityId::new(index), value));
271            }
272        }
273        None
274    }
275
276    fn size_hint(&self) -> (usize, Option<usize>) {
277        (self.remaining, Some(self.remaining))
278    }
279
280    fn count(self) -> usize {
281        self.remaining
282    }
283
284    fn nth(&mut self, n: usize) -> Option<Self::Item> {
285        if n >= self.remaining {
286            self.remaining = 0;
287            self.inner.by_ref().for_each(drop);
288            return None;
289        }
290
291        let mut skipped = 0;
292        for (index, value) in self.inner.by_ref() {
293            if let Some(value) = value.as_ref() {
294                if skipped == n {
295                    self.remaining -= n + 1;
296                    return Some((EntityId::new(index), value));
297                }
298                skipped += 1;
299            }
300        }
301
302        self.remaining = 0;
303        None
304    }
305}
306
307impl<'a, E: Entity, V> ExactSizeIterator for Iter<'a, E, V> {
308    fn len(&self) -> usize {
309        self.remaining
310    }
311}
312
313impl<'a, E: Entity, V> FusedIterator for Iter<'a, E, V> {}
314
315#[cfg(test)]
316mod tests {
317    use super::EntityMap;
318    use crate::define_entity;
319    use crate::entity::EntityId;
320
321    define_entity!(TestEntity);
322
323    #[test]
324    fn new_is_empty() {
325        let map = EntityMap::<TestEntity, i32>::new();
326
327        assert_eq!(map.len(), 0);
328        assert!(map.is_empty());
329        assert_eq!(map.capacity(), 0);
330    }
331
332    #[test]
333    fn with_capacity_sets_initial_capacity() {
334        let map = EntityMap::<TestEntity, i32>::with_capacity(8);
335
336        assert_eq!(map.len(), 0);
337        assert!(map.capacity() >= 8);
338    }
339
340    #[test]
341    fn insert_and_get_work_for_sparse_ids() {
342        let mut map = EntityMap::<TestEntity, &'static str>::new();
343        let id2 = EntityId::new(2);
344        let id5 = EntityId::new(5);
345
346        assert_eq!(map.insert(id2, "two"), None);
347        assert_eq!(map.insert(id5, "five"), None);
348
349        assert_eq!(map.len(), 2);
350        assert!(!map.is_empty());
351        assert_eq!(map.get(EntityId::new(0)), None);
352        assert_eq!(map.get(id2), Some(&"two"));
353        assert_eq!(map.get(id5), Some(&"five"));
354        assert_eq!(map.get(EntityId::new(6)), None);
355    }
356
357    #[test]
358    fn insert_replaces_existing_value_without_changing_len() {
359        let mut map = EntityMap::<TestEntity, i32>::new();
360        let id = EntityId::new(3);
361
362        assert_eq!(map.insert(id, 10), None);
363        assert_eq!(map.insert(id, 20), Some(10));
364
365        assert_eq!(map.len(), 1);
366        assert_eq!(map.get(id), Some(&20));
367    }
368
369    #[test]
370    fn get_mut_updates_value() {
371        let mut map = EntityMap::<TestEntity, i32>::new();
372        let id = EntityId::new(1);
373        let _ = map.insert(id, 10);
374
375        *map.get_mut(id).unwrap() = 99;
376
377        assert_eq!(map.get(id), Some(&99));
378        assert_eq!(map.get_mut(EntityId::new(7)), None);
379    }
380
381    #[test]
382    fn get_or_insert_inserts_missing_value_and_returns_mutable_reference() {
383        let mut map = EntityMap::<TestEntity, i32>::new();
384        let id = EntityId::new(3);
385
386        let value = map.get_or_insert(id, 10);
387        *value = 15;
388
389        assert_eq!(map.len(), 1);
390        assert_eq!(map.get(id), Some(&15));
391    }
392
393    #[test]
394    fn get_or_insert_does_not_replace_existing_value() {
395        let mut map = EntityMap::<TestEntity, i32>::new();
396        let id = EntityId::new(2);
397        let _ = map.insert(id, 20);
398
399        let value = map.get_or_insert(id, 99);
400
401        assert_eq!(*value, 20);
402        assert_eq!(map.len(), 1);
403        assert_eq!(map.get(id), Some(&20));
404    }
405
406    #[test]
407    fn get_or_insert_with_only_evaluates_closure_for_missing_key() {
408        let mut map = EntityMap::<TestEntity, i32>::new();
409        let missing_id = EntityId::new(1);
410        let existing_id = EntityId::new(4);
411        let _ = map.insert(existing_id, 40);
412        let mut calls = 0;
413
414        let inserted = map.get_or_insert_with(missing_id, || {
415            calls += 1;
416            10
417        });
418        assert_eq!(*inserted, 10);
419
420        let existing = map.get_or_insert_with(existing_id, || {
421            calls += 1;
422            99
423        });
424        assert_eq!(*existing, 40);
425
426        assert_eq!(calls, 1);
427        assert_eq!(map.len(), 2);
428    }
429
430    #[test]
431    fn contains_key_tracks_presence() {
432        let mut map = EntityMap::<TestEntity, i32>::new();
433        let id = EntityId::new(4);
434
435        assert!(!map.contains_key(id));
436        let _ = map.insert(id, 12);
437        assert!(map.contains_key(id));
438        assert!(!map.contains_key(EntityId::new(3)));
439    }
440
441    #[test]
442    fn remove_returns_value_and_decrements_len() {
443        let mut map = EntityMap::<TestEntity, i32>::new();
444        let id1 = EntityId::new(1);
445        let id4 = EntityId::new(4);
446        let _ = map.insert(id1, 10);
447        let _ = map.insert(id4, 40);
448
449        assert_eq!(map.remove(id1), Some(10));
450        assert_eq!(map.remove(id1), None);
451
452        assert_eq!(map.len(), 1);
453        assert_eq!(map.get(id1), None);
454        assert_eq!(map.get(id4), Some(&40));
455    }
456
457    #[test]
458    fn remove_missing_key_returns_none_without_changing_len() {
459        let mut map = EntityMap::<TestEntity, i32>::new();
460        let _ = map.insert(EntityId::new(1), 10);
461        let _ = map.insert(EntityId::new(8), 80);
462
463        assert_eq!(map.remove(EntityId::new(4)), None);
464
465        assert_eq!(map.len(), 2);
466        assert!(!map.is_empty());
467        assert_eq!(map.get(EntityId::new(1)), Some(&10));
468        assert_eq!(map.get(EntityId::new(8)), Some(&80));
469    }
470
471    #[test]
472    fn clear_removes_all_entries() {
473        let mut map = EntityMap::<TestEntity, i32>::new();
474        let _ = map.insert(EntityId::new(0), 1);
475        let _ = map.insert(EntityId::new(2), 2);
476
477        map.clear();
478
479        assert!(map.is_empty());
480        assert_eq!(map.len(), 0);
481        assert_eq!(map.get(EntityId::new(0)), None);
482        assert_eq!(map.get(EntityId::new(2)), None);
483    }
484
485    #[test]
486    fn reserve_and_shrink_to_fit_manage_capacity() {
487        let mut map = EntityMap::<TestEntity, i32>::new();
488        map.reserve(16);
489        assert!(map.capacity() >= 16);
490
491        let _ = map.insert(EntityId::new(10), 10);
492        let _ = map.insert(EntityId::new(20), 20);
493        assert_eq!(map.remove(EntityId::new(20)), Some(20));
494
495        map.shrink_to_fit();
496
497        assert_eq!(map.len(), 1);
498        assert_eq!(map.get(EntityId::new(10)), Some(&10));
499        assert_eq!(map.get(EntityId::new(20)), None);
500        assert!(map.capacity() >= 11);
501    }
502
503    #[test]
504    fn iter_yields_only_present_entries_in_entity_id_order() {
505        let mut map = EntityMap::<TestEntity, &'static str>::new();
506        let _ = map.insert(EntityId::new(3), "three");
507        let _ = map.insert(EntityId::new(0), "zero");
508        let _ = map.insert(EntityId::new(5), "five");
509
510        let items: Vec<_> = map.iter().collect();
511
512        assert_eq!(
513            items,
514            vec![
515                (EntityId::new(0), &"zero"),
516                (EntityId::new(3), &"three"),
517                (EntityId::new(5), &"five"),
518            ]
519        );
520    }
521
522    #[test]
523    fn iter_size_hint_len_and_count_track_remaining_entries() {
524        let mut map = EntityMap::<TestEntity, i32>::new();
525        let _ = map.insert(EntityId::new(1), 10);
526        let _ = map.insert(EntityId::new(4), 40);
527        let _ = map.insert(EntityId::new(7), 70);
528
529        let mut iter = map.iter();
530        assert_eq!(iter.size_hint(), (3, Some(3)));
531        assert_eq!(iter.len(), 3);
532
533        assert_eq!(iter.next(), Some((EntityId::new(1), &10)));
534        assert_eq!(iter.size_hint(), (2, Some(2)));
535        assert_eq!(iter.len(), 2);
536        assert_eq!(iter.count(), 2);
537    }
538
539    #[test]
540    fn iter_nth_skips_missing_entries_and_updates_remaining_len() {
541        let mut map = EntityMap::<TestEntity, i32>::new();
542        let _ = map.insert(EntityId::new(2), 20);
543        let _ = map.insert(EntityId::new(5), 50);
544        let _ = map.insert(EntityId::new(8), 80);
545
546        let mut iter = map.iter();
547        assert_eq!(iter.nth(1), Some((EntityId::new(5), &50)));
548        assert_eq!(iter.len(), 1);
549        assert_eq!(iter.next(), Some((EntityId::new(8), &80)));
550        assert_eq!(iter.next(), None);
551        assert_eq!(iter.next(), None);
552    }
553
554    #[test]
555    fn iter_nth_past_end_returns_none_and_exhausts_iterator() {
556        let mut map = EntityMap::<TestEntity, i32>::new();
557        let _ = map.insert(EntityId::new(1), 10);
558        let _ = map.insert(EntityId::new(3), 30);
559
560        let mut iter = map.iter();
561        assert_eq!(iter.nth(2), None);
562        assert_eq!(iter.len(), 0);
563        assert_eq!(iter.next(), None);
564    }
565
566    #[test]
567    fn debug_formats_like_a_map() {
568        let mut map = EntityMap::<TestEntity, i32>::new();
569        let _ = map.insert(EntityId::new(1), 10);
570        let _ = map.insert(EntityId::new(4), 40);
571
572        assert_eq!(
573            format!("{:?}", map),
574            "{TestEntityId(1): 10, TestEntityId(4): 40}"
575        );
576    }
577
578    #[test]
579    fn default_matches_new() {
580        let map = EntityMap::<TestEntity, i32>::default();
581
582        assert!(map.is_empty());
583        assert_eq!(map.len(), 0);
584    }
585
586    #[test]
587    fn extend_and_from_iter_insert_all_entries() {
588        let entries = vec![(EntityId::new(2), 20), (EntityId::new(5), 50)];
589
590        let mut map = EntityMap::<TestEntity, i32>::new();
591        map.extend(entries.clone());
592        assert_eq!(map.get(EntityId::new(2)), Some(&20));
593        assert_eq!(map.get(EntityId::new(5)), Some(&50));
594
595        let collected = EntityMap::<TestEntity, i32>::from_iter(entries);
596        assert_eq!(collected.get(EntityId::new(2)), Some(&20));
597        assert_eq!(collected.get(EntityId::new(5)), Some(&50));
598        assert_eq!(collected.len(), 2);
599    }
600
601    #[test]
602    fn into_iterator_for_reference_matches_iter() {
603        let mut map = EntityMap::<TestEntity, i32>::new();
604        let _ = map.insert(EntityId::new(0), 10);
605        let _ = map.insert(EntityId::new(2), 20);
606
607        let from_iter: Vec<_> = map.iter().collect();
608        let from_into_iter: Vec<_> = (&map).into_iter().collect();
609
610        assert_eq!(from_iter, from_into_iter);
611    }
612}
ixa/data_structures/entity_map.rs

ixa/data_structures/
entity_map.rs
