From 9b6611cd11199346cbe1f14ad44930347f90dec2 Mon Sep 17 00:00:00 2001
From: HampusM <hampus@hampusmat.com>
Date: Sat, 22 Jun 2024 16:15:12 +0200
Subject: feat(ecs): add query options filter entities

---
 ecs/src/query.rs | 52 +++++++++++++++++++++++++++++++++-------------------
 1 file changed, 33 insertions(+), 19 deletions(-)

(limited to 'ecs/src/query.rs')

diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index dc6b5f0..8a63256 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -1,6 +1,6 @@
 use std::any::Any;
 use std::collections::HashSet;
-use std::iter::{Flatten, Map};
+use std::iter::{Filter, Flatten, Map};
 use std::marker::PhantomData;
 use std::sync::{Arc, Weak};
 
@@ -15,6 +15,7 @@ use crate::component::{
     Sequence as ComponentSequence,
 };
 use crate::lock::{Lock, ReadGuard};
+use crate::query::options::Options;
 use crate::system::{
     NoInitParamFlag as NoInitSystemParamFlag,
     Param as SystemParam,
@@ -22,19 +23,22 @@ use crate::system::{
 };
 use crate::{EntityComponent, WorldData};
 
+pub mod options;
+
 #[derive(Debug)]
-pub struct Query<'world, Comps>
+pub struct Query<'world, Comps, OptionsT = ()>
 where
     Comps: ComponentSequence,
 {
     component_storage: ReadGuard<'world, ComponentStorage>,
     component_storage_lock: Weak<Lock<ComponentStorage>>,
-    comps_pd: PhantomData<Comps>,
+    _pd: PhantomData<(Comps, OptionsT)>,
 }
 
-impl<'world, Comps> Query<'world, Comps>
+impl<'world, Comps, OptionsT> Query<'world, Comps, OptionsT>
 where
     Comps: ComponentSequence,
+    OptionsT: Options,
 {
     /// Iterates over the entities matching this query.
     #[must_use]
@@ -51,14 +55,15 @@ where
                 .component_storage
                 .find_entities(&Comps::ids())
                 .map((|archetype| archetype.components.as_slice()) as ComponentIterMapFn)
-                .flatten(),
+                .flatten()
+                .filter(|components| OptionsT::entity_filter(*components)),
             comps_pd: PhantomData,
         }
     }
 
     /// Returns a weak reference query to the same components.
     #[must_use]
-    pub fn to_weak_ref(&self) -> WeakRef<Comps>
+    pub fn to_weak_ref(&self) -> WeakRef<Comps, OptionsT>
     {
         WeakRef {
             component_storage: self.component_storage_lock.clone(),
@@ -73,14 +78,15 @@ where
                 .read_nonblock()
                 .expect("Failed to acquire read-only component storage lock"),
             component_storage_lock: Arc::downgrade(component_storage),
-            comps_pd: PhantomData,
+            _pd: PhantomData,
         }
     }
 }
 
-impl<'world, Comps> IntoIterator for &'world Query<'world, Comps>
+impl<'world, Comps, OptionsT> IntoIterator for &'world Query<'world, Comps, OptionsT>
 where
     Comps: ComponentSequence,
+    OptionsT: Options,
 {
     type IntoIter = ComponentIter<'world, Comps>;
     type Item = Comps::Refs<'world>;
@@ -91,9 +97,11 @@ where
     }
 }
 
-unsafe impl<'world, Comps> SystemParam<'world> for Query<'world, Comps>
+unsafe impl<'world, Comps, OptionsT> SystemParam<'world>
+    for Query<'world, Comps, OptionsT>
 where
     Comps: ComponentSequence,
+    OptionsT: Options,
 {
     type Flags = NoInitSystemParamFlag;
     type Input = ();
@@ -161,15 +169,15 @@ where
 
 /// A entity query containing a weak reference to the world.
 #[derive(Debug)]
-pub struct WeakRef<Comps>
+pub struct WeakRef<Comps, OptionsT>
 where
     Comps: ComponentSequence,
 {
     component_storage: Weak<Lock<ComponentStorage>>,
-    comps_pd: PhantomData<Comps>,
+    comps_pd: PhantomData<(Comps, OptionsT)>,
 }
 
-impl<Comps> WeakRef<Comps>
+impl<Comps, OptionsT> WeakRef<Comps, OptionsT>
 where
     Comps: ComponentSequence,
 {
@@ -177,7 +185,7 @@ where
     ///
     /// Returns [`None`] if the [`World`] has been dropped.
     #[must_use]
-    pub fn access(&self) -> Option<Ref<'_, Comps>>
+    pub fn access(&self) -> Option<Ref<'_, Comps, OptionsT>>
     {
         Some(Ref {
             component_storage: self.component_storage.upgrade()?,
@@ -186,7 +194,7 @@ where
     }
 }
 
-impl<Comps> Clone for WeakRef<Comps>
+impl<Comps, OptionsT> Clone for WeakRef<Comps, OptionsT>
 where
     Comps: ComponentSequence,
 {
@@ -202,20 +210,21 @@ where
 /// Intermediate between [`Query`] and [`WeakRefQuery`]. Contains a strong reference to
 /// the world which is not allowed direct access to.
 #[derive(Debug, Clone)]
-pub struct Ref<'weak_ref, Comps>
+pub struct Ref<'weak_ref, Comps, OptionsT>
 where
     Comps: ComponentSequence,
 {
     component_storage: Arc<Lock<ComponentStorage>>,
-    _pd: PhantomData<&'weak_ref Comps>,
+    _pd: PhantomData<(&'weak_ref Comps, OptionsT)>,
 }
 
-impl<'weak_ref, Comps> Ref<'weak_ref, Comps>
+impl<'weak_ref, Comps, OptionsT> Ref<'weak_ref, Comps, OptionsT>
 where
     Comps: ComponentSequence,
+    OptionsT: Options,
 {
     #[must_use]
-    pub fn to_query(&self) -> Query<'_, Comps>
+    pub fn to_query(&self) -> Query<'_, Comps, OptionsT>
     {
         Query::new(&self.component_storage)
     }
@@ -223,9 +232,14 @@ where
 
 type ComponentIterMapFn = for<'a> fn(&'a Archetype) -> &'a [Vec<EntityComponent>];
 
+type ComponentIterFilterFn = for<'a, 'b> fn(&'a &'b Vec<EntityComponent>) -> bool;
+
 pub struct ComponentIter<'world, Comps>
 {
-    entities: Flatten<Map<ArchetypeRefIter<'world>, ComponentIterMapFn>>,
+    entities: Filter<
+        Flatten<Map<ArchetypeRefIter<'world>, ComponentIterMapFn>>,
+        ComponentIterFilterFn,
+    >,
     comps_pd: PhantomData<Comps>,
 }
 
-- 
cgit v1.2.3-18-g5258