diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs index 04731b45cffd0..813935de1cec9 100644 --- a/crates/bevy_ecs/src/lib.rs +++ b/crates/bevy_ecs/src/lib.rs @@ -26,7 +26,7 @@ pub mod prelude { component::Component, entity::Entity, event::{EventReader, EventWriter}, - query::{Added, ChangeTrackers, Changed, Or, QueryState, With, Without}, + query::{Added, AnyOf, ChangeTrackers, Changed, Or, QueryState, With, Without}, schedule::{ AmbiguitySetLabel, ExclusiveSystemDescriptorCoercion, ParallelSystemDescriptorCoercion, RunCriteria, RunCriteriaDescriptorCoercion, RunCriteriaLabel, RunCriteriaPiping, diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index a98cdd9dcd298..e01021e8564a9 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/crates/bevy_ecs/src/query/fetch.rs @@ -1103,7 +1103,120 @@ macro_rules! impl_tuple_fetch { }; } +/// The `AnyOf` query parameter fetches entities with any of the component types included in T. +/// +/// `Query>` is equivalent to `Query<(Option<&A>, Option<&B>, Option<&mut C>), (Or(With, With, With)>`. +/// Each of the components in `T` is returned as an `Option`, as with `Option` queries. +/// Entities are guaranteed to have at least one of the components in `T`. +pub struct AnyOf(T); + +macro_rules! impl_anytuple_fetch { + ($(($name: ident, $state: ident)),*) => { + #[allow(non_snake_case)] + impl<'w, 's, $($name: Fetch<'w, 's>),*> Fetch<'w, 's> for AnyOf<($(($name, bool),)*)> { + type Item = ($(Option<$name::Item>,)*); + type State = AnyOf<($($name::State,)*)>; + + #[allow(clippy::unused_unit)] + unsafe fn init(_world: &World, state: &Self::State, _last_change_tick: u32, _change_tick: u32) -> Self { + let ($($name,)*) = &state.0; + AnyOf(($(($name::init(_world, $name, _last_change_tick, _change_tick), false),)*)) + } + + + const IS_DENSE: bool = true $(&& $name::IS_DENSE)*; + + #[inline] + unsafe fn set_archetype(&mut self, _state: &Self::State, _archetype: &Archetype, _tables: &Tables) { + let ($($name,)*) = &mut self.0; + let ($($state,)*) = &_state.0; + $( + $name.1 = $state.matches_archetype(_archetype); + if $name.1 { + $name.0.set_archetype($state, _archetype, _tables); + } + )* + } + + #[inline] + unsafe fn set_table(&mut self, _state: &Self::State, _table: &Table) { + let ($($name,)*) = &mut self.0; + let ($($state,)*) = &_state.0; + $( + $name.1 = $state.matches_table(_table); + if $name.1 { + $name.0.set_table($state, _table); + } + )* + } + + #[inline] + #[allow(clippy::unused_unit)] + unsafe fn table_fetch(&mut self, _table_row: usize) -> Self::Item { + let ($($name,)*) = &mut self.0; + ($( + $name.1.then(|| $name.0.table_fetch(_table_row)), + )*) + } + + #[inline] + #[allow(clippy::unused_unit)] + unsafe fn archetype_fetch(&mut self, _archetype_index: usize) -> Self::Item { + let ($($name,)*) = &mut self.0; + ($( + $name.1.then(|| $name.0.archetype_fetch(_archetype_index)), + )*) + } + } + + // SAFETY: update_component_access and update_archetype_component_access are called for each item in the tuple + #[allow(non_snake_case)] + #[allow(clippy::unused_unit)] + unsafe impl<$($name: FetchState),*> FetchState for AnyOf<($($name,)*)> { + fn init(_world: &mut World) -> Self { + AnyOf(($($name::init(_world),)*)) + } + + fn update_component_access(&self, _access: &mut FilteredAccess) { + let ($($name,)*) = &self.0; + $($name.update_component_access(_access);)* + } + + fn update_archetype_component_access(&self, _archetype: &Archetype, _access: &mut Access) { + let ($($name,)*) = &self.0; + $( + if $name.matches_archetype(_archetype) { + $name.update_archetype_component_access(_archetype, _access); + } + )* + } + + fn matches_archetype(&self, _archetype: &Archetype) -> bool { + let ($($name,)*) = &self.0; + false $(|| $name.matches_archetype(_archetype))* + } + + fn matches_table(&self, _table: &Table) -> bool { + let ($($name,)*) = &self.0; + false $(|| $name.matches_table(_table))* + } + } + + impl<$($name: WorldQuery),*> WorldQuery for AnyOf<($($name,)*)> { + type Fetch = AnyOf<($(($name::Fetch, bool),)*)>; + type ReadOnlyFetch = AnyOf<($(($name::ReadOnlyFetch, bool),)*)>; + + type State = AnyOf<($($name::State,)*)>; + } + + /// SAFETY: each item in the tuple is read only + unsafe impl<$($name: ReadOnlyFetch),*> ReadOnlyFetch for AnyOf<($(($name, bool),)*)> {} + + }; +} + all_tuples!(impl_tuple_fetch, 0, 15, F, S); +all_tuples!(impl_anytuple_fetch, 0, 15, F, S); /// [`Fetch`] that does not actually fetch anything /// diff --git a/crates/bevy_ecs/src/query/mod.rs b/crates/bevy_ecs/src/query/mod.rs index dcc2416940cc2..0603a53bfd2ae 100644 --- a/crates/bevy_ecs/src/query/mod.rs +++ b/crates/bevy_ecs/src/query/mod.rs @@ -12,12 +12,15 @@ pub use state::*; #[cfg(test)] mod tests { + use super::AnyOf; use crate::{self as bevy_ecs, component::Component, world::World}; #[derive(Component, Debug, Eq, PartialEq)] struct A(usize); #[derive(Component, Debug, Eq, PartialEq)] struct B(usize); + #[derive(Component, Debug, Eq, PartialEq)] + struct C(usize); #[derive(Component, Debug, Eq, PartialEq)] #[component(storage = "SparseSet")] @@ -184,4 +187,21 @@ mod tests { let values = world.query::<&B>().iter(&world).collect::>(); assert_eq!(values, vec![&B(3)]); } + + #[test] + fn any_query() { + let mut world = World::new(); + + world.spawn().insert_bundle((A(1), B(2))); + world.spawn().insert_bundle((A(2),)); + world.spawn().insert_bundle((C(3),)); + + let values: Vec<(Option<&A>, Option<&B>)> = + world.query::>().iter(&world).collect(); + + assert_eq!( + values, + vec![(Some(&A(1)), Some(&B(2))), (Some(&A(2)), None),] + ); + } }