2 files changed, 116 insertions, 64 deletions

diff --git a/ecs/src/query/flexible.rs b/ecs/src/query/flexible.rs
index 7d24dc9..add30b0 100644
--- a/ecs/src/query/flexible.rs
+++ b/ecs/src/query/flexible.rs
@@ -1,26 +1,21 @@
 //! Low-level querying.
 use std::iter::{repeat_n, FlatMap, RepeatN, Zip};
 
-use crate::component::storage::archetype::{Archetype, ArchetypeEntity, EntityIter};
-use crate::component::storage::{
-    ArchetypeRefIter,
-    ArchetypeSearchTerms,
-    Storage as ComponentStorage,
-};
-use crate::lock::ReadGuard;
+use crate::component::storage::archetype::{Archetype, EntityIter};
+use crate::component::storage::{ArchetypeRefIter, ArchetypeSearchTerms};
+use crate::entity::Handle as EntityHandle;
 use crate::query::Terms;
-use crate::uid::Uid;
-use crate::{EntityComponent, World};
+use crate::World;
 
 /// Low-level entity query structure.
 #[derive(Debug)]
-pub struct Query<'world, 'terms>
+pub struct Query<'world, const MAX_TERM_CNT: usize>
 {
-    component_storage: ReadGuard<'world, ComponentStorage>,
-    terms: Terms<'terms>,
+    world: &'world World,
+    terms: Terms<MAX_TERM_CNT>,
 }
 
