From 54a9781d8d13eb1ef5def1236f9729d4b5f6d3ce Mon Sep 17 00:00:00 2001 From: devil-ira Date: Fri, 2 Sep 2022 10:52:05 +0200 Subject: [PATCH 1/6] compile fail tests --- ...y_iter_combinations_mut_iterator_safety.rs | 15 ---- ...er_combinations_mut_iterator_safety.stderr | 25 ------ .../ui/query_iter_many_mut_iterator_safety.rs | 15 ---- ...query_iter_many_mut_iterator_safety.stderr | 25 ------ .../ui/query_mut_iterator_impl_safety.rs | 22 ++++++ .../ui/query_mut_iterator_impl_safety.stderr | 77 +++++++++++++++++++ ...query_iter_join_map_mut_lifetime_safety.rs | 15 ++++ ...y_iter_join_map_mut_lifetime_safety.stderr | 5 ++ 8 files changed, 119 insertions(+), 80 deletions(-) delete mode 100644 crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_combinations_mut_iterator_safety.rs delete mode 100644 crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_combinations_mut_iterator_safety.stderr delete mode 100644 crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_many_mut_iterator_safety.rs delete mode 100644 crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_many_mut_iterator_safety.stderr create mode 100644 crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.rs create mode 100644 crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.stderr create mode 100644 crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_join_map_mut_lifetime_safety.rs create mode 100644 crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_join_map_mut_lifetime_safety.stderr diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_combinations_mut_iterator_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_combinations_mut_iterator_safety.rs deleted file mode 100644 index e88bc34cb4906..0000000000000 --- a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_combinations_mut_iterator_safety.rs +++ /dev/null @@ -1,15 +0,0 @@ -use bevy_ecs::prelude::*; - -#[derive(Component)] -struct A(usize); - -fn system(mut query: Query<&mut A>) { - let iter = query.iter_combinations_mut(); - - // This should fail to compile. - is_iterator(iter) -} - -fn is_iterator(_iter: T) {} - -fn main() {} diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_combinations_mut_iterator_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_combinations_mut_iterator_safety.stderr deleted file mode 100644 index 7cd36eadcff93..0000000000000 --- a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_combinations_mut_iterator_safety.stderr +++ /dev/null @@ -1,25 +0,0 @@ -error[E0277]: the trait bound `&mut A: ReadOnlyWorldQuery` is not satisfied - --> tests/ui/query_iter_combinations_mut_iterator_safety.rs:10:17 - | -10 | is_iterator(iter) - | ----------- ^^^^ the trait `ReadOnlyWorldQuery` is not implemented for `&mut A` - | | - | required by a bound introduced by this call - | - = help: the following other types implement trait `ReadOnlyWorldQuery`: - &T - () - (F0, F1) - (F0, F1, F2) - (F0, F1, F2, F3) - (F0, F1, F2, F3, F4) - (F0, F1, F2, F3, F4, F5) - (F0, F1, F2, F3, F4, F5, F6) - and 49 others - = note: `ReadOnlyWorldQuery` is implemented for `&A`, but not for `&mut A` - = note: required because of the requirements on the impl of `Iterator` for `QueryCombinationIter<'_, '_, &mut A, (), {_: usize}>` -note: required by a bound in `is_iterator` - --> tests/ui/query_iter_combinations_mut_iterator_safety.rs:13:19 - | -13 | fn is_iterator(_iter: T) {} - | ^^^^^^^^ required by this bound in `is_iterator` diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_many_mut_iterator_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_many_mut_iterator_safety.rs deleted file mode 100644 index d7f2cfa92c694..0000000000000 --- a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_many_mut_iterator_safety.rs +++ /dev/null @@ -1,15 +0,0 @@ -use bevy_ecs::prelude::*; - -#[derive(Component)] -struct A(usize); - -fn system(mut query: Query<&mut A>, e: Entity) { - let iter = query.iter_many_mut([e]); - - // This should fail to compile. - is_iterator(iter) -} - -fn is_iterator(_iter: T) {} - -fn main() {} diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_many_mut_iterator_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_many_mut_iterator_safety.stderr deleted file mode 100644 index 45840aa3230be..0000000000000 --- a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_iter_many_mut_iterator_safety.stderr +++ /dev/null @@ -1,25 +0,0 @@ -error[E0277]: the trait bound `&mut A: ReadOnlyWorldQuery` is not satisfied - --> tests/ui/query_iter_many_mut_iterator_safety.rs:10:17 - | -10 | is_iterator(iter) - | ----------- ^^^^ the trait `ReadOnlyWorldQuery` is not implemented for `&mut A` - | | - | required by a bound introduced by this call - | - = help: the following other types implement trait `ReadOnlyWorldQuery`: - &T - () - (F0, F1) - (F0, F1, F2) - (F0, F1, F2, F3) - (F0, F1, F2, F3, F4) - (F0, F1, F2, F3, F4, F5) - (F0, F1, F2, F3, F4, F5, F6) - and 49 others - = note: `ReadOnlyWorldQuery` is implemented for `&A`, but not for `&mut A` - = note: required because of the requirements on the impl of `Iterator` for `QueryManyIter<'_, '_, &mut A, (), std::array::IntoIter>` -note: required by a bound in `is_iterator` - --> tests/ui/query_iter_many_mut_iterator_safety.rs:13:19 - | -13 | fn is_iterator(_iter: T) {} - | ^^^^^^^^ required by this bound in `is_iterator` diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.rs new file mode 100644 index 0000000000000..2bd4bbeef4535 --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.rs @@ -0,0 +1,22 @@ +use bevy_ecs::prelude::*; + +#[derive(Component)] +struct A(usize); + +fn system(mut query: Query<&mut A>, e: Res) { + let iter = query.iter_combinations_mut(); + // This should fail to compile. + is_iterator(iter); + + let iter = query.iter_many_mut([*e]); + // This should fail to compile. + is_iterator(iter); + + let iter = query.iter_join_map_mut([*e], |e| *e); + // This should fail to compile. + is_iterator(iter); +} + +fn is_iterator(_iter: impl Iterator) {} + +fn main() {} diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.stderr new file mode 100644 index 0000000000000..159ac22f8f3ad --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.stderr @@ -0,0 +1,77 @@ +error[E0277]: the trait bound `&mut A: ReadOnlyWorldQuery` is not satisfied + --> tests/ui/query_mut_iterator_impl_safety.rs:9:17 + | +9 | is_iterator(iter); + | ----------- ^^^^ the trait `ReadOnlyWorldQuery` is not implemented for `&mut A` + | | + | required by a bound introduced by this call + | + = help: the following other types implement trait `ReadOnlyWorldQuery`: + &T + () + (F0, F1) + (F0, F1, F2) + (F0, F1, F2, F3) + (F0, F1, F2, F3, F4) + (F0, F1, F2, F3, F4, F5) + (F0, F1, F2, F3, F4, F5, F6) + and 49 others + = note: `ReadOnlyWorldQuery` is implemented for `&A`, but not for `&mut A` + = note: required because of the requirements on the impl of `Iterator` for `QueryCombinationIter<'_, '_, &mut A, (), {_: usize}>` +note: required by a bound in `is_iterator` + --> tests/ui/query_mut_iterator_impl_safety.rs:20:28 + | +20 | fn is_iterator(_iter: impl Iterator) {} + | ^^^^^^^^ required by this bound in `is_iterator` + +error[E0277]: the trait bound `&mut A: ReadOnlyWorldQuery` is not satisfied + --> tests/ui/query_mut_iterator_impl_safety.rs:13:17 + | +13 | is_iterator(iter); + | ----------- ^^^^ the trait `ReadOnlyWorldQuery` is not implemented for `&mut A` + | | + | required by a bound introduced by this call + | + = help: the following other types implement trait `ReadOnlyWorldQuery`: + &T + () + (F0, F1) + (F0, F1, F2) + (F0, F1, F2, F3) + (F0, F1, F2, F3, F4) + (F0, F1, F2, F3, F4, F5) + (F0, F1, F2, F3, F4, F5, F6) + and 49 others + = note: `ReadOnlyWorldQuery` is implemented for `&A`, but not for `&mut A` + = note: required because of the requirements on the impl of `Iterator` for `QueryManyIter<'_, '_, &mut A, (), std::array::IntoIter>` +note: required by a bound in `is_iterator` + --> tests/ui/query_mut_iterator_impl_safety.rs:20:28 + | +20 | fn is_iterator(_iter: impl Iterator) {} + | ^^^^^^^^ required by this bound in `is_iterator` + +error[E0277]: the trait bound `&mut A: ReadOnlyWorldQuery` is not satisfied + --> tests/ui/query_mut_iterator_impl_safety.rs:17:17 + | +17 | is_iterator(iter); + | ----------- ^^^^ the trait `ReadOnlyWorldQuery` is not implemented for `&mut A` + | | + | required by a bound introduced by this call + | + = help: the following other types implement trait `ReadOnlyWorldQuery`: + &T + () + (F0, F1) + (F0, F1, F2) + (F0, F1, F2, F3) + (F0, F1, F2, F3, F4) + (F0, F1, F2, F3, F4, F5) + (F0, F1, F2, F3, F4, F5, F6) + and 49 others + = note: `ReadOnlyWorldQuery` is implemented for `&A`, but not for `&mut A` + = note: required because of the requirements on the impl of `Iterator` for `QueryJoinMapIter<'_, '_, &mut A, (), std::array::IntoIter, [closure@$DIR/tests/ui/query_mut_iterator_impl_safety.rs:15:46: 15:52]>` +note: required by a bound in `is_iterator` + --> tests/ui/query_mut_iterator_impl_safety.rs:20:28 + | +20 | fn is_iterator(_iter: impl Iterator) {} + | ^^^^^^^^ required by this bound in `is_iterator` diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_join_map_mut_lifetime_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_join_map_mut_lifetime_safety.rs new file mode 100644 index 0000000000000..5b5e665fa8daa --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_join_map_mut_lifetime_safety.rs @@ -0,0 +1,15 @@ +use bevy_ecs::prelude::*; + +#[derive(Component)] +struct A(usize); + +fn system(mut query: Query<&mut A>, e: Res) { + let mut results = Vec::new(); + let mut iter = query.iter_join_map_mut([*e, *e], |e| *e); + while let Some(a) = iter.fetch_next() { + // this should fail to compile + results.push(a); + } +} + +fn main() {} diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_join_map_mut_lifetime_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_join_map_mut_lifetime_safety.stderr new file mode 100644 index 0000000000000..0238cf418a2ee --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_join_map_mut_lifetime_safety.stderr @@ -0,0 +1,5 @@ +error[E0499]: cannot borrow `iter` as mutable more than once at a time + --> tests/ui/system_query_iter_join_map_mut_lifetime_safety.rs:9:25 + | +9 | while let Some(a) = iter.fetch_next() { + | ^^^^^^^^^^^^^^^^^ `iter` was mutably borrowed here in the previous iteration of the loop From 4d7f1147ef1080990233bf8a8fbdde8d2654eabd Mon Sep 17 00:00:00 2001 From: devil-ira Date: Fri, 2 Sep 2022 10:52:15 +0200 Subject: [PATCH 2/6] test --- crates/bevy_ecs/src/system/mod.rs | 52 +++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/crates/bevy_ecs/src/system/mod.rs b/crates/bevy_ecs/src/system/mod.rs index aef76c277a5ec..e74d6b1715ec7 100644 --- a/crates/bevy_ecs/src/system/mod.rs +++ b/crates/bevy_ecs/src/system/mod.rs @@ -136,6 +136,7 @@ mod tests { bundle::Bundles, component::{Component, Components}, entity::{Entities, Entity}, + event::{EventReader, Events}, prelude::AnyOf, query::{Added, Changed, Or, With, Without}, schedule::{Schedule, Stage, SystemStage}, @@ -1187,4 +1188,55 @@ mod tests { } }); } + + #[test] + fn iter_join_map() { + struct DamageEvent { + target: Entity, + damage: f32, + } + + #[derive(Component, PartialEq, Debug)] + struct Health(f32); + + let mut world = World::new(); + + let e1 = world.spawn().insert(Health(200.)).id(); + let e2 = world.spawn().insert(Health(100.)).id(); + + let mut events = Events::::default(); + + events.extend([ + DamageEvent { + target: e1, + damage: 50., + }, + DamageEvent { + target: e2, + damage: 80., + }, + DamageEvent { + target: e1, + damage: 150., + }, + ]); + + world.insert_resource(events); + + fn process_damage_events( + mut health_query: Query<&mut Health>, + mut event_reader: EventReader, + ) { + let mut join = + health_query.iter_join_map_mut(event_reader.iter(), |event| event.target); + while let Some((mut health, event)) = join.fetch_next() { + health.0 -= event.damage; + } + } + run_system(&mut world, process_damage_events); + + run_system(&mut world, move |query: Query<&Health>| { + assert_eq!([&Health(0.), &Health(20.)], query.many([e1, e2])); + }); + } } From b5f79b145afead279dd88a8dba89d9d260d68bc6 Mon Sep 17 00:00:00 2001 From: devil-ira Date: Fri, 2 Sep 2022 10:53:14 +0200 Subject: [PATCH 3/6] Add `iter_join_map(_mut)`. --- crates/bevy_ecs/src/query/iter.rs | 149 ++++++++++++++++++++++++++++ crates/bevy_ecs/src/query/state.rs | 72 +++++++++++++- crates/bevy_ecs/src/system/query.rs | 95 +++++++++++++++++- 3 files changed, 312 insertions(+), 4 deletions(-) diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index b75b13aa54e92..8ba22f83b77ea 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -214,6 +214,155 @@ where { } +/// An [`Iterator`] over [`Query`](crate::system::Query) results a list mapped to [`Entity`]'s and the list items. +/// +/// This struct is created by the [`Query::iter_join_map`](crate::system::Query::iter_join_map) and [`Query::iter_join_map_mut`](crate::system::Query::iter_join_map_mut) methods. +pub struct QueryJoinMapIter<'w, 's, Q: WorldQuery, F: WorldQuery, I: Iterator, MapFn> +where + MapFn: FnMut(&I::Item) -> Entity, +{ + list: I, + map_f: MapFn, + entities: &'w Entities, + tables: &'w Tables, + archetypes: &'w Archetypes, + fetch: QueryFetch<'w, Q>, + filter: QueryFetch<'w, F>, + query_state: &'s QueryState, +} + +impl<'w, 's, Q: WorldQuery, F: WorldQuery, I: Iterator, MapFn> + QueryJoinMapIter<'w, 's, Q, F, I, MapFn> +where + MapFn: FnMut(&I::Item) -> Entity, +{ + /// # 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 `query_state.world_id`. Calling this on a `world` + /// with a mismatched [`WorldId`](crate::world::WorldId) is unsound. + pub(crate) unsafe fn new>( + world: &'w World, + query_state: &'s QueryState, + last_change_tick: u32, + change_tick: u32, + entity_map: II, + map_f: MapFn, + ) -> QueryJoinMapIter<'w, 's, Q, F, I, MapFn> { + let fetch = Q::init_fetch( + world, + &query_state.fetch_state, + last_change_tick, + change_tick, + ); + let filter = F::init_fetch( + world, + &query_state.filter_state, + last_change_tick, + change_tick, + ); + QueryJoinMapIter { + query_state, + entities: &world.entities, + archetypes: &world.archetypes, + tables: &world.storages.tables, + fetch, + filter, + list: entity_map.into_iter(), + map_f, + } + } + + /// SAFETY: + /// The lifetime here is not restrictive enough for Fetch with &mut access, + /// as calling `fetch_next_aliased_unchecked` multiple times can produce multiple + /// references to the same component, leading to unique reference aliasing. + /// + /// It is always safe for immutable borrows. + #[inline(always)] + unsafe fn fetch_next_aliased_unchecked(&mut self) -> Option<(QueryItem<'w, Q>, I::Item)> { + for item in self.list.by_ref() { + let location = match self.entities.get((self.map_f)(&item)) { + Some(location) => location, + None => continue, + }; + + if !self + .query_state + .matched_archetypes + .contains(location.archetype_id.index()) + { + continue; + } + + let archetype = &self.archetypes[location.archetype_id]; + + // SAFETY: `archetype` is from the world that `fetch/filter` were created for, + // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with + Q::set_archetype( + &mut self.fetch, + &self.query_state.fetch_state, + archetype, + self.tables, + ); + // SAFETY: `table` is from the world that `fetch/filter` were created for, + // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with + F::set_archetype( + &mut self.filter, + &self.query_state.filter_state, + archetype, + self.tables, + ); + // SAFETY: set_archetype was called prior. + // `location.index` is an archetype index row in range of the current archetype, because if it was not, the match above would have `continue`d + if F::archetype_filter_fetch(&mut self.filter, location.index) { + // SAFETY: set_archetype was called prior, `location.index` is an archetype index in range of the current archetype + return Some((Q::archetype_fetch(&mut self.fetch, location.index), item)); + } + } + None + } + + /// Get the next item from the iterator. + #[inline(always)] + pub fn fetch_next(&mut self) -> Option<(QueryItem<'_, Q>, I::Item)> { + // safety: we are limiting the returned reference to self, + // making sure this method cannot be called multiple times without getting rid + // of any previously returned unique references first, thus preventing aliasing. + unsafe { + self.fetch_next_aliased_unchecked() + .map(|(q_item, item)| (Q::shrink(q_item), item)) + } + } +} + +impl<'w, 's, Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, I: Iterator, MapFn> Iterator + for QueryJoinMapIter<'w, 's, Q, F, I, MapFn> +where + MapFn: FnMut(&I::Item) -> Entity, +{ + type Item = (QueryItem<'w, Q>, I::Item); + + #[inline(always)] + fn next(&mut self) -> Option { + // SAFETY: it is safe to alias for ReadOnlyWorldQuery + unsafe { self.fetch_next_aliased_unchecked() } + } + + fn size_hint(&self) -> (usize, Option) { + let (_, max_size) = self.list.size_hint(); + (0, max_size) + } +} + +// This is correct as QueryJoinMapIter always returns `None` once exhausted. +impl<'w, 's, Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, I: Iterator, MapFn> FusedIterator + for QueryJoinMapIter<'w, 's, Q, F, I, MapFn> +where + MapFn: FnMut(&I::Item) -> Entity, +{ +} + /// An iterator over `K`-sized combinations of query items without repetition. /// /// In this context, a combination is an unordered subset of `K` elements. diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index 9a1f08cde42d3..c6cdb1eaeea9d 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -3,7 +3,9 @@ use crate::{ component::ComponentId, entity::Entity, prelude::FromWorld, - query::{Access, FilteredAccess, QueryCombinationIter, QueryIter, WorldQuery}, + query::{ + Access, FilteredAccess, QueryCombinationIter, QueryIter, QueryJoinMapIter, WorldQuery, + }, storage::TableId, world::{World, WorldId}, }; @@ -642,6 +644,50 @@ impl QueryState { } } + /// Returns an [`Iterator`] over the inner join of the results of a query and list of items mapped to [`Entity`]'s. + /// + /// This can only be called for read-only queries, see [`Self::iter_list_mut`] for write-queries. + #[inline] + pub fn iter_join_map<'w, 's, I: IntoIterator, MapFn: FnMut(&I::Item) -> Entity>( + &'s mut self, + world: &'w World, + list: I, + map_f: MapFn, + ) -> QueryJoinMapIter<'w, 's, Q::ReadOnly, F::ReadOnly, I::IntoIter, MapFn> { + self.update_archetypes(world); + // SAFETY: Query has unique world access. + unsafe { + self.as_readonly().iter_join_map_unchecked_manual( + world, + world.last_change_tick(), + world.read_change_tick(), + list, + map_f, + ) + } + } + + /// Returns an iterator over the inner join of the results of a query and list of items mapped to [`Entity`]'s. + #[inline] + pub fn iter_join_map_mut<'w, 's, I: IntoIterator, MapFn: FnMut(&I::Item) -> Entity>( + &'s mut self, + world: &'w mut World, + list: I, + map_f: MapFn, + ) -> QueryJoinMapIter<'w, 's, Q, F, I::IntoIter, MapFn> { + self.update_archetypes(world); + // SAFETY: Query has unique world access. + unsafe { + self.iter_join_map_unchecked_manual( + world, + world.last_change_tick(), + world.read_change_tick(), + list, + map_f, + ) + } + } + /// Returns an [`Iterator`] over the query results for the given [`World`]. /// /// # Safety @@ -704,7 +750,6 @@ impl QueryState { /// /// 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 check for entity uniqueness /// This does not validate that `world.id()` matches `self.world_id`. Calling this on a `world` /// with a mismatched [`WorldId`] is unsound. #[inline] @@ -721,6 +766,29 @@ impl QueryState { QueryManyIter::new(world, self, entities, last_change_tick, change_tick) } + /// Returns an [`Iterator`] over each item in an inner join on [`Entity`] between the query and a list of items which are mapped to [`Entity`]'s. + /// + /// # 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. + #[inline] + pub(crate) unsafe fn iter_join_map_unchecked_manual<'w, 's, I: IntoIterator, MapFn>( + &'s self, + world: &'w World, + last_change_tick: u32, + change_tick: u32, + list: I, + map_f: MapFn, + ) -> QueryJoinMapIter<'w, 's, Q, F, I::IntoIter, MapFn> + where + MapFn: FnMut(&I::Item) -> Entity, + { + QueryJoinMapIter::new(world, self, last_change_tick, change_tick, list, map_f) + } + /// Returns an [`Iterator`] over all possible combinations of `K` query results for the /// given [`World`] without repetition. /// This can only be called for read-only queries. diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index ad6dcc02f1b45..5cf1bcfbbf053 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -2,8 +2,8 @@ use crate::{ component::Component, entity::Entity, query::{ - QueryCombinationIter, QueryEntityError, QueryItem, QueryIter, QueryManyIter, - QuerySingleError, QueryState, ROQueryItem, ReadOnlyWorldQuery, WorldQuery, + QueryCombinationIter, QueryEntityError, QueryItem, QueryIter, QueryJoinMapIter, + QueryManyIter, QuerySingleError, QueryState, ROQueryItem, ReadOnlyWorldQuery, WorldQuery, }, world::{Mut, World}, }; @@ -520,6 +520,97 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F> { } } + /// Returns an [`Iterator`] over the inner join of the results of a [`Query`] and list of items mapped to [`Entity`]'s. + /// + /// # Example + /// ``` + /// # use bevy_ecs::prelude::*; + /// # #[derive(Component)] + /// # struct Health { + /// # value: f32 + /// # } + /// #[derive(Component)] + /// struct DamageEvent { + /// target: Entity, + /// damage: f32, + /// } + /// + /// fn system( + /// mut damage_events: EventReader, + /// health_query: Query<&Health>, + /// ) { + /// for (health, event) in + /// health_query.iter_join_map(damage_events.iter(), |event| event.target) + /// { + /// println!("Entity has {} health and will take {} damage!", health.value, event.damage); + /// } + /// } + /// # bevy_ecs::system::assert_is_system(system); + /// ``` + #[inline] + pub fn iter_join_map Entity>( + &self, + list: I, + map_f: MapFn, + ) -> QueryJoinMapIter<'w, 's, Q::ReadOnly, F::ReadOnly, I::IntoIter, MapFn> { + // SAFETY: system runs without conflicts with other systems. + // same-system queries have runtime borrow checks when they conflict + unsafe { + self.state.as_readonly().iter_join_map_unchecked_manual( + self.world, + self.last_change_tick, + self.change_tick, + list, + map_f, + ) + } + } + + /// Returns an iterator over the inner join of the results of a [`Query`] and list of items mapped to [`Entity`]'s. + /// + /// # Example + /// ``` + /// # use bevy_ecs::prelude::*; + /// # #[derive(Component)] + /// # struct Health { + /// # value: f32 + /// # } + /// #[derive(Component)] + /// struct DamageEvent { + /// target: Entity, + /// damage: f32, + /// } + /// + /// fn system( + /// mut damage_events: EventReader, + /// mut health_query: Query<&mut Health>, + /// ) { + /// let mut join = health_query.iter_join_map_mut(damage_events.iter(), |event| event.target); + /// while let Some((mut health, event)) = join.fetch_next() { + /// health.value -= event.damage; + /// } + /// } + /// # bevy_ecs::system::assert_is_system(system); + /// ``` + #[inline] + pub fn iter_join_map_mut Entity>( + &mut self, + list: I, + map_f: MapFn, + ) -> QueryJoinMapIter<'_, '_, Q, F, I::IntoIter, MapFn> { + // SAFETY: system runs without conflicts with other systems. + // same-system queries have runtime borrow checks when they conflict + unsafe { + self.state.iter_join_map_unchecked_manual( + self.world, + self.last_change_tick, + self.change_tick, + list, + map_f, + ) + } + } + /// Returns an [`Iterator`] over the query results. /// /// # Safety From c5ceb8fc1fa2c258e076e7ff31f22ca995c38023 Mon Sep 17 00:00:00 2001 From: devil-ira Date: Fri, 2 Sep 2022 11:12:48 +0200 Subject: [PATCH 4/6] Fix compile fail tests --- .../tests/ui/query_mut_iterator_impl_safety.rs | 6 +++--- .../tests/ui/query_mut_iterator_impl_safety.stderr | 2 +- .../ui/system_query_iter_join_map_mut_lifetime_safety.rs | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.rs index 2bd4bbeef4535..0301821545cb2 100644 --- a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.rs +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.rs @@ -3,16 +3,16 @@ use bevy_ecs::prelude::*; #[derive(Component)] struct A(usize); -fn system(mut query: Query<&mut A>, e: Res) { +fn system(mut query: Query<&mut A>, e: Entity) { let iter = query.iter_combinations_mut(); // This should fail to compile. is_iterator(iter); - let iter = query.iter_many_mut([*e]); + let iter = query.iter_many_mut([e]); // This should fail to compile. is_iterator(iter); - let iter = query.iter_join_map_mut([*e], |e| *e); + let iter = query.iter_join_map_mut([e], |e| *e); // This should fail to compile. is_iterator(iter); } diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.stderr index 159ac22f8f3ad..97ab6f4d6ee93 100644 --- a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.stderr +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_mut_iterator_impl_safety.stderr @@ -69,7 +69,7 @@ error[E0277]: the trait bound `&mut A: ReadOnlyWorldQuery` is not satisfied (F0, F1, F2, F3, F4, F5, F6) and 49 others = note: `ReadOnlyWorldQuery` is implemented for `&A`, but not for `&mut A` - = note: required because of the requirements on the impl of `Iterator` for `QueryJoinMapIter<'_, '_, &mut A, (), std::array::IntoIter, [closure@$DIR/tests/ui/query_mut_iterator_impl_safety.rs:15:46: 15:52]>` + = note: required because of the requirements on the impl of `Iterator` for `QueryJoinMapIter<'_, '_, &mut A, (), std::array::IntoIter, [closure@$DIR/tests/ui/query_mut_iterator_impl_safety.rs:15:45: 15:51]>` note: required by a bound in `is_iterator` --> tests/ui/query_mut_iterator_impl_safety.rs:20:28 | diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_join_map_mut_lifetime_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_join_map_mut_lifetime_safety.rs index 5b5e665fa8daa..0b45766286869 100644 --- a/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_join_map_mut_lifetime_safety.rs +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_iter_join_map_mut_lifetime_safety.rs @@ -3,9 +3,9 @@ use bevy_ecs::prelude::*; #[derive(Component)] struct A(usize); -fn system(mut query: Query<&mut A>, e: Res) { +fn system(mut query: Query<&mut A>, e: Entity) { let mut results = Vec::new(); - let mut iter = query.iter_join_map_mut([*e, *e], |e| *e); + let mut iter = query.iter_join_map_mut([e], |e| *e); while let Some(a) = iter.fetch_next() { // this should fail to compile results.push(a); From dee495e53240cf416d5e1849360f5f4e9f5c3ce7 Mon Sep 17 00:00:00 2001 From: ira Date: Fri, 2 Sep 2022 12:25:37 +0200 Subject: [PATCH 5/6] Update crates/bevy_ecs/src/query/iter.rs Co-authored-by: Tristan Guichaoua <33934311+tguichaoua@users.noreply.github.com> --- crates/bevy_ecs/src/query/iter.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 8ba22f83b77ea..73411deea0df8 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -241,12 +241,12 @@ where /// have unique access to the components they query. /// This does not validate that `world.id()` matches `query_state.world_id`. Calling this on a `world` /// with a mismatched [`WorldId`](crate::world::WorldId) is unsound. - pub(crate) unsafe fn new>( + pub(crate) unsafe fn new( world: &'w World, query_state: &'s QueryState, last_change_tick: u32, change_tick: u32, - entity_map: II, + entity_map: impl IntoIterator, map_f: MapFn, ) -> QueryJoinMapIter<'w, 's, Q, F, I, MapFn> { let fetch = Q::init_fetch( From fb71220c1bbffa94262b2b033d480b441a94d88f Mon Sep 17 00:00:00 2001 From: devil-ira Date: Fri, 2 Sep 2022 12:45:36 +0200 Subject: [PATCH 6/6] Fix doc --- crates/bevy_ecs/src/query/state.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index c6cdb1eaeea9d..2a902835e4c9c 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -646,7 +646,7 @@ impl QueryState { /// Returns an [`Iterator`] over the inner join of the results of a query and list of items mapped to [`Entity`]'s. /// - /// This can only be called for read-only queries, see [`Self::iter_list_mut`] for write-queries. + /// This can only be called for read-only queries, see [`Self::iter_join_map_mut`] for write-queries. #[inline] pub fn iter_join_map<'w, 's, I: IntoIterator, MapFn: FnMut(&I::Item) -> Entity>( &'s mut self,