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

3 files changed, 162 insertions, 64 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 })),
+                    }],
+                },
             ],
         });
 
diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index beb2478..60d4210 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -6,10 +6,13 @@ use std::sync::Arc;
 
 use crate::component::storage::{
     Archetype,
+    ArchetypeEntity,
     ArchetypeRefIter,
+    EntityIter,
     Storage as ComponentStorage,
 };
 use crate::component::{Metadata as ComponentMetadata, Sequence as ComponentSequence};
+use crate::entity::Uid as EntityUid;
 use crate::lock::{Lock, ReadGuard};
 use crate::query::options::Options;
 use crate::system::{
@@ -17,7 +20,7 @@ use crate::system::{
     Param as SystemParam,
     System,
 };
-use crate::{EntityComponent, World, WorldData};
+use crate::{World, WorldData};
 
 pub mod options;
 
@@ -49,13 +52,26 @@ where
             entities: self
                 .component_storage
                 .find_entities(Comps::metadata())
-                .map((|archetype| archetype.components.as_slice()) as ComponentIterMapFn)
+                .map((|archetype| archetype.entities()) as ComponentIterMapFn)
                 .flatten()
-                .filter(|components| OptionsT::entity_filter(*components)),
+                .filter(|entity| OptionsT::entity_filter(entity.components())),
             comps_pd: PhantomData,
         }
     }
 
+    /// Returns the UID of the entity at the given query iteration index.
+    pub fn entity_uid(&self, entity_index: usize) -> Option<EntityUid>
+    {
+        Some(
+            self.component_storage
+                .find_entities(Comps::metadata())
+                .map(|archetype| archetype.entities())
+                .flatten()
+                .nth(entity_index)?
+                .uid(),
+        )
+    }
+
     pub(crate) fn new(component_storage: &'world Arc<Lock<ComponentStorage>>) -> Self
     {
         Self {
@@ -145,9 +161,9 @@ where
     }
 }
 
-type ComponentIterMapFn = for<'a> fn(&'a Archetype) -> &'a [Vec<EntityComponent>];
+type ComponentIterMapFn = for<'a> fn(&'a Archetype) -> EntityIter<'a>;
 
-type ComponentIterFilterFn = for<'a, 'b> fn(&'a &'b Vec<EntityComponent>) -> bool;
+type ComponentIterFilterFn = for<'a, 'b> fn(&'a &'b ArchetypeEntity) -> bool;
 
 pub struct ComponentIter<'world, Comps>
 {
@@ -166,7 +182,9 @@ where
 
     fn next(&mut self) -> Option<Self::Item>
     {
-        Some(Comps::from_components(self.entities.next()?.iter()))
+        Some(Comps::from_components(
+            self.entities.next()?.components().iter(),
+        ))
     }
 }
 
diff --git a/ecs/src/relationship.rs b/ecs/src/relationship.rs
index 9004f0f..964a47f 100644
--- a/ecs/src/relationship.rs
+++ b/ecs/src/relationship.rs
@@ -174,7 +174,7 @@ where
 
         let component = ComponentRefMut::new(
             entity
-                .get(component_index)?
+                .get_component(component_index)?
                 .component
                 .write_nonblock()
                 .unwrap_or_else(|_| {