-impl<'world, 'terms> Query<'world, 'terms>
+impl<'world, const MAX_TERM_CNT: usize> Query<'world, MAX_TERM_CNT>
 {
     /// Iterates over the entities matching this query.
     #[must_use]
@@ -28,6 +23,8 @@ impl<'world, 'terms> Query<'world, 'terms>
     {
         Iter {
             iter: self
+                .world
+                .data
                 .component_storage
                 .search_archetypes(ArchetypeSearchTerms {
                     required_components: &self.terms.required_components,
@@ -42,59 +39,37 @@ impl<'world, 'terms> Query<'world, 'terms>
         }
     }
 
-    pub(crate) fn new(world: &'world World, terms: Terms<'terms>) -> Self
+    pub(crate) fn new(world: &'world World, terms: Terms<MAX_TERM_CNT>) -> Self
     {
-        Self {
-            component_storage: world
-                .data
-                .component_storage
-                .read_nonblock()
-                .expect("Failed to acquire read-only component storage lock"),
-            terms,
-        }
+        Self { world, terms }
     }
 }
 
-pub struct Iter<'query>
-{
-    iter: QueryEntityIter<'query>,
-}
-
-impl<'query> Iterator for Iter<'query>
+impl<'query, const MAX_TERM_CNT: usize> IntoIterator for &'query Query<'_, MAX_TERM_CNT>
 {
+    type IntoIter = Iter<'query>;
     type Item = EntityHandle<'query>;
 
-    fn next(&mut self) -> Option<Self::Item>
+    fn into_iter(self) -> Self::IntoIter
     {
-        let (archetype, entity) = self.iter.next()?;
-
-        Some(EntityHandle { archetype, entity })
+        self.iter()
     }
 }
 
-pub struct EntityHandle<'query>
+pub struct Iter<'query>
 {
-    archetype: &'query Archetype,
-    entity: &'query ArchetypeEntity,
+    iter: QueryEntityIter<'query>,
 }
 
-impl<'query> EntityHandle<'query>
+impl<'query> Iterator for Iter<'query>
 {
-    /// Returns the [`Uid`] of this entity.
-    #[inline]
-    #[must_use]
-    pub fn uid(&self) -> Uid
-    {
-        self.entity.uid
-    }
+    type Item = EntityHandle<'query>;
 
-    #[inline]
-    #[must_use]
-    pub fn get_component(&self, component_uid: Uid) -> Option<&'query EntityComponent>
+    fn next(&mut self) -> Option<Self::Item>
     {
-        let index = self.archetype.get_index_for_component(component_uid)?;
+        let (archetype, entity) = self.iter.next()?;
 
-        Some(self.entity.components.get(index).unwrap())
+        Some(EntityHandle::new(archetype, entity))
     }
 }
 
diff --git a/ecs/src/query/term.rs b/ecs/src/query/term.rs
index 7f24147..9c772da 100644
--- a/ecs/src/query/term.rs
+++ b/ecs/src/query/term.rs
@@ -1,38 +1,115 @@
+use std::any::type_name;
 use std::marker::PhantomData;
 
-use crate::component::Component;
-use crate::query::{TermWithoutField, TermsBuilder, TermsBuilderInterface};
+use crate::component::{
+    Component,
+    Handle as ComponentHandle,
+    HandleMut as ComponentHandleMut,
+};
+use crate::query::{
+    TermWithField,
+    TermWithoutField,
+    TermsBuilder,
+    TermsBuilderInterface,
+};
+use crate::uid::With as WithUid;
 
-pub struct With<ComponentT>
+pub struct With<WithUidT>
 where
-    ComponentT: Component,
+    WithUidT: WithUid,
 {
-    _pd: PhantomData<ComponentT>,
+    _pd: PhantomData<WithUidT>,
 }
 
-impl<ComponentT> TermWithoutField for With<ComponentT>
+impl<WithUidT> TermWithoutField for With<WithUidT>
 where
-    ComponentT: Component,
+    WithUidT: WithUid,
 {
-    fn apply_to_terms_builder(terms_builder: &mut TermsBuilder<'_>)
+    fn apply_to_terms_builder<const MAX_TERM_CNT: usize>(
+        terms_builder: &mut TermsBuilder<MAX_TERM_CNT>,
+    )
     {
-        terms_builder.with::<ComponentT>();
+        terms_builder.with::<WithUidT>();
     }
 }
 
-pub struct Without<ComponentT>
+pub struct Without<WithUidT>
 where
-    ComponentT: Component,
+    WithUidT: WithUid,
 {
-    _pd: PhantomData<ComponentT>,
+    _pd: PhantomData<WithUidT>,
 }
 
-impl<ComponentT> TermWithoutField for Without<ComponentT>
+impl<WithUidT> TermWithoutField for Without<WithUidT>
 where
-    ComponentT: Component,
+    WithUidT: WithUid,
 {
-    fn apply_to_terms_builder(terms_builder: &mut TermsBuilder<'_>)
+    fn apply_to_terms_builder<const MAX_TERM_CNT: usize>(
+        terms_builder: &mut TermsBuilder<MAX_TERM_CNT>,
+    )
     {
-        terms_builder.without::<ComponentT>();
+        terms_builder.without::<WithUidT>();
+    }
+}
+
+impl<ComponentT: Component> TermWithField for Option<&ComponentT>
+{
+    type Field<'a> = Option<ComponentHandle<'a, ComponentT>>;
+
+    fn apply_to_terms_builder<const MAX_TERM_CNT: usize>(
+        _terms_builder: &mut TermsBuilder<MAX_TERM_CNT>,
+    )
+    {
+    }
+
+    fn get_field<'world>(
+        entity_handle: &crate::entity::Handle<'world>,
+        _world: &'world crate::World,
+    ) -> Self::Field<'world>
+    {
+        Some(
+            ComponentHandle::<'world, ComponentT>::from_entity_component_ref(
+                entity_handle
+                    .get_matching_components(ComponentT::id())
+                    .next()?,
+            )
+            .unwrap_or_else(|err| {
+                panic!(
+                    "Creating handle to component {} failed: {err}",
+                    type_name::<ComponentT>()
+                );
+            }),
+        )
+    }
+}
+
+impl<ComponentT: Component> TermWithField for Option<&mut ComponentT>
+{
+    type Field<'a> = Option<ComponentHandleMut<'a, ComponentT>>;
+
+    fn apply_to_terms_builder<const MAX_TERM_CNT: usize>(
+        _terms_builder: &mut TermsBuilder<MAX_TERM_CNT>,
+    )
+    {
+    }
+
+    fn get_field<'world>(
+        entity_handle: &crate::entity::Handle<'world>,
+        _world: &'world crate::World,
+    ) -> Self::Field<'world>
+    {
+        Some(
+            ComponentHandleMut::<'world, ComponentT>::from_entity_component_ref(
+                entity_handle
+                    .get_matching_components(ComponentT::id())
+                    .next()?,
+            )
+            .unwrap_or_else(|err| {
+                panic!(
+                    "Creating handle to component {} failed: {err}",
+                    type_name::<ComponentT>()
+                );
+            }),
+        )
     }
 }