[WIP] World query filter proposal prototype implementation.

crates/bevy_ecs/macros/src/fetch.rs

@@ -25,7 +25,7 @@ mod field_attr_keywords {
 
 pub static WORLD_QUERY_ATTRIBUTE_NAME: &str = "world_query";
 
-pub fn derive_world_query_impl(ast: DeriveInput) -> TokenStream {
+pub fn derive_world_query_impl(ast: DeriveInput, impl_filter: bool) -> TokenStream {
     let visibility = ast.vis;
 
     let mut fetch_struct_attributes = FetchStructAttributes::default();
@@ -316,13 +316,49 @@ pub fn derive_world_query_impl(ast: DeriveInput) -> TokenStream {
         }
     };
 
+    let impl_world_query_filter = if impl_filter {
+        quote! {
+            impl #user_impl_generics #path::query::WorldQueryFilter for #struct_name #user_ty_generics #user_where_clauses {}
+        }
+    } else {
+        quote! {}
+    };
+
+    let assert_world_query_filter = if impl_filter {
+        quote! {
+            fn assert_world_query_filter<T>()
+            where
+                T: #path::query::WorldQueryFilter,
+            {
+            }
+        }
+    } else {
+        quote! {}
+    };
+
+    let world_query_filter_asserts = if impl_filter {
+        quote! {
+            // Statically checks that the field implements `WorldQueryFilter`.
+            // We need this to make sure that we can't implement `WorldQueryFilter` for structs with non-`WorldQueryFilter` fields. I.e.:
+            // ```
+            // #[derive(WorldQueryFilter)]
+            // pub struct Foo { a: &'static mut MyComponent }
+            // ```
+            #( assert_world_query_filter::<#field_types>(); )*
+        }
+    } else {
+        quote! {}
+    };
+
     let tokens = TokenStream::from(quote! {
         #fetch_impl
 
         #state_impl
 
         #read_only_fetch_impl
 
+        #impl_world_query_filter
+
         impl #user_impl_generics #path::query::WorldQuery for #struct_name #user_ty_generics #user_where_clauses {
             type State = #state_struct_name #user_ty_generics;
             fn shrink<'__wlong: '__wshort, '__wshort>(item: #path::query::#item_type_alias<'__wlong, Self>)
@@ -356,9 +392,12 @@ pub fn derive_world_query_impl(ast: DeriveInput) -> TokenStream {
             {
             }
 
+            #assert_world_query_filter
+
             // We generate a readonly assertion for every struct member.
             fn assert_all #user_impl_generics_with_world () #user_where_clauses_with_world {
                 #read_only_asserts
+                #world_query_filter_asserts
             }
         };
 

crates/bevy_ecs/macros/src/lib.rs

@@ -431,7 +431,14 @@ pub fn derive_system_param(input: TokenStream) -> TokenStream {
 #[proc_macro_derive(WorldQuery, attributes(world_query))]
 pub fn derive_world_query(input: TokenStream) -> TokenStream {
     let ast = parse_macro_input!(input as DeriveInput);
-    derive_world_query_impl(ast)
+    derive_world_query_impl(ast, false)
+}
+
+/// Implement `WorldQueryFilter` to use a struct as a parameter in a query
+#[proc_macro_derive(WorldQueryFilter, attributes(world_query))]
+pub fn derive_world_query_filter(input: TokenStream) -> TokenStream {
+    let ast = parse_macro_input!(input as DeriveInput);
+    derive_world_query_impl(ast, true)
 }
 
 #[proc_macro_derive(SystemLabel)]

crates/bevy_ecs/src/lib.rs

@@ -56,7 +56,7 @@ mod tests {
         bundle::Bundle,
         component::{Component, ComponentId},
         entity::Entity,
-        query::{Added, ChangeTrackers, Changed, FilteredAccess, With, Without, WorldQuery},
+        query::{Added, ChangeTrackers, Changed, FilteredAccess, With, Without, WorldQueryFilter},
         world::{Mut, World},
     };
     use bevy_tasks::{ComputeTaskPool, TaskPool};
@@ -900,7 +900,7 @@ mod tests {
             }
         }
 
-        fn get_filtered<F: WorldQuery>(world: &mut World) -> Vec<Entity> {
+        fn get_filtered<F: WorldQueryFilter>(world: &mut World) -> Vec<Entity> {
             world
                 .query_filtered::<Entity, F>()
                 .iter(world)

crates/bevy_ecs/src/query/fetch.rs

@@ -8,7 +8,7 @@ use crate::{
     world::{Mut, World},
 };
 use bevy_ecs_macros::all_tuples;
-pub use bevy_ecs_macros::WorldQuery;
+pub use bevy_ecs_macros::{WorldQuery, WorldQueryFilter};
 use bevy_ptr::{ThinSlicePtr, UnsafeCellDeref};
 use std::{cell::UnsafeCell, marker::PhantomData};
 
@@ -28,19 +28,11 @@ use std::{cell::UnsafeCell, marker::PhantomData};
 /// - `Option<WQ>`: Queries the inner [`WorldQuery`] `WQ` but instead of discarding the entity if the world
 ///     query fails it returns [`None`]. See [`Query`](crate::system::Query).
 /// - `(WQ1, WQ2, ...)`: Queries all contained world queries allowing to query for more than one thing.
-///     This is the `And` operator for filters. See [`Or`].
+///     This is the `And` operator for filters. See [`Or`](crate::query::Or).
 /// - `ChangeTrackers<C>`: See the docs of [`ChangeTrackers`].
 /// - [`Entity`]: Using the entity type as a world query will grant access to the entity that is
 ///     being queried for. See [`Entity`].
 ///
-/// Bevy also offers a few filters like [`Added`](crate::query::Added), [`Changed`](crate::query::Changed),
-/// [`With`](crate::query::With), [`Without`](crate::query::Without) and [`Or`].
-/// For more information on these consult the item's corresponding documentation.
-///
-/// [`Or`]: crate::query::Or
-///
-/// # Derive
-///
 /// This trait can be derived with the [`derive@super::WorldQuery`] macro.
 ///
 /// You may want to implement a custom query with the derive macro for the following reasons:
@@ -260,14 +252,28 @@ use std::{cell::UnsafeCell, marker::PhantomData};
 ///
 /// # bevy_ecs::system::assert_is_system(my_system);
 /// ```
+pub trait WorldQuery: for<'w> WorldQueryGats<'w, _State = Self::State> {
+    type State: FetchState;
+
+    /// This function manually implements variance for the query items.
+    fn shrink<'wlong: 'wshort, 'wshort>(item: QueryItem<'wlong, Self>) -> QueryItem<'wshort, Self>;
+}
+
+/// Types that can be used as a query filter for queries from a [`World`].
+///
+/// Bevy offers a few filters like [`Added`](crate::query::Added), [`Changed`](crate::query::Changed),
+/// [`With`](crate::query::With), [`Without`](crate::query::Without) and [`Or`].
+/// For more information on these consult the item's corresponding documentation.
 ///
-/// ## Filters
+/// [`Or`]: crate::query::Or
 ///
-/// Using [`derive@super::WorldQuery`] macro we can create our own query filters.
+/// # Derive
+///
+/// Using [`derive@super::WorldQueryFilter`] macro we can create our own query filters.
 ///
 /// ```
 /// # use bevy_ecs::prelude::*;
-/// use bevy_ecs::{query::WorldQuery, component::Component};
+/// use bevy_ecs::{query::WorldQueryFilter, component::Component};
 ///
 /// #[derive(Component)]
 /// struct Foo;
@@ -278,7 +284,7 @@ use std::{cell::UnsafeCell, marker::PhantomData};
 /// #[derive(Component)]
 /// struct Qux;
 ///
-/// #[derive(WorldQuery)]
+/// #[derive(WorldQueryFilter)]
 /// struct MyFilter<T: Component, P: Component> {
 ///     _foo: With<Foo>,
 ///     _bar: With<Bar>,
@@ -292,12 +298,22 @@ use std::{cell::UnsafeCell, marker::PhantomData};
 ///
 /// # bevy_ecs::system::assert_is_system(my_system);
 /// ```
-pub trait WorldQuery: for<'w> WorldQueryGats<'w, _State = Self::State> {
-    type State: FetchState;
-
-    /// This function manually implements variance for the query items.
-    fn shrink<'wlong: 'wshort, 'wshort>(item: QueryItem<'wlong, Self>) -> QueryItem<'wshort, Self>;
-}
+///
+/// **Note:** component types themselves can't be used as filters:
+///
+/// ```compile_fail
+/// # use bevy_ecs::prelude::*;
+/// use bevy_ecs::query::WorldQueryFilter;
+///
+/// #[derive(Component)]
+/// struct Foo;
+///
+/// #[derive(WorldQueryFilter)]
+/// struct FooFilter {
+///     foo: &'static Foo,
+/// }
+/// ```
+pub trait WorldQueryFilter: WorldQuery {}
 
 /// The [`Fetch`] of a [`WorldQuery`], which declares which data it needs access to
 pub type QueryFetch<'w, Q> = <Q as WorldQueryGats<'w>>::Fetch;
@@ -1555,6 +1571,11 @@ macro_rules! impl_tuple_fetch {
             }
         }
 
+        #[allow(non_snake_case)]
+        #[allow(clippy::unused_unit)]
+        impl<$($name: WorldQueryFilter),*> WorldQueryFilter for ($($name,)*) {
+        }
+
         /// SAFETY: each item in the tuple is read only
         unsafe impl<'w, $($name: ReadOnlyFetch),*> ReadOnlyFetch for ($($name,)*) {}
 

crates/bevy_ecs/src/query/filter.rs

@@ -4,7 +4,7 @@ use crate::{
     entity::Entity,
     query::{
         debug_checked_unreachable, Access, Fetch, FetchState, FilteredAccess, QueryFetch,
-        ROQueryFetch, WorldQuery, WorldQueryGats,
+        ROQueryFetch, WorldQuery, WorldQueryFilter, WorldQueryGats,
     },
     storage::{ComponentSparseSet, Table, Tables},
     world::World,
@@ -55,6 +55,8 @@ impl<T: Component> WorldQuery for With<T> {
     }
 }
 
+impl<T: Component> WorldQueryFilter for With<T> {}
+
 /// The [`Fetch`] of [`With`].
 #[doc(hidden)]
 pub struct WithFetch<T> {
@@ -195,6 +197,8 @@ impl<T: Component> WorldQuery for Without<T> {
     }
 }
 
+impl<T: Component> WorldQueryFilter for Without<T> {}
+
 /// The [`Fetch`] of [`Without`].
 #[doc(hidden)]
 pub struct WithoutFetch<T> {
@@ -351,6 +355,11 @@ macro_rules! impl_query_filter_tuple {
             }
         }
 
+        #[allow(unused_variables)]
+        #[allow(non_snake_case)]
+        impl<$($filter: WorldQuery),*> WorldQueryFilter for Or<($($filter,)*)> {
+        }
+
         #[allow(unused_variables)]
         #[allow(non_snake_case)]
         impl<'w, $($filter: WorldQueryGats<'w>),*> WorldQueryGats<'w> for Or<($($filter,)*)> {
@@ -523,6 +532,9 @@ macro_rules! impl_tick_filter {
             }
         }
 
+        impl<T: Component> WorldQueryFilter for $name<T> {
+        }
+
         // SAFETY: this reads the T component. archetype component access and component access are updated to reflect that
         unsafe impl<T: Component> FetchState for $state_name<T> {
             fn init(world: &mut World) -> Self {