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    pub fn new() -> Self {
79        Self {
80            data: Vec::new(),
81            len: 0,
82            _phantom: PhantomData,
83        }
84    }
85
86    /// Creates an empty `EntityMap` with space for at least `capacity` values.
87    pub fn with_capacity(capacity: usize) -> Self {
88        Self {
89            data: Vec::with_capacity(capacity),
90            len: 0,
91            _phantom: PhantomData,
92        }
93    }
94
95    /// Cheap conversion to an `EntityVec<E, Option<V>>`
96    pub fn into_entity_vec(self) -> EntityVec<E, Option<V>> {
97        self.data.into()
98    }
99
100    /// Returns the number of stored key-value pairs.
101    #[inline]
102    pub fn len(&self) -> usize {
103        self.len
104    }
105
106    /// Returns the capacity of the backing storage.
107    #[inline]
108    pub fn capacity(&self) -> usize {
109        self.data.capacity()
110    }
111
112    /// Returns `true` if the map contains no values.
113    #[inline]
114    pub fn is_empty(&self) -> bool {
115        self.len == 0
116    }
117
118    /// Reserves capacity for at least `additional` more values.
119    pub fn reserve(&mut self, additional: usize) {
120        self.data.reserve(additional);
121    }
122
123    /// Shrinks the backing vector to fit the highest occupied entity ID.
124    pub fn shrink_to_fit(&mut self) {
125        while self.data.last().is_some_and(Option::is_none) {
126            self.data.pop();
127        }
128        self.data.shrink_to_fit();
129    }
130
131    /// Returns `true` if `entity_id` is present in the map.
132    pub fn contains_key(&self, entity_id: EntityId<E>) -> bool {
133        self.get(entity_id).is_some()
134    }
135
136    /// Returns the value for `entity_id`, or `None` if not present.
137    pub fn get(&self, entity_id: EntityId<E>) -> Option<&V> {
138        self.data.get(entity_id.0).and_then(Option::as_ref)
139    }
140
141    /// Returns the value for `entity_id` mutably, or `None` if not present.
142    pub fn get_mut(&mut self, entity_id: EntityId<E>) -> Option<&mut V> {
143        self.data.get_mut(entity_id.0).and_then(Option::as_mut)
144    }
145
146    /// Inserts `value` for `entity_id`, returning the previous value if one existed.
147    pub fn insert(&mut self, entity_id: EntityId<E>, value: V) -> Option<V> {
148        if entity_id.0 >= self.data.len() {
149            self.data.resize_with(entity_id.0 + 1, || None);
150        }
151
152        let slot = &mut self.data[entity_id.0];
153        let previous = slot.replace(value);
154        if previous.is_none() {
155            self.len += 1;
156        }
157        previous
158    }
159
160    /// Returns the value for `entity_id`, inserting `value` if it is not already present.
161    pub fn get_or_insert(&mut self, entity_id: EntityId<E>, value: V) -> &mut V {
162        self.get_or_insert_with(entity_id, || value)
163    }
164
165    /// Returns the value for `entity_id`, inserting a value from `f` if it is not already present.
166    pub fn get_or_insert_with<F>(&mut self, entity_id: EntityId<E>, f: F) -> &mut V
167    where
168        F: FnOnce() -> V,
169    {
170        if entity_id.0 >= self.data.len() {
171            self.data.resize_with(entity_id.0 + 1, || None);
172        }
173
174        let slot = &mut self.data[entity_id.0];
175        if slot.is_none() {
176            *slot = Some(f());
177            self.len += 1;
178        }
179
180        slot.as_mut().unwrap()
181    }
182
183    /// Removes and returns the value for `entity_id`, if present.
184    pub fn remove(&mut self, entity_id: EntityId<E>) -> Option<V> {
185        let removed = self.data.get_mut(entity_id.0).and_then(Option::take);
186        if removed.is_some() {
187            self.len -= 1;
188        }
189        removed
190    }
191
192    /// Clears the map, removing all key-value pairs.
193    pub fn clear(&mut self) {
194        self.data.clear();
195        self.len = 0;
196    }
197
198    /// Returns an iterator over `(EntityId<E>, &V)` pairs.
199    pub fn iter(&self) -> Iter<'_, E, V> {
200        Iter {
201            inner: self.data.iter().enumerate(),
202            remaining: self.len,
203            _phantom: PhantomData,
204        }
205    }
206}
207
208impl<E: Entity, V> Default for EntityMap<E, V> {
209    fn default() -> Self {
210        Self::new()
211    }
212}
213
214impl<E: Entity, V: Debug> Debug for EntityMap<E, V> {
215    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
216        f.debug_map().entries(self.iter()).finish()
217    }
218}
219
220impl<E: Entity, V> Extend<(EntityId<E>, V)> for EntityMap<E, V> {
221    fn extend<I: IntoIterator<Item = (EntityId<E>, V)>>(&mut self, iter: I) {
222        for (entity_id, value) in iter {
223            let _ = self.insert(entity_id, value);
224        }
225    }
226}
227
228impl<E: Entity, V> FromIterator<(EntityId<E>, V)> for EntityMap<E, V> {
229    fn from_iter<I: IntoIterator<Item = (EntityId<E>, V)>>(iter: I) -> Self {
230        let mut map = Self::new();
231        map.extend(iter);
232        map
233    }
234}
235
236impl<'a, E: Entity, V> IntoIterator for &'a EntityMap<E, V> {
237    type Item = (EntityId<E>, &'a V);
238    type IntoIter = Iter<'a, E, V>;
239
240    fn into_iter(self) -> Self::IntoIter {
241        self.iter()
242    }
243}
244
245/// Iterator over `(EntityId<E>, &V)` pairs from an `EntityMap<E, V>`.
246pub struct Iter<'a, E: Entity, V> {
247    inner: std::iter::Enumerate<std::slice::Iter<'a, Option<V>>>,
248    remaining: usize,
249    _phantom: PhantomData<E>,
250}
251
252impl<'a, E: Entity, V> Iterator for Iter<'a, E, V> {
253    type Item = (EntityId<E>, &'a V);
254
255    fn next(&mut self) -> Option<Self::Item> {
256        for (index, value) in self.inner.by_ref() {
257            if let Some(value) = value.as_ref() {
258                self.remaining -= 1;
259                return Some((EntityId::new(index), value));
260            }
261        }
262        None
263    }
264
265    fn size_hint(&self) -> (usize, Option<usize>) {
266        (self.remaining, Some(self.remaining))
267    }
268
269    fn count(self) -> usize {
270        self.remaining
271    }
272
273    fn nth(&mut self, n: usize) -> Option<Self::Item> {
274        if n >= self.remaining {
275            self.remaining = 0;
276            self.inner.by_ref().for_each(drop);
277            return None;
278        }
279
280        let mut skipped = 0;
281        for (index, value) in self.inner.by_ref() {
282            if let Some(value) = value.as_ref() {
283                if skipped == n {
284                    self.remaining -= n + 1;
285                    return Some((EntityId::new(index), value));
286                }
287                skipped += 1;
288            }
289        }
290
291        self.remaining = 0;
292        None
293    }
294}
295
296impl<'a, E: Entity, V> ExactSizeIterator for Iter<'a, E, V> {
297    fn len(&self) -> usize {
298        self.remaining
299    }
300}
301
302impl<'a, E: Entity, V> FusedIterator for Iter<'a, E, V> {}
303
304#[cfg(test)]
305mod tests {
306    use super::EntityMap;
307    use crate::define_entity;
308    use crate::entity::EntityId;
309
310    define_entity!(TestEntity);
311
312    #[test]
313    fn new_is_empty() {
314        let map = EntityMap::<TestEntity, i32>::new();
315
316        assert_eq!(map.len(), 0);
317        assert!(map.is_empty());
318        assert_eq!(map.capacity(), 0);
319    }
320
321    #[test]
322    fn with_capacity_sets_initial_capacity() {
323        let map = EntityMap::<TestEntity, i32>::with_capacity(8);
324
325        assert_eq!(map.len(), 0);
326        assert!(map.capacity() >= 8);
327    }
328
329    #[test]
330    fn insert_and_get_work_for_sparse_ids() {
331        let mut map = EntityMap::<TestEntity, &'static str>::new();
332        let id2 = EntityId::new(2);
333        let id5 = EntityId::new(5);
334
335        assert_eq!(map.insert(id2, "two"), None);
336        assert_eq!(map.insert(id5, "five"), None);
337
338        assert_eq!(map.len(), 2);
339        assert!(!map.is_empty());
340        assert_eq!(map.get(EntityId::new(0)), None);
341        assert_eq!(map.get(id2), Some(&"two"));
342        assert_eq!(map.get(id5), Some(&"five"));
343        assert_eq!(map.get(EntityId::new(6)), None);
344    }
345
346    #[test]
347    fn insert_replaces_existing_value_without_changing_len() {
348        let mut map = EntityMap::<TestEntity, i32>::new();
349        let id = EntityId::new(3);
350
351        assert_eq!(map.insert(id, 10), None);
352        assert_eq!(map.insert(id, 20), Some(10));
353
354        assert_eq!(map.len(), 1);
355        assert_eq!(map.get(id), Some(&20));
356    }
357
358    #[test]
359    fn get_mut_updates_value() {
360        let mut map = EntityMap::<TestEntity, i32>::new();
361        let id = EntityId::new(1);
362        let _ = map.insert(id, 10);
363
364        *map.get_mut(id).unwrap() = 99;
365
366        assert_eq!(map.get(id), Some(&99));
367        assert_eq!(map.get_mut(EntityId::new(7)), None);
368    }
369
370    #[test]
371    fn get_or_insert_inserts_missing_value_and_returns_mutable_reference() {
372        let mut map = EntityMap::<TestEntity, i32>::new();
373        let id = EntityId::new(3);
374
375        let value = map.get_or_insert(id, 10);
376        *value = 15;
377
378        assert_eq!(map.len(), 1);
379        assert_eq!(map.get(id), Some(&15));
380    }
381
382    #[test]
383    fn get_or_insert_does_not_replace_existing_value() {
384        let mut map = EntityMap::<TestEntity, i32>::new();
385        let id = EntityId::new(2);
386        let _ = map.insert(id, 20);
387
388        let value = map.get_or_insert(id, 99);
389
390        assert_eq!(*value, 20);
391        assert_eq!(map.len(), 1);
392        assert_eq!(map.get(id), Some(&20));
393    }
394
395    #[test]
396    fn get_or_insert_with_only_evaluates_closure_for_missing_key() {
397        let mut map = EntityMap::<TestEntity, i32>::new();
398        let missing_id = EntityId::new(1);
399        let existing_id = EntityId::new(4);
400        let _ = map.insert(existing_id, 40);
401        let mut calls = 0;
402
403        let inserted = map.get_or_insert_with(missing_id, || {
404            calls += 1;
405            10
406        });
407        assert_eq!(*inserted, 10);
408
409        let existing = map.get_or_insert_with(existing_id, || {
410            calls += 1;
411            99
412        });
413        assert_eq!(*existing, 40);
414
415        assert_eq!(calls, 1);
416        assert_eq!(map.len(), 2);
417    }
418
419    #[test]
420    fn contains_key_tracks_presence() {
421        let mut map = EntityMap::<TestEntity, i32>::new();
422        let id = EntityId::new(4);
423
424        assert!(!map.contains_key(id));
425        let _ = map.insert(id, 12);
426        assert!(map.contains_key(id));
427        assert!(!map.contains_key(EntityId::new(3)));
428    }
429
430    #[test]
431    fn remove_returns_value_and_decrements_len() {
432        let mut map = EntityMap::<TestEntity, i32>::new();
433        let id1 = EntityId::new(1);
434        let id4 = EntityId::new(4);
435        let _ = map.insert(id1, 10);
436        let _ = map.insert(id4, 40);
437
438        assert_eq!(map.remove(id1), Some(10));
439        assert_eq!(map.remove(id1), None);
440
441        assert_eq!(map.len(), 1);
442        assert_eq!(map.get(id1), None);
443        assert_eq!(map.get(id4), Some(&40));
444    }
445
446    #[test]
447    fn remove_missing_key_returns_none_without_changing_len() {
448        let mut map = EntityMap::<TestEntity, i32>::new();
449        let _ = map.insert(EntityId::new(1), 10);
450        let _ = map.insert(EntityId::new(8), 80);
451
452        assert_eq!(map.remove(EntityId::new(4)), None);
453
454        assert_eq!(map.len(), 2);
455        assert!(!map.is_empty());
456        assert_eq!(map.get(EntityId::new(1)), Some(&10));
457        assert_eq!(map.get(EntityId::new(8)), Some(&80));
458    }
459
460    #[test]
461    fn clear_removes_all_entries() {
462        let mut map = EntityMap::<TestEntity, i32>::new();
463        let _ = map.insert(EntityId::new(0), 1);
464        let _ = map.insert(EntityId::new(2), 2);
465
466        map.clear();
467
468        assert!(map.is_empty());
469        assert_eq!(map.len(), 0);
470        assert_eq!(map.get(EntityId::new(0)), None);
471        assert_eq!(map.get(EntityId::new(2)), None);
472    }
473
474    #[test]
475    fn reserve_and_shrink_to_fit_manage_capacity() {
476        let mut map = EntityMap::<TestEntity, i32>::new();
477        map.reserve(16);
478        assert!(map.capacity() >= 16);
479
480        let _ = map.insert(EntityId::new(10), 10);
481        let _ = map.insert(EntityId::new(20), 20);
482        assert_eq!(map.remove(EntityId::new(20)), Some(20));
483
484        map.shrink_to_fit();
485
486        assert_eq!(map.len(), 1);
487        assert_eq!(map.get(EntityId::new(10)), Some(&10));
488        assert_eq!(map.get(EntityId::new(20)), None);
489        assert!(map.capacity() >= 11);
490    }
491
492    #[test]
493    fn iter_yields_only_present_entries_in_entity_id_order() {
494        let mut map = EntityMap::<TestEntity, &'static str>::new();
495        let _ = map.insert(EntityId::new(3), "three");
496        let _ = map.insert(EntityId::new(0), "zero");
497        let _ = map.insert(EntityId::new(5), "five");
498
499        let items: Vec<_> = map.iter().collect();
500
501        assert_eq!(
502            items,
503            vec![
504                (EntityId::new(0), &"zero"),
505                (EntityId::new(3), &"three"),
506                (EntityId::new(5), &"five"),
507            ]
508        );
509    }
510
511    #[test]
512    fn iter_size_hint_len_and_count_track_remaining_entries() {
513        let mut map = EntityMap::<TestEntity, i32>::new();
514        let _ = map.insert(EntityId::new(1), 10);
515        let _ = map.insert(EntityId::new(4), 40);
516        let _ = map.insert(EntityId::new(7), 70);
517
518        let mut iter = map.iter();
519        assert_eq!(iter.size_hint(), (3, Some(3)));
520        assert_eq!(iter.len(), 3);
521
522        assert_eq!(iter.next(), Some((EntityId::new(1), &10)));
523        assert_eq!(iter.size_hint(), (2, Some(2)));
524        assert_eq!(iter.len(), 2);
525        assert_eq!(iter.count(), 2);
526    }
527
528    #[test]
529    fn iter_nth_skips_missing_entries_and_updates_remaining_len() {
530        let mut map = EntityMap::<TestEntity, i32>::new();
531        let _ = map.insert(EntityId::new(2), 20);
532        let _ = map.insert(EntityId::new(5), 50);
533        let _ = map.insert(EntityId::new(8), 80);
534
535        let mut iter = map.iter();
536        assert_eq!(iter.nth(1), Some((EntityId::new(5), &50)));
537        assert_eq!(iter.len(), 1);
538        assert_eq!(iter.next(), Some((EntityId::new(8), &80)));
539        assert_eq!(iter.next(), None);
540        assert_eq!(iter.next(), None);
541    }
542
543    #[test]
544    fn iter_nth_past_end_returns_none_and_exhausts_iterator() {
545        let mut map = EntityMap::<TestEntity, i32>::new();
546        let _ = map.insert(EntityId::new(1), 10);
547        let _ = map.insert(EntityId::new(3), 30);
548
549        let mut iter = map.iter();
550        assert_eq!(iter.nth(2), None);
551        assert_eq!(iter.len(), 0);
552        assert_eq!(iter.next(), None);
553    }
554
555    #[test]
556    fn debug_formats_like_a_map() {
557        let mut map = EntityMap::<TestEntity, i32>::new();
558        let _ = map.insert(EntityId::new(1), 10);
559        let _ = map.insert(EntityId::new(4), 40);
560
561        assert_eq!(
562            format!("{:?}", map),
563            "{TestEntityId(1): 10, TestEntityId(4): 40}"
564        );
565    }
566
567    #[test]
568    fn default_matches_new() {
569        let map = EntityMap::<TestEntity, i32>::default();
570
571        assert!(map.is_empty());
572        assert_eq!(map.len(), 0);
573    }
574
575    #[test]
576    fn extend_and_from_iter_insert_all_entries() {
577        let entries = vec![(EntityId::new(2), 20), (EntityId::new(5), 50)];
578
579        let mut map = EntityMap::<TestEntity, i32>::new();
580        map.extend(entries.clone());
581        assert_eq!(map.get(EntityId::new(2)), Some(&20));
582        assert_eq!(map.get(EntityId::new(5)), Some(&50));
583
584        let collected = EntityMap::<TestEntity, i32>::from_iter(entries);
585        assert_eq!(collected.get(EntityId::new(2)), Some(&20));
586        assert_eq!(collected.get(EntityId::new(5)), Some(&50));
587        assert_eq!(collected.len(), 2);
588    }
589
590    #[test]
591    fn into_iterator_for_reference_matches_iter() {
592        let mut map = EntityMap::<TestEntity, i32>::new();
593        let _ = map.insert(EntityId::new(0), 10);
594        let _ = map.insert(EntityId::new(2), 20);
595
596        let from_iter: Vec<_> = map.iter().collect();
597        let from_into_iter: Vec<_> = (&map).into_iter().collect();
598
599        assert_eq!(from_iter, from_into_iter);
600    }
601}
ixa/data_structures/entity_map.rs

ixa/data_structures/
entity_map.rs
