feat(ecs): add query function to get entity UID by index

1 files changed, 137 insertions, 57 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 44addb7..2bd3826 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -96,8 +96,11 @@ impl Storage
             )
             .expect("Archetype is gone");
 
-        archetype.components.push(
-            components
+        let entity_uid = EntityUid::new_unique();
+
+        archetype.entities.push(ArchetypeEntity {
+            uid: entity_uid,
+            components: components
                 .into_iter()
                 .map(|component| EntityComponent {
                     id: component.id(),
@@ -105,13 +108,11 @@ impl Storage
                     component: Lock::new(component),
                 })
                 .collect(),
-        );
-
-        let entity_uid = EntityUid::new_unique();
+        });
 
         archetype
             .entity_lookup
-            .insert(entity_uid, archetype.components.len() - 1);
+            .insert(entity_uid, archetype.entities.len() - 1);
 
         (archetype_id, entity_uid)
     }
@@ -180,7 +181,7 @@ pub struct Archetype
 {
     component_ids: HashMap<ComponentId, usize>,
     entity_lookup: HashMap<EntityUid, usize>,
-    pub components: Vec<Vec<EntityComponent>>,
+    entities: Vec<ArchetypeEntity>,
 }
 
 impl Archetype
@@ -194,7 +195,7 @@ impl Archetype
                 .map(|(index, component_id)| (component_id, index))
                 .collect(),
             entity_lookup: HashMap::new(),
-            components: Vec::new(),
+            entities: Vec::new(),
         }
     }
 
@@ -212,11 +213,16 @@ impl Archetype
         }
     }
 
-    pub fn get_entity(&self, entity_uid: EntityUid) -> Option<&[EntityComponent]>
+    pub fn get_entity(&self, entity_uid: EntityUid) -> Option<&ArchetypeEntity>
     {
         let entity_index = *self.entity_lookup.get(&entity_uid)?;
 
-        Some(&self.components.get(entity_index)?)
+        self.entities.get(entity_index)
+    }
+
+    pub fn entities(&self) -> EntityIter<'_>
+    {
+        EntityIter { iter: self.entities.iter() }
     }
 
     pub fn get_index_for_component(&self, component_id: &ComponentId) -> Option<usize>
@@ -226,6 +232,31 @@ impl Archetype
 }
 
 #[derive(Debug)]
+pub struct ArchetypeEntity
+{
+    uid: EntityUid,
+    components: Vec<EntityComponent>,
+}
+
+impl ArchetypeEntity
+{
+    pub fn uid(&self) -> EntityUid
+    {
+        self.uid
+    }
+
+    pub fn components(&self) -> &[EntityComponent]
+    {
+        &self.components
+    }
+
+    pub fn get_component(&self, index: usize) -> Option<&EntityComponent>
+    {
+        self.components.get(index)
+    }
+}
+
+#[derive(Debug)]
 pub struct ArchetypeRefIter<'component_storage>
 {
     inner: SliceIter<'component_storage, usize>,
@@ -256,6 +287,22 @@ impl<'component_storage> Iterator for ArchetypeRefIter<'component_storage>
     }
 }
 
