forestia/bevy
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Always collect() when using QueryIterMany::sort methods. (#16844)

Browse Source 
# 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.
This commit is contained in:
Chris Russell 2024-12-16 19:06:33 -05:00 committed by GitHub
parent 5f4b5a37f1
commit 8b33b91836
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 114 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

144
crates/bevy_ecs/src/query/iter.rs
View File

@ -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<Item: Borrow<Entity>>>
/// # let entity_list: Vec<Entity> = Vec::new();
/// // We need to collect the internal iterator before iterating mutably
/// let mut parent_query_iter = query.iter_many_mut(entity_list)
/// .sort::<Entity>()
/// .collect_inner();
/// .sort::<Entity>();
///
/// 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<Item: Borrow<Entity>>>
.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::<Vec<_>>()
.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<Item: Borrow<Entity>>>
.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::<Vec<_>>()
.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<Item: Borrow<Entity>>>
};
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::<Vec<_>>()
.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<Item: Borrow<Entity>>>
};
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::<Vec<_>>()
.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<Item: Borrow<Entity>>>
};
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::<Vec<_>>()
.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<Item: Borrow<Entity>>>
};
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::<Vec<_>>()
.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<Item: Borrow<Entity>>>
};
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::<Vec<_>>()
.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<Item = Entity>>
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<Entity>> {
QuerySortedManyIter {
entity_iter: self.entity_iter.collect::<Vec<_>>().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<Entity>> {
/// 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<D::Item<'_>> {
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<Item = Entity>>
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<D::Item<'_>> {
let entity = self.entity_iter.next_back()?;
@ -3091,4 +3122,57 @@ mod tests {
iter_2.sort::<Entity>();
}
// 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() {}
}
}
}