1 files changed, 105 insertions, 25 deletions

diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index 4e3e4a1..3c59554 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -5,7 +5,7 @@ use std::collections::{HashMap, HashSet};
 use std::fmt::Debug;
 use std::hash::Hash;
 use std::marker::PhantomData;
-use std::slice::IterMut as SliceIterMut;
+use std::slice::{Iter as SliceIter, IterMut as SliceIterMut};
 
 use crate::component::{Component, Sequence as ComponentSequence};
 use crate::system::{
@@ -128,9 +128,18 @@ impl<'world, Comps> Query<'world, Comps>
 where
     Comps: ComponentSequence,
 {
-    pub fn iter_mut(&mut self) -> QueryComponentIter<Comps>
+    pub fn iter(&self) -> QueryComponentIter<Comps>
     {
         QueryComponentIter {
+            entity_iter: self.component_storage.entities.iter(),
+            component_type_ids: Comps::type_ids(),
+            comps_pd: PhantomData,
+        }
+    }
+
+    pub fn iter_mut(&mut self) -> QueryComponentMutIter<Comps>
+    {
+        QueryComponentMutIter {
             entity_iter: self.component_storage.entities.iter_mut(),
             component_type_ids: Comps::type_ids(),
             comps_pd: PhantomData,
@@ -138,11 +147,24 @@ where
     }
 }
 
-impl<'world, Comps> IntoIterator for &'world mut Query<'world, Comps>
+impl<'world, Comps> IntoIterator for &'world Query<'world, Comps>
 where
-    Comps: ComponentSequence + 'world,
+    Comps: ComponentSequence,
 {
     type IntoIter = QueryComponentIter<'world, Comps>;
+    type Item = Comps::Refs<'world>;
+
+    fn into_iter(self) -> Self::IntoIter
+    {
+        self.iter()
+    }
+}
+
+impl<'world, Comps> IntoIterator for &'world mut Query<'world, Comps>
+where
+    Comps: ComponentSequence,
+{
+    type IntoIter = QueryComponentMutIter<'world, Comps>;
     type Item = Comps::MutRefs<'world>;
 
     fn into_iter(self) -> Self::IntoIter
@@ -217,7 +239,7 @@ impl QueryComponentIds
 
 pub struct QueryComponentIter<'world, Comps>
 {
-    entity_iter: SliceIterMut<'world, Entity>,
+    entity_iter: SliceIter<'world, Entity>,
     component_type_ids: Vec<TypeId>,
     comps_pd: PhantomData<Comps>,
 }
@@ -226,33 +248,91 @@ impl<'world, Comps> Iterator for QueryComponentIter<'world, Comps>
 where
     Comps: ComponentSequence + 'world,
 {
+    type Item = Comps::Refs<'world>;
+
+    fn next(&mut self) -> Option<Self::Item>
+    {
+        let matching_entity = find_entity_with_components::<&Entity>(
+            &mut self.entity_iter,
+            &self.component_type_ids,
+        )?;
+
+        Some(Comps::from_components(&matching_entity.components))
+    }
+}
+
+pub struct QueryComponentMutIter<'world, Comps>
+{
+    entity_iter: SliceIterMut<'world, Entity>,
+    component_type_ids: Vec<TypeId>,
+    comps_pd: PhantomData<Comps>,
+}
+
+impl<'world, Comps> Iterator for QueryComponentMutIter<'world, Comps>
+where
+    Comps: ComponentSequence + 'world,
+{
     type Item = Comps::MutRefs<'world>;
 
     fn next(&mut self) -> Option<Self::Item>
     {
-        // TODO: This is a really dumb and slow way to do this. Refactor the world
-        // to store components in archetypes
-        let entity =
-            self.entity_iter.find(|entity| {
-                let entity_components: HashSet<_> = entity
-                    .components
-                    .iter()
-                    .map(|component| component.as_ref().type_id())
-                    .collect();
-
-                if self.component_type_ids.iter().all(|component_type_id| {
-                    entity_components.contains(component_type_id)
-                }) {
-                    return true;
-                }
-
-                false
-            })?;
-
-        Some(Comps::from_components(&mut entity.components))
+        let matching_entity = find_entity_with_components::<&mut Entity>(
+            &mut self.entity_iter,
+            &self.component_type_ids,
+        )?;
+
+        Some(Comps::from_components_mut(&mut matching_entity.components))
+    }
+}
+
+trait EntityRef
+{
+    fn components(&self) -> &[Box<dyn Component>];
+}
+
+impl EntityRef for &Entity
+{
+    fn components(&self) -> &[Box<dyn Component>]
+    {
+        &self.components
     }
 }
 
+impl EntityRef for &mut Entity
+{
+    fn components(&self) -> &[Box<dyn Component>]
+    {
+        &self.components
+    }
+}
+
+fn find_entity_with_components<EntityRefT>(
+    entity_iter: &mut impl Iterator<Item = EntityRefT>,
+    component_type_ids: &[TypeId],
+) -> Option<EntityRefT>
+where
+    EntityRefT: EntityRef,
+{
+    // TODO: This is a really dumb and slow way to do this. Refactor the world
+    // to store components in archetypes
+    entity_iter.find(|entity| {
+        let entity_components: HashSet<_> = entity
+            .components()
+            .iter()
+            .map(|component| component.as_ref().type_id())
+            .collect();
+
+        if component_type_ids
+            .iter()
+            .all(|component_type_id| entity_components.contains(component_type_id))
+        {
+            return true;
+        }
+
+        false
+    })
+}
+
 #[derive(Debug)]
 pub struct ComponentStorage
 {