+#[derive(Debug)]
+pub struct EntityIter<'archetype>
+{
+    iter: SliceIter<'archetype, ArchetypeEntity>,
+}
+
+impl<'archetype> Iterator for EntityIter<'archetype>
+{
+    type Item = &'archetype ArchetypeEntity;
+
+    fn next(&mut self) -> Option<Self::Item>
+    {
+        self.iter.next()
+    }
+}
+
 #[cfg(test)]
 mod tests
 {
@@ -265,6 +312,7 @@ mod tests
 
     use super::{Archetype, Storage};
     use crate::archetype::Id as ArchetypeId;
+    use crate::component::storage::ArchetypeEntity;
     use crate::component::{
         Id as ComponentId,
         IsOptional as ComponentIsOptional,
@@ -321,14 +369,14 @@ mod tests
         assert_eq!(archetype.component_ids.len(), 2);
 
         // One entity
-        assert_eq!(archetype.components.len(), 1);
+        assert_eq!(archetype.entities.len(), 1);
 
         let entity_components = archetype
-            .components
+            .entities
             .first()
             .expect("Expected a entity in archetype");
 
-        assert_eq!(entity_components.len(), 2);
+        assert_eq!(entity_components.components.len(), 2);
 
         assert_eq!(component_storage.archetype_lookup.len(), 1);
 
@@ -362,6 +410,10 @@ mod tests
     {
         let mut component_storage = Storage::default();
 
+        let entity_uid_a = EntityUid::new_unique();
+        let entity_uid_b = EntityUid::new_unique();
+        let entity_uid_c = EntityUid::new_unique();
+
         component_storage.archetypes.push(Archetype {
             component_ids: HashMap::from([
                 (ComponentId::of::<IronBoots>(), 0),
@@ -369,29 +421,45 @@ mod tests
                 (ComponentId::of::<Hookshot>(), 2),
             ]),
             entity_lookup: HashMap::from([
-                (EntityUid::new_unique(), 0),
-                (EntityUid::new_unique(), 1),
-                (EntityUid::new_unique(), 2),
+                (entity_uid_a, 0),
+                (entity_uid_b, 1),
+                (entity_uid_c, 2),
             ]),
-            components: vec![
-                vec![EntityComponent {
-                    id: ComponentId::of::<IronBoots>(),
-                    name: "Iron boots",
-                    component: Lock::new(Box::new(IronBoots)),
-                }],
-                vec![EntityComponent {
-                    id: ComponentId::of::<HealthPotion>(),
-                    name: "Health potion",
-                    component: Lock::new(Box::new(HealthPotion { _hp_restoration: 20 })),
-                }],
-                vec![EntityComponent {
-                    id: ComponentId::of::<Hookshot>(),
-                    name: "Hookshot",
-                    component: Lock::new(Box::new(Hookshot { _range: 67 })),
-                }],
+            entities: vec![
+                ArchetypeEntity {
+                    uid: entity_uid_a,
+                    components: vec![EntityComponent {
+                        id: ComponentId::of::<IronBoots>(),
+                        name: "Iron boots",
+                        component: Lock::new(Box::new(IronBoots)),
+                    }],
+                },
+                ArchetypeEntity {
+                    uid: entity_uid_b,
+                    components: vec![EntityComponent {
+                        id: ComponentId::of::<HealthPotion>(),
+                        name: "Health potion",
+                        component: Lock::new(Box::new(HealthPotion {
+                            _hp_restoration: 20,
+                        })),
+                    }],
+                },
+                ArchetypeEntity {
+                    uid: entity_uid_c,
+                    components: vec![EntityComponent {
+                        id: ComponentId::of::<Hookshot>(),
+                        name: "Hookshot",
+                        component: Lock::new(Box::new(Hookshot { _range: 67 })),
+                    }],
+                },
             ],
         });
 
+        let entity_uid_d = EntityUid::new_unique();
+        let entity_uid_e = EntityUid::new_unique();
+        let entity_uid_f = EntityUid::new_unique();
+        let entity_uid_g = EntityUid::new_unique();
+
         component_storage.archetypes.push(Archetype {
             component_ids: HashMap::from([
                 (ComponentId::of::<DekuNut>(), 0),
@@ -400,32 +468,44 @@ mod tests
                 (ComponentId::of::<Hookshot>(), 3),
             ]),
             entity_lookup: HashMap::from([
-                (EntityUid::new_unique(), 0),
-                (EntityUid::new_unique(), 1),
-                (EntityUid::new_unique(), 2),
-                (EntityUid::new_unique(), 3),
+                (entity_uid_d, 0),
+                (entity_uid_e, 1),
+                (entity_uid_f, 2),
+                (entity_uid_g, 3),
             ]),
-            components: vec![
-                vec![EntityComponent {
-                    id: ComponentId::of::<DekuNut>(),
-                    name: "Deku nut",
-                    component: Lock::new(Box::new(DekuNut { _throwing_damage: 5 })),
-                }],
-                vec![EntityComponent {
-                    id: ComponentId::of::<IronBoots>(),
-                    name: "Iron boots",
-                    component: Lock::new(Box::new(IronBoots)),
-                }],
-                vec![EntityComponent {
-                    id: ComponentId::of::<Bow>(),
-                    name: "Bow",
-                    component: Lock::new(Box::new(Bow { _damage: 20 })),
-                }],
-                vec![EntityComponent {
-                    id: ComponentId::of::<Hookshot>(),
-                    name: "Hookshot",
-                    component: Lock::new(Box::new(Hookshot { _range: 67 })),
-                }],
+            entities: vec![
+                ArchetypeEntity {
+                    uid: entity_uid_d,
+                    components: vec![EntityComponent {
+                        id: ComponentId::of::<DekuNut>(),
+                        name: "Deku nut",
+                        component: Lock::new(Box::new(DekuNut { _throwing_damage: 5 })),
+                    }],
+                },
+                ArchetypeEntity {
+                    uid: entity_uid_e,
+                    components: vec![EntityComponent {
+                        id: ComponentId::of::<IronBoots>(),
+                        name: "Iron boots",
+                        component: Lock::new(Box::new(IronBoots)),
+                    }],
+                },
+                ArchetypeEntity {
+                    uid: entity_uid_f,
+                    components: vec![EntityComponent {
+                        id: ComponentId::of::<Bow>(),
+                        name: "Bow",
+                        component: Lock::new(Box::new(Bow { _damage: 20 })),
+                    }],
+                },
+                ArchetypeEntity {
+                    uid: entity_uid_g,
+                    components: vec![EntityComponent {
+                        id: ComponentId::of::<Hookshot>(),
+                        name: "Hookshot",
+                        component: Lock::new(Box::new(Hookshot { _range: 67 })),
+                    }],
+                },
             ],
         });