1 files changed, 78 insertions, 43 deletions
diff --git a/engine-ecs/src/query/term.rs b/engine-ecs/src/query/term.rs
index 7c4503a..04bca1d 100644
--- a/engine-ecs/src/query/term.rs
+++ b/engine-ecs/src/query/term.rs
@@ -1,17 +1,19 @@
 use std::any::type_name;
 use std::marker::PhantomData;
 
+use crate::World;
 use crate::component::{
     Component,
     Handle as ComponentHandle,
     HandleMut as ComponentHandleMut,
 };
+use crate::entity::Handle as EntityHandle;
 use crate::query::{
-    TermWithField,
-    TermWithoutField,
+    Term,
     TermsBuilder,
     TermsBuilderInterface,
 };
+use crate::tuple::Tuple;
 use crate::uid::With as WithUid;
 
 pub struct With<WithUidT>
@@ -21,16 +23,27 @@ where
     _pd: PhantomData<WithUidT>,
 }
 
-impl<WithUidT> TermWithoutField for With<WithUidT>
+impl<'world, WithUidT> Term<'world> for With<WithUidT>
 where
     WithUidT: WithUid,
 {
+    type AddField<Fields: Tuple> = Fields;
+
     fn apply_to_terms_builder<const MAX_TERM_CNT: usize>(
         terms_builder: &mut TermsBuilder<MAX_TERM_CNT>,
     )
     {
         terms_builder.present([WithUidT::uid()]);
     }
+
+    fn add_field<Fields: Tuple>(
+        _entity_handle: &EntityHandle<'world>,
+        _world: &'world World,
+        fields: Fields,
+    ) -> Self::AddField<Fields>
+    {
+        fields
+    }
 }
 
 pub struct Without<WithUidT>
@@ -40,21 +53,33 @@ where
     _pd: PhantomData<WithUidT>,
 }
 
-impl<WithUidT> TermWithoutField for Without<WithUidT>
+impl<'world, WithUidT> Term<'world> for Without<WithUidT>
 where
     WithUidT: WithUid,
 {
+    type AddField<Fields: Tuple> = Fields;
+
     fn apply_to_terms_builder<const MAX_TERM_CNT: usize>(
         terms_builder: &mut TermsBuilder<MAX_TERM_CNT>,
     )
     {
         terms_builder.absent([WithUidT::uid()]);
     }
+
+    fn add_field<Fields: Tuple>(
+        _entity_handle: &EntityHandle<'world>,
+        _world: &'world World,
+        fields: Fields,
+    ) -> Self::AddField<Fields>
+    {
+        fields
+    }
 }
 
-impl<ComponentT: Component> TermWithField for Option<&ComponentT>
+impl<'world, ComponentT: Component> Term<'world> for Option<&ComponentT>
 {
-    type Field<'a> = Option<ComponentHandle<'a, ComponentT>>;
+    type AddField<Fields: Tuple> =
+        Fields::WithElementAtEnd<Option<ComponentHandle<'world, ComponentT>>>;
 
     fn apply_to_terms_builder<const MAX_TERM_CNT: usize>(
         _terms_builder: &mut TermsBuilder<MAX_TERM_CNT>,
@@ -62,30 +87,36 @@ impl<ComponentT: Component> TermWithField for Option<&ComponentT>
     {
     }
 
-    fn get_field<'world>(
-        entity_handle: &crate::entity::Handle<'world>,
-        _world: &'world crate::World,
-    ) -> Self::Field<'world>
+    fn add_field<Fields: Tuple>(
+        entity_handle: &EntityHandle<'world>,
+        _world: &'world World,
+        fields: Fields,
+    ) -> Self::AddField<Fields>
     {
-        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>()
-                );
-            }),
-        )
+        let component = (|| {
+            let comp_ref = &entity_handle
+                .get_matching_components(ComponentT::id())
+                .next()?;
+
+            match ComponentHandle::<ComponentT>::from_entity_component_ref(comp_ref) {
+                Ok(component) => Some(component),
+                Err(err) => {
+                    panic!(
+                        "Creating handle to component {} failed: {err}",
+                        type_name::<ComponentT>()
+                    );
+                }
+            }
+        })();
+
+        fields.with_elem(component)
     }
 }
 
-impl<ComponentT: Component> TermWithField for Option<&mut ComponentT>
+impl<'world, ComponentT: Component> Term<'world> for Option<&mut ComponentT>
 {
-    type Field<'a> = Option<ComponentHandleMut<'a, ComponentT>>;
+    type AddField<Fields: Tuple> =
+        Fields::WithElementAtEnd<Option<ComponentHandleMut<'world, ComponentT>>>;
 
     fn apply_to_terms_builder<const MAX_TERM_CNT: usize>(
         _terms_builder: &mut TermsBuilder<MAX_TERM_CNT>,
@@ -93,24 +124,28 @@ impl<ComponentT: Component> TermWithField for Option<&mut ComponentT>
     {
     }
 
-    fn get_field<'world>(
-        entity_handle: &crate::entity::Handle<'world>,
-        world: &'world crate::World,
-    ) -> Self::Field<'world>
+    fn add_field<Fields: Tuple>(
+        entity_handle: &EntityHandle<'world>,
+        world: &'world World,
+        fields: Fields,
+    ) -> Self::AddField<Fields>
     {
-        Some(
-            ComponentHandleMut::<'world, ComponentT>::from_entity_component_ref(
-                &entity_handle
-                    .get_matching_components(ComponentT::id())
-                    .next()?,
-                world,
-            )
-            .unwrap_or_else(|err| {
-                panic!(
-                    "Creating handle to component {} failed: {err}",
-                    type_name::<ComponentT>()
-                );
-            }),
-        )
+        let component = (|| {
+            let comp_ref = &entity_handle
+                .get_matching_components(ComponentT::id())
+                .next()?;
+
+            match ComponentHandleMut::<ComponentT>::from_entity_component_ref(comp_ref, world) {
+                Ok(component) => Some(component),
+                Err(err) => {
+                    panic!(
+                        "Creating mut handle to component {} failed: {err}",
+                        type_name::<ComponentT>()
+                    );
+                }
+            }
+        })();
+
+        fields.with_elem(component)
     }
 }