feat(ecs): add query options filter entities

3 files changed, 102 insertions, 20 deletions

diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index b883171..1f745b9 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -15,6 +15,7 @@ use crate::component::{Component, Id as ComponentId, Sequence as ComponentSequen
 use crate::event::{Event, Id as EventId, Ids, Sequence as EventSequence};
 use crate::extension::{Collector as ExtensionCollector, Extension};
 use crate::lock::Lock;
+use crate::query::options::Options as QueryOptions;
 use crate::sole::Sole;
 use crate::system::{System, TypeErased as TypeErasedSystem};
 use crate::type_name::TypeName;
@@ -121,9 +122,10 @@ impl World
         drop(event);
     }
 
-    pub fn query<Comps>(&self) -> Query<Comps>
+    pub fn query<Comps, OptionsT>(&self) -> Query<Comps, OptionsT>
     where
         Comps: ComponentSequence,
+        OptionsT: QueryOptions,
     {
         Query::new(&self.data.component_storage)
     }
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>,
 }
 
diff --git a/ecs/src/query/options.rs b/ecs/src/query/options.rs
new file mode 100644
index 0000000..d895073
--- /dev/null
+++ b/ecs/src/query/options.rs
@@ -0,0 +1,66 @@
+use std::collections::HashSet;
+use std::marker::PhantomData;
+
+use crate::component::{Component, Id as ComponentId};
+use crate::EntityComponent;
+
+/// Query options.
+pub trait Options
+{
+    fn entity_filter<'component>(
+        components: impl IntoIterator<Item = &'component EntityComponent>,
+    ) -> bool;
+}
+
+impl Options for ()
+{
+    fn entity_filter<'component>(
+        _: impl IntoIterator<Item = &'component EntityComponent>,
+    ) -> bool
+    {
+        true
+    }
+}
+
+pub struct With<ComponentT>
+where
+    ComponentT: Component,
+{
+    _pd: PhantomData<ComponentT>,
+}
+
+impl<ComponentT> Options for With<ComponentT>
+where
+    ComponentT: Component,
+{
+    fn entity_filter<'component>(
+        components: impl IntoIterator<Item = &'component EntityComponent>,
+    ) -> bool
+    {
+        let ids_set = components
+            .into_iter()
+            .map(|component| component.id)
+            .collect::<HashSet<_>>();
+
+        ids_set.contains(&ComponentId::of::<ComponentT>())
+    }
+}
+
+pub struct Not<OptionsT>
+where
+    OptionsT: Options,
+{
+    _pd: PhantomData<OptionsT>,
+}
+
+impl<OptionsT> Options for Not<OptionsT>
+where
+    OptionsT: Options,
+{
+    fn entity_filter<'component>(
+        components: impl IntoIterator<Item = &'component EntityComponent>,
+    ) -> bool
+    {
+        !OptionsT::entity_filter(components)
+    }
+}