Query Filters

crates/bevy_ecs/hecs/macros/src/lib.rs

@@ -281,38 +281,41 @@ pub fn impl_query_set(_input: TokenStream) -> TokenStream {
     let mut tokens = TokenStream::new();
     let max_queries = 4;
     let queries = get_idents(|i| format!("Q{}", i), max_queries);
+    let filters = get_idents(|i| format!("F{}", i), max_queries);
     let lifetimes = get_lifetimes(|i| format!("'q{}", i), max_queries);
     let mut query_fns = Vec::new();
     let mut query_fn_muts = Vec::new();
     for i in 0..max_queries {
         let query = &queries[i];
         let lifetime = &lifetimes[i];
+        let filter = &filters[i];
         let fn_name = Ident::new(&format!("q{}", i), Span::call_site());
         let fn_name_mut = Ident::new(&format!("q{}_mut", i), Span::call_site());
         let index = Index::from(i);
         query_fns.push(quote! {
-            pub fn #fn_name(&self) -> &Query<#lifetime, #query> {
+            pub fn #fn_name(&self) -> &Query<#lifetime, #query, #filter> {
                 &self.value.#index
             }
         });
         query_fn_muts.push(quote! {
-            pub fn #fn_name_mut(&mut self) -> &mut Query<#lifetime, #query> {
+            pub fn #fn_name_mut(&mut self) -> &mut Query<#lifetime, #query, #filter> {
                 &mut self.value.#index
             }
         });
     }
 
     for query_count in 1..=max_queries {
         let query = &queries[0..query_count];
+        let filter = &filters[0..query_count];
         let lifetime = &lifetimes[0..query_count];
         let query_fn = &query_fns[0..query_count];
         let query_fn_mut = &query_fn_muts[0..query_count];
         tokens.extend(TokenStream::from(quote! {
-            impl<#(#lifetime,)* #(#query: HecsQuery,)*> QueryTuple for (#(Query<#lifetime, #query>,)*) {
+            impl<#(#lifetime,)* #(#query: HecsQuery,)* #(#filter: QueryFilter,)*> QueryTuple for (#(Query<#lifetime, #query, #filter>,)*) {
                 unsafe fn new(world: &World, component_access: &TypeAccess<ArchetypeComponent>) -> Self {
                     (
                         #(
-                            Query::<#query>::new(
+                            Query::<#query, #filter>::new(
                                 std::mem::transmute(world),
                                 std::mem::transmute(component_access),
                             ),
@@ -322,12 +325,12 @@ pub fn impl_query_set(_input: TokenStream) -> TokenStream {
 
                 fn get_accesses() -> Vec<QueryAccess> {
                     vec![
-                        #(<#query::Fetch as Fetch>::access(),)*
+                        #(QueryAccess::union(vec![<#query::Fetch as Fetch>::access(), #filter::access()]),)*
                     ]
                 }
             }
 
-            impl<#(#lifetime,)* #(#query: HecsQuery,)*> QuerySet<(#(Query<#lifetime, #query>,)*)> {
+            impl<#(#lifetime,)* #(#query: HecsQuery,)* #(#filter: QueryFilter,)*> QuerySet<(#(Query<#lifetime, #query, #filter>,)*)> {
                 #(#query_fn)*
                 #(#query_fn_mut)*
             }

crates/bevy_ecs/hecs/src/access.rs

@@ -282,10 +282,11 @@ impl<T: Hash + Eq + PartialEq + Copy> TypeAccess<T> {
 
 #[cfg(test)]
 mod tests {
-    use crate::{ArchetypeComponent, Entity, Fetch, Query, TypeAccess, With, World};
+    use crate::{ArchetypeComponent, Entity, Fetch, Query, QueryAccess, TypeAccess, World};
     use std::vec;
 
     struct A;
+    #[derive(Clone, Eq, PartialEq, Debug)]
     struct B;
     struct C;
 
@@ -344,7 +345,7 @@ mod tests {
         );
 
         let mut a_with_b_type_access = TypeAccess::default();
-        <With<B, &A> as Query>::Fetch::access()
+        QueryAccess::with::<B>(<&A as Query>::Fetch::access())
             .get_world_archetype_access(&world, Some(&mut a_with_b_type_access));
 
         assert_eq!(
@@ -353,7 +354,7 @@ mod tests {
         );
 
         let mut a_with_b_option_c_type_access = TypeAccess::default();
-        <With<B, (&A, Option<&mut C>)> as Query>::Fetch::access()
+        QueryAccess::with::<B>(<(&A, Option<&mut C>) as Query>::Fetch::access())
             .get_world_archetype_access(&world, Some(&mut a_with_b_option_c_type_access));
 
         assert_eq!(

crates/bevy_ecs/hecs/src/filter.rs

@@ -0,0 +1,245 @@
+use crate::{archetype::Archetype, Component, QueryAccess};
+use core::{any::TypeId, marker::PhantomData, ptr::NonNull};
+use std::{boxed::Box, vec};
+
+pub trait QueryFilter: Sized {
+    type EntityFilter: EntityFilter;
+    fn access() -> QueryAccess;
+    fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter>;
+}
+
+pub trait EntityFilter: Sized {
+    const DANGLING: Self;
+
+    /// # Safety
+    /// This might access archetype data in an unsafe manner. In general filters should be read-only and they should only access
+    /// the data they have claimed in `access()`.
+    unsafe fn matches_entity(&self, _offset: usize) -> bool;
+}
+
+pub struct AnyEntityFilter;
+
+impl EntityFilter for AnyEntityFilter {
+    const DANGLING: Self = AnyEntityFilter;
+
+    #[inline]
+    unsafe fn matches_entity(&self, _offset: usize) -> bool {
+        true
+    }
+}
+
+pub struct Or<T>(pub T);
+
+/// Query transformer that retrieves components of type `T` that have been mutated since the start of the frame.
+/// Added components do not count as mutated.
+pub struct Mutated<T>(NonNull<bool>, PhantomData<T>);
+
+/// Query transformer that retrieves components of type `T` that have been added since the start of the frame.
+pub struct Added<T>(NonNull<bool>, PhantomData<T>);
+
+/// Query transformer that retrieves components of type `T` that have either been mutated or added since the start of the frame.
+pub struct Changed<T>(NonNull<bool>, NonNull<bool>, PhantomData<T>);
+
+impl QueryFilter for () {
+    type EntityFilter = AnyEntityFilter;
+
+    fn access() -> QueryAccess {
+        QueryAccess::None
+    }
+
+    #[inline]
+    fn get_entity_filter(_archetype: &Archetype) -> Option<Self::EntityFilter> {
+        Some(AnyEntityFilter)
+    }
+}
+
+impl<T: Component> QueryFilter for Added<T> {
+    type EntityFilter = Self;
+
+    fn access() -> QueryAccess {
+        QueryAccess::read::<T>()
+    }
+
+    #[inline]
+    fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
+        archetype
+            .get_type_state(TypeId::of::<T>())
+            .map(|state| Added(state.added(), Default::default()))
+    }
+}
+
+impl<T: Component> EntityFilter for Added<T> {
+    const DANGLING: Self = Added(NonNull::dangling(), PhantomData::<T>);
+
+    #[inline]
+    unsafe fn matches_entity(&self, offset: usize) -> bool {
+        *self.0.as_ptr().add(offset)
+    }
+}
+
+impl<T: Component> QueryFilter for Mutated<T> {
+    type EntityFilter = Self;
+
+    fn access() -> QueryAccess {
+        QueryAccess::read::<T>()
+    }
+
+    #[inline]
+    fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
+        archetype
+            .get_type_state(TypeId::of::<T>())
+            .map(|state| Mutated(state.mutated(), Default::default()))
+    }
+}
+
+impl<T: Component> EntityFilter for Mutated<T> {
+    const DANGLING: Self = Mutated(NonNull::dangling(), PhantomData::<T>);
+
+    unsafe fn matches_entity(&self, offset: usize) -> bool {
+        *self.0.as_ptr().add(offset)
+    }
+}
+
+impl<T: Component> QueryFilter for Changed<T> {
+    type EntityFilter = Self;
+
+    fn access() -> QueryAccess {
+        QueryAccess::read::<T>()
+    }
+
+    #[inline]
+    fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
+        archetype
+            .get_type_state(TypeId::of::<T>())
+            .map(|state| Changed(state.added(), state.mutated(), Default::default()))
+    }
+}
+
+impl<T: Component> EntityFilter for Changed<T> {
+    const DANGLING: Self = Changed(NonNull::dangling(), NonNull::dangling(), PhantomData::<T>);
+
+    #[inline]
+    unsafe fn matches_entity(&self, offset: usize) -> bool {
+        *self.0.as_ptr().add(offset) || *self.1.as_ptr().add(offset)
+    }
+}
+
+pub struct Without<T>(PhantomData<T>);
+
+impl<T: Component> QueryFilter for Without<T> {
+    type EntityFilter = AnyEntityFilter;
+
+    fn access() -> QueryAccess {
+        QueryAccess::without::<T>(QueryAccess::None)
+    }
+
+    #[inline]
+    fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
+        if archetype.has_type(TypeId::of::<T>()) {
+            None
+        } else {
+            Some(AnyEntityFilter)
+        }
+    }
+}
+
+pub struct With<T>(PhantomData<T>);
+
+impl<T: Component> QueryFilter for With<T> {
+    type EntityFilter = AnyEntityFilter;
+
+    fn access() -> QueryAccess {
+        QueryAccess::with::<T>(QueryAccess::None)
+    }
+
+    #[inline]
+    fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
+        if archetype.has_type(TypeId::of::<T>()) {
+            Some(AnyEntityFilter)
+        } else {
+            None
+        }
+    }
+}
+
+macro_rules! impl_query_filter_tuple {
+    ($($filter: ident),*) => {
+        #[allow(unused_variables)]
+        #[allow(non_snake_case)]
+        impl<$($filter: QueryFilter),*> QueryFilter for ($($filter,)*) {
+            type EntityFilter = ($($filter::EntityFilter,)*);
+
+            fn access() -> QueryAccess {
+                QueryAccess::union(vec![
+                    $($filter::access(),)+
+                ])
+            }
+
+            fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
+                Some(($($filter::get_entity_filter(archetype)?,)*))
+            }
+
+        }
+
+        #[allow(unused_variables)]
+        #[allow(non_snake_case)]
+        impl<$($filter: EntityFilter),*> EntityFilter for ($($filter,)*) {
+            const DANGLING: Self = ($($filter::DANGLING,)*);
+            unsafe fn matches_entity(&self, offset: usize) -> bool {
+                let ($($filter,)*) = self;
+                true $(&& $filter.matches_entity(offset))*
+            }
+        }
+
+        #[allow(unused_variables)]
+        #[allow(non_snake_case)]
+        impl<$($filter: QueryFilter),*> QueryFilter for Or<($($filter,)*)> {
+            type EntityFilter = Or<($(Option<<$filter as QueryFilter>::EntityFilter>,)*)>;
+            fn access() -> QueryAccess {
+                QueryAccess::union(vec![
+                    $(QueryAccess::Optional(Box::new($filter::access())),)+
+                ])
+            }
+
+            fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
+                let mut matches_something = false;
+                $(
+                    let $filter = $filter::get_entity_filter(archetype);
+                    matches_something = matches_something || $filter.is_some();
+                )*
+                if matches_something {
+                    Some(Or(($($filter,)*)))
+                } else {
+                    None
+                }
+            }
+
+        }
+        #[allow(unused_variables)]
+        #[allow(non_snake_case)]
+        impl<$($filter: EntityFilter),*> EntityFilter for Or<($(Option<$filter>,)*)> {
+            const DANGLING: Self = Or(($(Some($filter::DANGLING),)*));
+            unsafe fn matches_entity(&self, offset: usize) -> bool {
+                let Or(($($filter,)*)) = self;
+                false $(|| $filter.as_ref().map_or(false, |filter|filter.matches_entity(offset)))*
+            }
+        }
+    };
+}
+
+impl_query_filter_tuple!(A);
+impl_query_filter_tuple!(A, B);
+impl_query_filter_tuple!(A, B, C);
+impl_query_filter_tuple!(A, B, C, D);
+impl_query_filter_tuple!(A, B, C, D, E);
+impl_query_filter_tuple!(A, B, C, D, E, F);
+impl_query_filter_tuple!(A, B, C, D, E, F, G);
+impl_query_filter_tuple!(A, B, C, D, E, F, G, H);
+impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I);
+impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J);
+impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J, K);
+impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J, K, L);
+impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J, K, L, M);
+impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J, K, L, M, N);
+impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O);
+impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P);

crates/bevy_ecs/hecs/src/lib.rs

@@ -69,6 +69,7 @@ mod borrow;
 mod bundle;
 mod entities;
 mod entity_builder;
+mod filter;
 mod query;
 #[cfg(feature = "serde")]
 mod serde;
@@ -80,10 +81,8 @@ pub use borrow::{AtomicBorrow, Ref, RefMut};
 pub use bundle::{Bundle, DynamicBundle, MissingComponent};
 pub use entities::{Entity, EntityReserver, Location, NoSuchEntity};
 pub use entity_builder::{BuiltEntity, EntityBuilder};
-pub use query::{
-    Added, Batch, BatchedIter, Changed, Mut, Mutated, Or, Query, QueryIter, ReadOnlyFetch, With,
-    Without,
-};
+pub use filter::{Added, Changed, EntityFilter, Mutated, Or, QueryFilter, With, Without};
+pub use query::{Batch, BatchedIter, Mut, Query, QueryIter, ReadOnlyFetch};
 pub use world::{ArchetypesGeneration, Component, ComponentError, SpawnBatchIter, World};
 
 // Unstable implementation details needed by the macros