From 8b33b9183600bbe7ca2161858a848f9669392921 Mon Sep 17 00:00:00 2001 From: Chris Russell <8494645+chescock@users.noreply.github.com> Date: Mon, 16 Dec 2024 19:06:33 -0500 Subject: [PATCH] Always `collect()` when using `QueryIterMany::sort` methods. (#16844) # Objective When calling any of the `sort` methods on a `QueryManyIter` with mutable data, `collect_inner()` must be called before fetching items. Remove the need for that call. ## Solution Have the `sort` methods `collect()` the entity list into a `Vec` before returning. --- crates/bevy_ecs/src/query/iter.rs | 144 +++++++++++++++++++++++------- 1 file changed, 114 insertions(+), 30 deletions(-) diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 4de38258e9edc..021737b36f88d 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -16,7 +16,6 @@ use core::{ }; use super::{QueryData, QueryFilter, ReadOnlyQueryData}; -use alloc::vec::IntoIter; /// An [`Iterator`] over query results of a [`Query`](crate::system::Query). /// @@ -1453,8 +1452,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> /// # let entity_list: Vec = Vec::new(); /// // We need to collect the internal iterator before iterating mutably /// let mut parent_query_iter = query.iter_many_mut(entity_list) - /// .sort::() - /// .collect_inner(); + /// .sort::(); /// /// let mut scratch_value = 0; /// while let Some(mut part_value) = parent_query_iter.fetch_next_back() @@ -1500,7 +1498,14 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> .map(|(key, entity)| (key, NeutralOrd(entity))) .collect(); keyed_query.sort(); - let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity.0); + // Re-collect into a `Vec` to eagerly drop the lens items. + // They must be dropped before `fetch_next` is called since they may alias + // with other data items if there are duplicate entities in `entity_iter`. + let entity_iter = keyed_query + .into_iter() + .map(|(.., entity)| entity.0) + .collect::>() + .into_iter(); // SAFETY: // `self.world` has permission to access the required components. // Each lens query item is dropped before the respective actual query item is accessed. @@ -1589,7 +1594,14 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> .map(|(key, entity)| (key, NeutralOrd(entity))) .collect(); keyed_query.sort_unstable(); - let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity.0); + // Re-collect into a `Vec` to eagerly drop the lens items. + // They must be dropped before `fetch_next` is called since they may alias + // with other data items if there are duplicate entities in `entity_iter`. + let entity_iter = keyed_query + .into_iter() + .map(|(.., entity)| entity.0) + .collect::>() + .into_iter(); // SAFETY: // `self.world` has permission to access the required components. // Each lens query item is dropped before the respective actual query item is accessed. @@ -1681,7 +1693,14 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> }; let mut keyed_query: Vec<_> = query_lens.collect(); keyed_query.sort_by(|(key_1, _), (key_2, _)| compare(key_1, key_2)); - let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity); + // Re-collect into a `Vec` to eagerly drop the lens items. + // They must be dropped before `fetch_next` is called since they may alias + // with other data items if there are duplicate entities in `entity_iter`. + let entity_iter = keyed_query + .into_iter() + .map(|(.., entity)| entity) + .collect::>() + .into_iter(); // SAFETY: // `self.world` has permission to access the required components. // Each lens query item is dropped before the respective actual query item is accessed. @@ -1739,7 +1758,14 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> }; let mut keyed_query: Vec<_> = query_lens.collect(); keyed_query.sort_unstable_by(|(key_1, _), (key_2, _)| compare(key_1, key_2)); - let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity); + // Re-collect into a `Vec` to eagerly drop the lens items. + // They must be dropped before `fetch_next` is called since they may alias + // with other data items if there are duplicate entities in `entity_iter`. + let entity_iter = keyed_query + .into_iter() + .map(|(.., entity)| entity) + .collect::>() + .into_iter(); // SAFETY: // `self.world` has permission to access the required components. // Each lens query item is dropped before the respective actual query item is accessed. @@ -1863,7 +1889,14 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> }; let mut keyed_query: Vec<_> = query_lens.collect(); keyed_query.sort_by_key(|(lens, _)| f(lens)); - let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity); + // Re-collect into a `Vec` to eagerly drop the lens items. + // They must be dropped before `fetch_next` is called since they may alias + // with other data items if there are duplicate entities in `entity_iter`. + let entity_iter = keyed_query + .into_iter() + .map(|(.., entity)| entity) + .collect::>() + .into_iter(); // SAFETY: // `self.world` has permission to access the required components. // Each lens query item is dropped before the respective actual query item is accessed. @@ -1924,7 +1957,14 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> }; let mut keyed_query: Vec<_> = query_lens.collect(); keyed_query.sort_unstable_by_key(|(lens, _)| f(lens)); - let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity); + // Re-collect into a `Vec` to eagerly drop the lens items. + // They must be dropped before `fetch_next` is called since they may alias + // with other data items if there are duplicate entities in `entity_iter`. + let entity_iter = keyed_query + .into_iter() + .map(|(.., entity)| entity) + .collect::>() + .into_iter(); // SAFETY: // `self.world` has permission to access the required components. // Each lens query item is dropped before the respective actual query item is accessed. @@ -1985,7 +2025,14 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> }; let mut keyed_query: Vec<_> = query_lens.collect(); keyed_query.sort_by_cached_key(|(lens, _)| f(lens)); - let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity); + // Re-collect into a `Vec` to eagerly drop the lens items. + // They must be dropped before `fetch_next` is called since they may alias + // with other data items if there are duplicate entities in `entity_iter`. + let entity_iter = keyed_query + .into_iter() + .map(|(.., entity)| entity) + .collect::>() + .into_iter(); // SAFETY: // `self.world` has permission to access the required components. // Each lens query item is dropped before the respective actual query item is accessed. @@ -2178,25 +2225,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> unsafe { D::fetch(&mut self.fetch, entity, location.table_row) } } - /// Collects the internal [`I`](QuerySortedManyIter) once. - /// [`fetch_next`](QuerySortedManyIter) and [`fetch_next_back`](QuerySortedManyIter) require this to be called first. - #[inline(always)] - pub fn collect_inner(self) -> QuerySortedManyIter<'w, 's, D, F, IntoIter> { - QuerySortedManyIter { - entity_iter: self.entity_iter.collect::>().into_iter(), - entities: self.entities, - tables: self.tables, - archetypes: self.archetypes, - fetch: self.fetch, - query_state: self.query_state, - } - } -} - -impl<'w, 's, D: QueryData, F: QueryFilter> QuerySortedManyIter<'w, 's, D, F, IntoIter> { /// Get next result from the query - /// [`collect_inner`](QuerySortedManyIter) needs to be called before this method becomes available. - /// This is done to prevent mutable aliasing. #[inline(always)] pub fn fetch_next(&mut self) -> Option> { let entity = self.entity_iter.next()?; @@ -2210,10 +2239,12 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QuerySortedManyIter<'w, 's, D, F, Int // `entity` is passed from `entity_iter` the first time. unsafe { D::shrink(self.fetch_next_aliased_unchecked(entity)).into() } } +} +impl<'w, 's, D: QueryData, F: QueryFilter, I: DoubleEndedIterator> + QuerySortedManyIter<'w, 's, D, F, I> +{ /// Get next result from the query - /// [`collect_inner`](QuerySortedManyIter) needs to be called before this method becomes available. - /// This is done to prevent mutable aliasing. #[inline(always)] pub fn fetch_next_back(&mut self) -> Option> { let entity = self.entity_iter.next_back()?; @@ -3091,4 +3122,57 @@ mod tests { iter_2.sort::(); } + + // This test should be run with miri to check for UB caused by aliasing. + // The lens items created during the sort must not be live at the same time as the mutable references returned from the iterator. + #[test] + fn query_iter_many_sorts_duplicate_entities_no_ub() { + #[derive(Component, Ord, PartialOrd, Eq, PartialEq)] + struct C(usize); + + let mut world = World::new(); + let id = world.spawn(C(10)).id(); + let mut query_state = world.query::<&mut C>(); + + { + let mut query = query_state.iter_many_mut(&mut world, [id, id]).sort::<&C>(); + while query.fetch_next().is_some() {} + } + { + let mut query = query_state + .iter_many_mut(&mut world, [id, id]) + .sort_unstable::<&C>(); + while query.fetch_next().is_some() {} + } + { + let mut query = query_state + .iter_many_mut(&mut world, [id, id]) + .sort_by::<&C>(Ord::cmp); + while query.fetch_next().is_some() {} + } + { + let mut query = query_state + .iter_many_mut(&mut world, [id, id]) + .sort_unstable_by::<&C>(Ord::cmp); + while query.fetch_next().is_some() {} + } + { + let mut query = query_state + .iter_many_mut(&mut world, [id, id]) + .sort_by_key::<&C, _>(|d| d.0); + while query.fetch_next().is_some() {} + } + { + let mut query = query_state + .iter_many_mut(&mut world, [id, id]) + .sort_unstable_by_key::<&C, _>(|d| d.0); + while query.fetch_next().is_some() {} + } + { + let mut query = query_state + .iter_many_mut(&mut world, [id, id]) + .sort_by_cached_key::<&C, _>(|d| d.0); + while query.fetch_next().is_some() {} + } + } }