Skip to main content

libdqg/ecs/
storage.rs

1use super::entity::Entity;
2
3/// Sparse component storage keyed by [`Entity`]. Indexed directly by the entity's index, with
4/// the stored generation checked on every access so a stale `Entity` (held past a despawn) never
5/// reads or writes whatever new component ends up in the recycled slot.
6#[derive(Clone)]
7pub struct ComponentStore<T> {
8    slots: Vec<Option<(u32, T)>>,
9}
10
11impl<T> Default for ComponentStore<T> {
12    // Written by hand rather than `#[derive(Default)]`: the derive would add a `T: Default`
13    // bound even though an empty `Vec` needs no such thing from its element type.
14    fn default() -> Self {
15        Self::new()
16    }
17}
18
19impl<T> ComponentStore<T> {
20    pub fn new() -> Self {
21        Self { slots: Vec::new() }
22    }
23
24    pub fn insert(&mut self, entity: Entity, value: T) {
25        let index = entity.index as usize;
26        if index >= self.slots.len() {
27            self.slots.resize_with(index + 1, || None);
28        }
29        self.slots[index] = Some((entity.generation, value));
30    }
31
32    pub fn remove(&mut self, entity: Entity) -> Option<T> {
33        let slot = self.slots.get_mut(entity.index as usize)?;
34        if slot.as_ref().map(|(generation, _)| *generation) == Some(entity.generation) {
35            slot.take().map(|(_, value)| value)
36        } else {
37            None
38        }
39    }
40
41    pub fn get(&self, entity: Entity) -> Option<&T> {
42        self.slots.get(entity.index as usize)?.as_ref().and_then(|(generation, value)| {
43            (*generation == entity.generation).then_some(value)
44        })
45    }
46
47    pub fn get_mut(&mut self, entity: Entity) -> Option<&mut T> {
48        self.slots.get_mut(entity.index as usize)?.as_mut().and_then(|(generation, value)| {
49            (*generation == entity.generation).then_some(value)
50        })
51    }
52
53    pub fn iter(&self) -> impl Iterator<Item = (Entity, &T)> {
54        self.slots.iter().enumerate().filter_map(|(i, slot)| {
55            slot.as_ref().map(|(generation, value)| (Entity { index: i as u32, generation: *generation }, value))
56        })
57    }
58
59    pub fn iter_mut(&mut self) -> impl Iterator<Item = (Entity, &mut T)> {
60        self.slots.iter_mut().enumerate().filter_map(|(i, slot)| {
61            slot.as_mut().map(|(generation, value)| (Entity { index: i as u32, generation: *generation }, value))
62        })
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69    use crate::ecs::entity::EntityAllocator;
70
71    #[test]
72    fn stale_entity_cannot_read_recycled_slot() {
73        let mut allocator = EntityAllocator::new();
74        let mut store = ComponentStore::new();
75
76        let a = allocator.spawn();
77        store.insert(a, "a");
78        allocator.despawn(a);
79
80        let b = allocator.spawn();
81        store.insert(b, "b");
82
83        assert_eq!(store.get(a), None);
84        assert_eq!(store.get(b), Some(&"b"));
85    }
86
87    #[test]
88    fn insert_then_get_round_trips() {
89        let mut allocator = EntityAllocator::new();
90        let mut store = ComponentStore::new();
91        let a = allocator.spawn();
92        store.insert(a, 42);
93        assert_eq!(store.get(a), Some(&42));
94    }
95}