Skip to content

Commit

Permalink
Added QueryState::many_for_each_unchecked_manual and co.
Browse files Browse the repository at this point in the history
  • Loading branch information
tim-blackbird committed Jun 2, 2022
1 parent f68ac41 commit f36f62f
Show file tree
Hide file tree
Showing 2 changed files with 186 additions and 26 deletions.
150 changes: 147 additions & 3 deletions crates/bevy_ecs/src/query/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ use bevy_tasks::{ComputeTaskPool, TaskPool};
#[cfg(feature = "trace")]
use bevy_utils::tracing::Instrument;
use fixedbitset::FixedBitSet;
use std::{fmt, ops::Deref};
use std::{borrow::Borrow, fmt, ops::Deref};

use super::{QueryFetch, QueryItem, ROQueryFetch, ROQueryItem};

Expand Down Expand Up @@ -693,6 +693,93 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
);
}

/// Runs `func` on each query result for the given [`World`]. This is faster than the equivalent
/// iter() method, but cannot be chained like a normal [`Iterator`].
///
/// This can only be called for read-only queries, see [`Self::for_each_mut`] for write-queries.
#[inline]
pub fn many_for_each<
'w,
E: Borrow<Entity>,
EntityList: IntoIterator<Item = E>,
FN: FnMut(ROQueryItem<'w, Q>),
>(
&mut self,
world: &'w World,
entities: EntityList,
func: FN,
) {
// SAFETY: query is read only
unsafe {
self.update_archetypes(world);
self.many_for_each_unchecked_manual::<ROQueryFetch<Q>, E, EntityList, FN>(
world,
entities,
func,
world.last_change_tick(),
world.read_change_tick(),
);
}
}

/// Runs `func` on each query result for the given [`World`]. This is faster than the equivalent
/// `iter_mut()` method, but cannot be chained like a normal [`Iterator`].
#[inline]
pub fn many_for_each_mut<
'w,
E: Borrow<Entity>,
EntityList: IntoIterator<Item = E>,
FN: FnMut(QueryItem<'w, Q>),
>(
&mut self,
world: &'w mut World,
entities: EntityList,
func: FN,
) {
// SAFETY: query has unique world access
unsafe {
self.update_archetypes(world);
self.many_for_each_unchecked_manual::<QueryFetch<Q>, E, EntityList, FN>(
world,
entities,
func,
world.last_change_tick(),
world.read_change_tick(),
);
}
}

/// Runs `func` on each query result for the given [`World`]. This is faster than the equivalent
/// iter() method, but cannot be chained like a normal [`Iterator`].
///
/// This can only be called for read-only queries.
///
/// # Safety
///
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
/// have unique access to the components they query.
#[inline]
pub unsafe fn many_for_each_unchecked<
'w,
E: Borrow<Entity>,
EntityList: IntoIterator<Item = E>,
FN: FnMut(QueryItem<'w, Q>),
>(
&mut self,
world: &'w World,
entities: EntityList,
func: FN,
) {
self.update_archetypes(world);
self.many_for_each_unchecked_manual::<QueryFetch<Q>, E, EntityList, FN>(
world,
entities,
func,
world.last_change_tick(),
world.read_change_tick(),
);
}

/// Runs `func` on each query result in parallel.
///
/// This can only be called for read-only queries, see [`Self::par_for_each_mut`] for
Expand Down Expand Up @@ -797,7 +884,7 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
change_tick: u32,
) {
// NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
// QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
// QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::many_for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
let mut fetch = QF::init(world, &self.fetch_state, last_change_tick, change_tick);
let mut filter = <QueryFetch<F> as Fetch>::init(
world,
Expand Down Expand Up @@ -839,6 +926,63 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
}
}

/// Runs `func` on each query result for the given [`World`], where the last change and
/// the current change tick are given. This is faster than the equivalent
/// iter() method, but cannot be chained like a normal [`Iterator`].
///
/// # Safety
///
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
/// have unique access to the components they query.
/// This does not validate that `world.id()` matches `self.world_id`. Calling this on a `world`
/// with a mismatched [`WorldId`] is unsound.
pub(crate) unsafe fn many_for_each_unchecked_manual<
'w,
QF: Fetch<'w, State = Q::State>,
E: Borrow<Entity>,
EntityList: IntoIterator<Item = E>,
FN: FnMut(QF::Item),
>(
&self,
world: &'w World,
entities: EntityList,
mut func: FN,
last_change_tick: u32,
change_tick: u32,
) {
// NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
// QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::many_for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
let mut fetch = QF::init(world, &self.fetch_state, last_change_tick, change_tick);
let mut filter = <QueryFetch<F> as Fetch>::init(
world,
&self.filter_state,
last_change_tick,
change_tick,
);
let archetypes = &world.archetypes;
let tables = &world.storages().tables;

for location in entities
.into_iter()
.filter_map(|e| world.entities.get(*e.borrow()))
{
if !self
.matched_archetypes
.contains(location.archetype_id.index())
{
continue;
}

let archetype = &archetypes[location.archetype_id];

fetch.set_archetype(&self.fetch_state, archetype, tables);
filter.set_archetype(&self.filter_state, archetype, tables);
if filter.archetype_filter_fetch(location.index) {
func(fetch.archetype_fetch(location.index))
}
}
}

/// Runs `func` on each query result in parallel for the given [`World`], where the last change and
/// the current change tick are given. This is faster than the equivalent
/// iter() method, but cannot be chained like a normal [`Iterator`].
Expand Down Expand Up @@ -866,7 +1010,7 @@ impl<Q: WorldQuery, F: WorldQuery> QueryState<Q, F> {
change_tick: u32,
) {
// NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
// QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
// QueryIter, QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::many_for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
self.task_pool
.as_ref()
.expect("Cannot iterate query in parallel. No ComputeTaskPool initialized.")
Expand Down
62 changes: 39 additions & 23 deletions crates/bevy_ecs/src/system/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1156,12 +1156,14 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> {
/// }
/// # bevy_ecs::system::assert_is_system(system);
/// ```
pub fn many_iter<I, EntityList, E>(&self, entities: EntityList) -> QueryManyIter<'_, '_, I, Q, F, E>
where
pub fn many_iter<
E: Borrow<Entity>,
I: Iterator<Item = E>,
EntityList: IntoIterator<IntoIter = I>,
{
>(
&self,
entities: EntityList,
) -> QueryManyIter<'_, '_, I, Q, F, E> {
QueryManyIter {
entities: entities.into_iter(),
query: self,
Expand Down Expand Up @@ -1195,16 +1197,23 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> {
/// }
/// # bevy_ecs::system::assert_is_system(system);
/// ```
pub fn many_for_each<EntityList, E>(&self, entities: EntityList, f: impl Fn(ROQueryItem<'_, Q>))
where
E: Borrow<Entity>,
EntityList: IntoIterator<Item = E>,
{
entities.into_iter().map(|e| *e.borrow()).for_each(|input| {
if let Ok(item) = self.get(input) {
f(item);
}
});
pub fn many_for_each<E: Borrow<Entity>, EntityList: IntoIterator<Item = E>>(
&self,
entities: EntityList,
f: impl Fn(ROQueryItem<'_, Q>),
) {
// SAFE: system runs without conflicts with other systems.
// same-system queries have runtime borrow checks when they conflict
unsafe {
self.state
.many_for_each_unchecked_manual::<ROQueryFetch<Q>, E, EntityList, _>(
self.world,
entities,
f,
self.last_change_tick,
self.change_tick,
);
};
}

/// Calls a closure on each result of [`Query`] where the entities match.
Expand Down Expand Up @@ -1235,16 +1244,23 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> {
/// }
/// # bevy_ecs::system::assert_is_system(system);
/// ```
pub fn many_for_each_mut<EntityList, E>(&mut self, entities: EntityList, mut f: impl FnMut(QueryItem<'_, Q>))
where
E: Borrow<Entity>,
EntityList: IntoIterator<Item = E>,
{
entities.into_iter().map(|e| *e.borrow()).for_each(|entity| {
if let Ok(item) = self.get_mut(entity) {
f(item);
}
});
pub fn many_for_each_mut<E: Borrow<Entity>, EntityList: IntoIterator<Item = E>>(
&mut self,
entities: EntityList,
f: impl FnMut(QueryItem<'_, Q>),
) {
// SAFE: system runs without conflicts with other systems.
// same-system queries have runtime borrow checks when they conflict
unsafe {
self.state
.many_for_each_unchecked_manual::<QueryFetch<Q>, E, EntityList, _>(
self.world,
entities,
f,
self.last_change_tick,
self.change_tick,
);
};
}
}

Expand Down

0 comments on commit f36f62f

Please sign in to comment.