From 3c502c4a6cdf7fe4d6c3e102cacca3f238ca5d2b Mon Sep 17 00:00:00 2001 From: Victoron <59878206+Victoronz@users.noreply.github.com> Date: Wed, 22 May 2024 21:33:04 +0200 Subject: [PATCH 1/3] add entities_all_unique to QueryManyIter --- crates/bevy_ecs/src/query/iter.rs | 192 +++++++++++++++++++++++++++++- 1 file changed, 191 insertions(+), 1 deletion(-) diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index d09c756491b4e..12164073d6e5b 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -1,7 +1,7 @@ use crate::{ archetype::{Archetype, ArchetypeEntity, Archetypes}, component::Tick, - entity::{Entities, Entity}, + entity::{Entities, Entity, EntityHashSet}, query::{ArchetypeFilter, DebugCheckedUnwrap, QueryState, StorageId}, storage::{Table, TableRow, Tables}, world::unsafe_world_cell::UnsafeWorldCell, @@ -13,6 +13,7 @@ use std::{ iter::FusedIterator, mem::MaybeUninit, ops::Range, + vec::IntoIter, }; use super::{QueryData, QueryFilter, ReadOnlyQueryData}; @@ -1123,6 +1124,83 @@ where // of any previously returned unique references first, thus preventing aliasing. unsafe { self.fetch_next_aliased_unchecked().map(D::shrink) } } + + /// Checks for uniqueness in the `Entity` iterator `I`, returning a new Iterator on success. + /// Return `self` on failure. + /// This new iterator allows for mutable iteration without `fetch_next`. + /// # Example + /// ``` + /// # use bevy_ecs::prelude::*; + /// # use std::ops::{Deref, DerefMut}; + /// # + /// # #[derive(Component, Clone, Copy)] + /// # struct PartValue(usize); + /// # + /// # impl Deref for PartValue { + /// # type Target = usize; + /// # + /// # fn deref(&self) -> &Self::Target { + /// # &self.0 + /// # } + /// # } + /// # + /// # impl DerefMut for PartValue { + /// # fn deref_mut(&mut self) -> &mut Self::Target { + /// # &mut self.0 + /// # } + /// # } + /// # + /// # let mut world = World::new(); + /// # + /// // Mutable `Iterator` trait iteration. + /// fn system(mut query: Query<&mut PartValue>) { + /// # let entity_list: Vec = Vec::new(); + /// # + /// let mut unique_iter = query.iter_many_mut(entity_list) + /// .entities_all_unique() + /// .expect("the entity_list only contains unique entities"); + /// + /// for mut part_value in unique_iter { + /// **part_value += 1; + /// } + /// } + /// # + /// # let mut schedule = Schedule::default(); + /// # schedule.add_systems((system)); + /// # schedule.run(&mut world); + /// ``` + #[inline(always)] + pub fn entities_all_unique( + self, + ) -> Result< + QueryManyUniqueIter<'w, 's, D, F, IntoIter>, + QueryManyIter<'w, 's, D, F, IntoIter>, + > { + let mut used = EntityHashSet::default(); + let entities: Vec<_> = self.entity_iter.collect(); + + if entities.iter().all(move |e| used.insert(*e.borrow())) { + return Ok(QueryManyUniqueIter { + entity_iter: entities.into_iter(), + entities: self.entities, + tables: self.tables, + archetypes: self.archetypes, + fetch: self.fetch, + filter: self.filter, + query_state: self.query_state, + }); + } + + Err(QueryManyIter { + entity_iter: entities.into_iter(), + entities: self.entities, + tables: self.tables, + archetypes: self.archetypes, + fetch: self.fetch, + filter: self.filter, + query_state: self.query_state, + }) + } } impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator> Iterator @@ -1161,6 +1239,118 @@ where } } +/// An [`Iterator`] over the query items generated from an iterator of unique [`Entity`]s. +/// +/// Items are returned in the order of the provided iterator. +/// Entities that don't match the query are skipped. +/// +/// In contrast with `QueryManyIter`, this allows for mutable iteration without a `fetch_next` method. +/// +/// This struct is created by the [`QueryManyIter::entities_all_unique`] method. +pub struct QueryManyUniqueIter<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> +where + I::Item: Borrow, +{ + entity_iter: I, + entities: &'w Entities, + tables: &'w Tables, + archetypes: &'w Archetypes, + fetch: D::Fetch<'w>, + filter: F::Fetch<'w>, + query_state: &'s QueryState, +} + +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> QueryManyUniqueIter<'w, 's, D, F, I> +where + I::Item: Borrow, +{ + // Entities are guaranteed to be unique, so no lifetime shrinking needed for mutable iteration. + #[inline(always)] + fn fetch_next(&mut self) -> Option> { + for entity in self.entity_iter.by_ref() { + let entity = *entity.borrow(); + let Some(location) = self.entities.get(entity) else { + continue; + }; + + if !self + .query_state + .matched_archetypes + .contains(location.archetype_id.index()) + { + continue; + } + let (archetype, table); + // SAFETY: + // `tables` and `archetypes` belong to the same world that the [`QueryIter`] + // was initialized for. + unsafe { + archetype = self + .archetypes + .get(location.archetype_id) + .debug_checked_unwrap(); + table = self.tables.get(location.table_id).debug_checked_unwrap(); + } + // 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 + unsafe { + D::set_archetype( + &mut self.fetch, + &self.query_state.fetch_state, + archetype, + table, + ); + } + // 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 + unsafe { + F::set_archetype( + &mut self.filter, + &self.query_state.filter_state, + archetype, + table, + ); + } + + // SAFETY: set_archetype was called prior. + // `location.archetype_row` is an archetype index row in range of the current archetype, because if it was not, the match above would have `continue`d + if unsafe { F::filter_fetch(&mut self.filter, entity, location.table_row) } { + // SAFETY: + // - set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype + // - fetch is only called once for each entity. + return Some(unsafe { D::fetch(&mut self.fetch, entity, location.table_row) }); + } + } + None + } +} + +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Iterator + for QueryManyUniqueIter<'w, 's, D, F, I> +where + I::Item: Borrow, +{ + type Item = D::Item<'w>; + + #[inline(always)] + fn next(&mut self) -> Option { + self.fetch_next() + } + + fn size_hint(&self) -> (usize, Option) { + let (_, max_size) = self.entity_iter.size_hint(); + (0, max_size) + } +} + +// This is correct as [`QueryManyIter`] always returns `None` once exhausted. +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> FusedIterator + for QueryManyUniqueIter<'w, 's, D, F, I> +where + I::Item: Borrow, +{ +} + /// An iterator over `K`-sized combinations of query items without repetition. /// /// A combination is an arrangement of a collection of items where order does not matter. From 1d63a3e76ecfb450454f80c8323ee793f8758133 Mon Sep 17 00:00:00 2001 From: Victoron <59878206+Victoronz@users.noreply.github.com> Date: Tue, 2 Jul 2024 19:28:38 +0200 Subject: [PATCH 2/3] make QueryManyUniqueIter a wrapper around QueryManyIter also adds: a Debug impl associated type bounds in the struct and impl definitions a missing unsafe block in QueryManyIter --- crates/bevy_ecs/src/query/iter.rs | 127 ++++++++---------------------- 1 file changed, 34 insertions(+), 93 deletions(-) diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 12164073d6e5b..76e1e3d81607c 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -1077,11 +1077,17 @@ where continue; } - let archetype = self - .archetypes - .get(location.archetype_id) - .debug_checked_unwrap(); - let table = self.tables.get(location.table_id).debug_checked_unwrap(); + let (archetype, table); + // SAFETY: + // `tables` and `archetypes` belong to the same world that the [`QueryIter`] + // was initialized for. + unsafe { + archetype = self + .archetypes + .get(location.archetype_id) + .debug_checked_unwrap(); + table = self.tables.get(location.table_id).debug_checked_unwrap(); + } // 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 @@ -1180,7 +1186,7 @@ where let entities: Vec<_> = self.entity_iter.collect(); if entities.iter().all(move |e| used.insert(*e.borrow())) { - return Ok(QueryManyUniqueIter { + return Ok(QueryManyUniqueIter(QueryManyIter { entity_iter: entities.into_iter(), entities: self.entities, tables: self.tables, @@ -1188,7 +1194,7 @@ where fetch: self.fetch, filter: self.filter, query_state: self.query_state, - }); + })); } Err(QueryManyIter { @@ -1247,110 +1253,45 @@ where /// In contrast with `QueryManyIter`, this allows for mutable iteration without a `fetch_next` method. /// /// This struct is created by the [`QueryManyIter::entities_all_unique`] method. -pub struct QueryManyUniqueIter<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> -where - I::Item: Borrow, -{ - entity_iter: I, - entities: &'w Entities, - tables: &'w Tables, - archetypes: &'w Archetypes, - fetch: D::Fetch<'w>, - filter: F::Fetch<'w>, - query_state: &'s QueryState, -} - -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> QueryManyUniqueIter<'w, 's, D, F, I> -where - I::Item: Borrow, -{ - // Entities are guaranteed to be unique, so no lifetime shrinking needed for mutable iteration. - #[inline(always)] - fn fetch_next(&mut self) -> Option> { - for entity in self.entity_iter.by_ref() { - let entity = *entity.borrow(); - let Some(location) = self.entities.get(entity) else { - continue; - }; - - if !self - .query_state - .matched_archetypes - .contains(location.archetype_id.index()) - { - continue; - } - let (archetype, table); - // SAFETY: - // `tables` and `archetypes` belong to the same world that the [`QueryIter`] - // was initialized for. - unsafe { - archetype = self - .archetypes - .get(location.archetype_id) - .debug_checked_unwrap(); - table = self.tables.get(location.table_id).debug_checked_unwrap(); - } - // 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 - unsafe { - D::set_archetype( - &mut self.fetch, - &self.query_state.fetch_state, - archetype, - table, - ); - } - // 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 - unsafe { - F::set_archetype( - &mut self.filter, - &self.query_state.filter_state, - archetype, - table, - ); - } - - // SAFETY: set_archetype was called prior. - // `location.archetype_row` is an archetype index row in range of the current archetype, because if it was not, the match above would have `continue`d - if unsafe { F::filter_fetch(&mut self.filter, entity, location.table_row) } { - // SAFETY: - // - set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype - // - fetch is only called once for each entity. - return Some(unsafe { D::fetch(&mut self.fetch, entity, location.table_row) }); - } - } - None - } -} - -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Iterator +pub struct QueryManyUniqueIter< + 'w, + 's, + D: QueryData, + F: QueryFilter, + I: Iterator>, +>(QueryManyIter<'w, 's, D, F, I>); + +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> Iterator for QueryManyUniqueIter<'w, 's, D, F, I> -where - I::Item: Borrow, { type Item = D::Item<'w>; #[inline(always)] fn next(&mut self) -> Option { - self.fetch_next() + // SAFETY: Entities are guaranteed to be unique, thus do not alias. + unsafe { self.0.fetch_next_aliased_unchecked() } } fn size_hint(&self) -> (usize, Option) { - let (_, max_size) = self.entity_iter.size_hint(); + let (_, max_size) = self.0.entity_iter.size_hint(); (0, max_size) } } // This is correct as [`QueryManyIter`] always returns `None` once exhausted. -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> FusedIterator +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> FusedIterator for QueryManyUniqueIter<'w, 's, D, F, I> -where - I::Item: Borrow, { } +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> Debug + for QueryManyUniqueIter<'w, 's, D, F, I> +{ + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("QueryManyUniqueIter").finish() + } +} + /// An iterator over `K`-sized combinations of query items without repetition. /// /// A combination is an arrangement of a collection of items where order does not matter. From 122dc8df6882ce38a784d837dd7846a6125c7a20 Mon Sep 17 00:00:00 2001 From: Victoron <59878206+Victoronz@users.noreply.github.com> Date: Tue, 2 Jul 2024 19:44:54 +0200 Subject: [PATCH 3/3] add dedup_entities method --- crates/bevy_ecs/src/query/iter.rs | 67 +++++++++++++++++++++++++++++-- 1 file changed, 63 insertions(+), 4 deletions(-) diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 76e1e3d81607c..f20d82252ee68 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -1131,9 +1131,9 @@ where unsafe { self.fetch_next_aliased_unchecked().map(D::shrink) } } - /// Checks for uniqueness in the `Entity` iterator `I`, returning a new Iterator on success. - /// Return `self` on failure. - /// This new iterator allows for mutable iteration without `fetch_next`. + /// Checks for uniqueness in the `Entity` iterator `I`, returning a new `Iterator` on success. + /// Returns `self` on failure. + /// The new iterator allows for mutable iteration without `fetch_next`. /// # Example /// ``` /// # use bevy_ecs::prelude::*; @@ -1207,6 +1207,65 @@ where query_state: self.query_state, }) } + + /// Deduplicates the `Entity` iterator `I`, returning a new Iterator over unique entities. + /// The new iterator allows for mutable iteration without `fetch_next`. + /// # Example + /// ``` + /// # use bevy_ecs::prelude::*; + /// # use std::ops::{Deref, DerefMut}; + /// # + /// # #[derive(Component, Clone, Copy)] + /// # struct PartValue(usize); + /// # + /// # impl Deref for PartValue { + /// # type Target = usize; + /// # + /// # fn deref(&self) -> &Self::Target { + /// # &self.0 + /// # } + /// # } + /// # + /// # impl DerefMut for PartValue { + /// # fn deref_mut(&mut self) -> &mut Self::Target { + /// # &mut self.0 + /// # } + /// # } + /// # + /// # let mut world = World::new(); + /// # + /// // Mutable `Iterator` trait iteration. + /// fn system(mut query: Query<&mut PartValue>) { + /// # let entity_list: Vec = Vec::new(); + /// # + /// let mut unique_iter = query.iter_many_mut(entity_list) + /// .dedup_entities(); + /// + /// for mut part_value in unique_iter { + /// **part_value += 1; + /// } + /// } + /// # + /// # let mut schedule = Schedule::default(); + /// # schedule.add_systems((system)); + /// # schedule.run(&mut world); + /// ``` + #[inline(always)] + pub fn dedup_entities( + self, + ) -> QueryManyUniqueIter<'w, 's, D, F, impl Iterator>> { + let mut used = EntityHashSet::default(); + + QueryManyUniqueIter(QueryManyIter { + entity_iter: self.entity_iter.filter(move |e| used.insert(*e.borrow())), + entities: self.entities, + tables: self.tables, + archetypes: self.archetypes, + fetch: self.fetch, + filter: self.filter, + query_state: self.query_state, + }) + } } impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator> Iterator @@ -1252,7 +1311,7 @@ where /// /// In contrast with `QueryManyIter`, this allows for mutable iteration without a `fetch_next` method. /// -/// This struct is created by the [`QueryManyIter::entities_all_unique`] method. +/// This struct is created by the [`QueryManyIter::entities_all_unique`] and [`QueryManyIter::dedup_entities`] methods. pub struct QueryManyUniqueIter< 'w, 's,