1 files changed, 24 insertions, 6 deletions

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(),
+        ))
     }
 }