From 39f9060410b2010aa3dde47b7aeee659f1a7a3fb Mon Sep 17 00:00:00 2001 From: aviac Date: Mon, 12 Aug 2024 07:58:07 +0200 Subject: [PATCH] implement bare PoC query predicate filter --- crates/bevy_ecs/src/lib.rs | 5 +- crates/bevy_ecs/src/query/mod.rs | 33 ++++++++- crates/bevy_ecs/src/query/predicate.rs | 92 ++++++++++++++++++++++++++ 3 files changed, 128 insertions(+), 2 deletions(-) create mode 100644 crates/bevy_ecs/src/query/predicate.rs diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs index 9a9a533da09b5..3ff26cc1cd893 100644 --- a/crates/bevy_ecs/src/lib.rs +++ b/crates/bevy_ecs/src/lib.rs @@ -52,7 +52,10 @@ pub mod prelude { entity::{Entity, EntityMapper}, event::{Event, EventMutator, EventReader, EventWriter, Events}, observer::{Observer, Trigger}, - query::{Added, AnyOf, Changed, Has, Or, QueryBuilder, QueryState, With, Without}, + query::{ + Added, AnyOf, Changed, Has, Or, Predicate, PredicateFilter, QueryBuilder, QueryState, + With, Without, + }, removal_detection::RemovedComponents, schedule::{ apply_deferred, common_conditions::*, Condition, IntoSystemConfigs, IntoSystemSet, diff --git a/crates/bevy_ecs/src/query/mod.rs b/crates/bevy_ecs/src/query/mod.rs index f1719d8ba5292..d2e1c6c7326b9 100644 --- a/crates/bevy_ecs/src/query/mod.rs +++ b/crates/bevy_ecs/src/query/mod.rs @@ -7,6 +7,7 @@ mod fetch; mod filter; mod iter; mod par_iter; +mod predicate; mod state; mod world_query; @@ -18,6 +19,7 @@ pub use fetch::*; pub use filter::*; pub use iter::*; pub use par_iter::*; +pub use predicate::*; pub use state::*; pub use world_query::*; @@ -71,7 +73,7 @@ impl DebugCheckedUnwrap for Option { mod tests { use bevy_ecs_macros::{QueryData, QueryFilter}; - use crate::prelude::{AnyOf, Changed, Entity, Or, QueryState, With, Without}; + use crate::prelude::{AnyOf, Changed, Entity, Or, Predicate, QueryState, With, Without}; use crate::query::{ArchetypeFilter, Has, QueryCombinationIter, ReadOnlyQueryData}; use crate::schedule::{IntoSystemConfigs, Schedule}; use crate::system::{IntoSystem, Query, System, SystemState}; @@ -81,6 +83,8 @@ mod tests { use std::fmt::Debug; use std::hash::Hash; + use super::PredicateFilter; + #[derive(Component, Debug, Hash, Eq, PartialEq, Clone, Copy, PartialOrd, Ord)] struct A(usize); #[derive(Component, Debug, Hash, Eq, PartialEq, Clone, Copy)] @@ -757,4 +761,31 @@ mod tests { let values = world.query::<&B>().iter(&world).collect::>(); assert_eq!(values, vec![&B(2)]); } + + #[derive(Component, Debug, Clone, Copy)] + pub struct E(usize); + + #[derive(Debug, Clone, Component, Default)] + pub struct NonZeroE; + + impl PredicateFilter for NonZeroE { + type Data = &'static E; + fn filter_predicate(item: crate::query::QueryItem) -> bool { + item.0 != 0 + } + } + + #[test] + fn query_predicate() { + let mut world = World::new(); + world.spawn((A(1), E(0))); + world.spawn((A(2), E(1))); + world.spawn((A(3), E(2))); + + let values = world + .query_filtered::<&A, Predicate>() + .iter(&mut world) + .collect::>(); + assert_eq!(values, vec![&A(2), &A(3)]); + } } diff --git a/crates/bevy_ecs/src/query/predicate.rs b/crates/bevy_ecs/src/query/predicate.rs new file mode 100644 index 0000000000000..ba8e226aa57d2 --- /dev/null +++ b/crates/bevy_ecs/src/query/predicate.rs @@ -0,0 +1,92 @@ +use std::marker::PhantomData; + +use super::{QueryData, QueryFilter, QueryItem, WorldQuery}; + +pub trait PredicateFilter { + type Data: QueryData; + fn filter_predicate(item: QueryItem) -> bool; +} + +pub struct Predicate

(PhantomData

); + +unsafe impl> WorldQuery for Predicate

{ + type Item<'a> = D::Item<'a>; + + type Fetch<'a> = D::Fetch<'a>; + + type State = D::State; + + fn shrink<'wlong: 'wshort, 'wshort>(item: Self::Item<'wlong>) -> Self::Item<'wshort> { + D::shrink(item) + } + + unsafe fn init_fetch<'w>( + world: crate::world::unsafe_world_cell::UnsafeWorldCell<'w>, + state: &Self::State, + last_run: crate::component::Tick, + this_run: crate::component::Tick, + ) -> Self::Fetch<'w> { + D::init_fetch(world, state, last_run, this_run) + } + + const IS_DENSE: bool = D::IS_DENSE; + + unsafe fn set_archetype<'w>( + fetch: &mut Self::Fetch<'w>, + state: &Self::State, + archetype: &'w crate::archetype::Archetype, + table: &'w crate::storage::Table, + ) { + D::set_archetype(fetch, state, archetype, table) + } + + unsafe fn set_table<'w>( + fetch: &mut Self::Fetch<'w>, + state: &Self::State, + table: &'w crate::storage::Table, + ) { + D::set_table(fetch, state, table) + } + + unsafe fn fetch<'w>( + fetch: &mut Self::Fetch<'w>, + entity: crate::entity::Entity, + table_row: crate::storage::TableRow, + ) -> Self::Item<'w> { + D::fetch(fetch, entity, table_row) + } + + fn update_component_access( + state: &Self::State, + access: &mut super::FilteredAccess, + ) { + D::update_component_access(state, access) + } + + fn init_state(world: &mut crate::prelude::World) -> Self::State { + D::init_state(world) + } + + fn get_state(components: &crate::component::Components) -> Option { + D::get_state(components) + } + + fn matches_component_set( + state: &Self::State, + set_contains_id: &impl Fn(crate::component::ComponentId) -> bool, + ) -> bool { + D::matches_component_set(state, set_contains_id) + } +} + +impl> QueryFilter for Predicate

{ + const IS_ARCHETYPAL: bool = false; + + unsafe fn filter_fetch( + fetch: &mut Self::Fetch<'_>, + entity: crate::entity::Entity, + table_row: crate::storage::TableRow, + ) -> bool { + P::filter_predicate(Self::fetch(fetch, entity, table_row)) + } +}