1 files changed, 46 insertions, 19 deletions

diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index 88ca180..de6c832 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -1,8 +1,11 @@
-use std::any::{Any, TypeId};
+use std::any::{type_name, Any, TypeId};
 use std::collections::HashSet;
+use std::iter::{Flatten, Map};
 use std::marker::PhantomData;
+use std::slice::Iter as SliceIter;
 use std::sync::{Arc, Weak};
 
+use crate::component::storage::Archetype;
 use crate::component::{
     IsOptional as ComponentIsOptional,
     Sequence as ComponentSequence,
@@ -13,7 +16,7 @@ use crate::system::{
     Param as SystemParam,
     System,
 };
-use crate::{ComponentStorage, WorldData};
+use crate::{ComponentStorage, EntityComponent, WorldData};
 
 #[derive(Debug)]
 pub struct Query<'world, Comps>
@@ -34,10 +37,17 @@ where
     where
         'this: 'world,
     {
+        #[cfg(feature = "debug")]
+        tracing::debug!("Searching for {}", type_name::<Comps>());
+
         ComponentIter {
-            component_storage: &self.component_storage,
-            current_entity_index: 0,
-            component_type_ids: Comps::type_ids(),
+            entities: self
+                .component_storage
+                .find_entities(&Comps::type_ids())
+                .unwrap_or_else(|| panic!("Could not find {:?}", type_name::<Comps>()))
+                .iter()
+                .map((|archetype| archetype.components.as_slice()) as ComponentIterMapFn)
+                .flatten(),
             comps_pd: PhantomData,
         }
     }
@@ -118,6 +128,33 @@ where
             component_type_ids: Comps::type_ids(),
         })
     }
+
+    fn handle_pre_run(world_data: &WorldData)
+    {
+        let mut component_storage_lock = world_data
+            .component_storage
+            .write_nonblock()
+            .expect("Failed to acquire read-write component storage lock");
+
+        #[cfg(feature = "debug")]
+        tracing::debug!(
+            "Adding archetypes lookup entry for components: ({})",
+            type_name::<Comps>()
+        );
+
+        component_storage_lock.add_archetype_lookup_entry(
+            &Comps::type_ids()
+                .into_iter()
+                .filter_map(|(component_id, is_optional)| {
+                    if is_optional == ComponentIsOptional::Yes {
+                        return None;
+                    }
+
+                    Some(component_id)
+                })
+                .collect::<Vec<_>>(),
+        );
+    }
 }
 
 /// A entity query containing a weak reference to the world.
@@ -182,11 +219,11 @@ where
     }
 }
 
+type ComponentIterMapFn = for<'a> fn(&'a &'a Archetype) -> &'a [Vec<EntityComponent>];
+
 pub struct ComponentIter<'world, Comps>
 {
-    component_storage: &'world ComponentStorage,
-    current_entity_index: usize,
-    component_type_ids: Vec<(TypeId, ComponentIsOptional)>,
+    entities: Flatten<Map<SliceIter<'world, &'world Archetype>, ComponentIterMapFn>>,
     comps_pd: PhantomData<Comps>,
 }
 
@@ -198,17 +235,7 @@ where
 
     fn next(&mut self) -> Option<Self::Item>
     {
-        let (matching_entity_index, matching_entity_components) =
-            self.component_storage.find_entity_with_components(
-                self.current_entity_index,
-                &self.component_type_ids,
-            )?;
-
-        self.current_entity_index = matching_entity_index + 1;
-
-        Some(Comps::from_components(
-            matching_entity_components.iter().map(|component| component),
-        ))
+        Some(Comps::from_components(self.entities.next()?.iter()))
     }
 }