Skip to content

implement get_many_unique #18315

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Mar 16, 2025
Merged

implement get_many_unique #18315

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
e1b1439
implement get_many_unique
Victoronz Mar 14, 2025
a2f9630
fix missed renames and improve docs
Victoronz Mar 14, 2025
5e6e71c
infer Entity properly
Victoronz Mar 14, 2025
70ba56b
Merge branch 'main' into get-many-unique
Victoronz Mar 15, 2025
f5c4c9c
Account for the new type default
Victoronz Mar 15, 2025
b6f4c6a
doc fixes
Victoronz Mar 15, 2025
6993690
elide remaining UniqueEntityArray defaults uses
Victoronz Mar 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 85 additions & 2 deletions crates/bevy_ecs/src/query/state.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use crate::{
archetype::{Archetype, ArchetypeComponentId, ArchetypeGeneration, ArchetypeId},
component::{ComponentId, Tick},
entity::{Entity, EntityBorrow, EntitySet},
entity::{unique_array::UniqueEntityArray, Entity, EntityBorrow, EntitySet},
entity_disabling::DefaultQueryFilters,
prelude::FromWorld,
query::{Access, FilteredAccess, QueryCombinationIter, QueryIter, QueryParIter, WorldQuery},
Expand Down Expand Up @@ -997,6 +997,44 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
self.query(world).get_many_inner(entities)
}

/// Returns the read-only query results for the given [`UniqueEntityArray`].
///
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is
/// returned instead.
///
/// # Examples
///
/// ```
/// use bevy_ecs::{prelude::*, query::QueryEntityError, entity::{EntitySetIterator, unique_array::UniqueEntityArray, unique_vec::UniqueEntityVec}};
///
/// #[derive(Component, PartialEq, Debug)]
/// struct A(usize);
///
/// let mut world = World::new();
/// let entity_set: UniqueEntityVec = world.spawn_batch((0..3).map(A)).collect_set();
/// let entity_set: UniqueEntityArray<3> = entity_set.try_into().unwrap();
///
/// world.spawn(A(73));
///
/// let mut query_state = world.query::<&A>();
///
/// let component_values = query_state.get_many_unique(&world, entity_set).unwrap();
///
/// assert_eq!(component_values, [&A(0), &A(1), &A(2)]);
///
/// let wrong_entity = Entity::from_raw(365);
///
/// assert_eq!(match query_state.get_many_unique(&mut world, UniqueEntityArray::from([wrong_entity])).unwrap_err() {QueryEntityError::EntityDoesNotExist(error) => error.entity, _ => panic!()}, wrong_entity);
/// ```
#[inline]
pub fn get_many_unique<'w, const N: usize>(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm hoping to deprecate the methods like this on QueryState, so I'd mildly prefer not to add these in the first place. The functionality would still be available as .query(world).get_many_unique_inner(entities).

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I was a little confused why this function needed to be implemented for both the state module and query module. It seems like it would only really be needed in the query module

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I was a little confused why this function needed to be implemented for both the state module and query module.

For context: Until recently, the actual implementation of all of these methods was on QueryState, and the Query methods just called those. I added the QueryState::query methods and moved the implementations over, with the eventual goal of getting rid of most of the other methods on QueryState.

Copy link
Contributor Author

@Victoronz Victoronz Mar 14, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am aware that we're trying to deprecate them!
However in the spirit of being incremental, I find it more confusing to not add them here when the other methods do not yet have the #[deprecated] moniker. Once we do deprecate them before 0.16, then feel free to remove these again!

Given that sometimes PRs can miss a cut-off, or contributors don't have time before an RC/release, leaving changes unfinished seems like an uneccessary risk, when both adding and removing is simple implementation-wise.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in the spirit of being incremental, I find it more confusing to not add them here

Well, in that case, you forgot get_many_unique_manual(&self, &World, entities) and get_many_unique_unchecked(&mut self, UnsafeWorldCell, entities) :P

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in the spirit of being incremental, I find it more confusing to not add them here

Well, in that case, you forgot get_many_unique_manual(&self, &World, entities) and get_many_unique_unchecked(&mut self, UnsafeWorldCell, entities) :P

These already don't exist for get_many, so I haven't added them here.
These inconsistencies are not nice, but we can at least not make them worse in the interim :P

&mut self,
world: &'w World,
entities: UniqueEntityArray<N, Entity>,
) -> Result<[ROQueryItem<'w, D>; N], QueryEntityError> {
self.query(world).get_many_unique_inner(entities)
}

/// Gets the query result for the given [`World`] and [`Entity`].
///
/// This is always guaranteed to run in `O(1)` time.
Expand Down Expand Up @@ -1053,7 +1091,52 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
world: &'w mut World,
entities: [Entity; N],
) -> Result<[D::Item<'w>; N], QueryEntityError> {
self.query_mut(world).get_many_inner(entities)
self.query_mut(world).get_many_mut_inner(entities)
}

/// Returns the query results for the given [`UniqueEntityArray`].
///
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is
/// returned instead.
///
/// ```
/// use bevy_ecs::{prelude::*, query::QueryEntityError, entity::{EntitySetIterator, unique_array::UniqueEntityArray, unique_vec::UniqueEntityVec}};
///
/// #[derive(Component, PartialEq, Debug)]
/// struct A(usize);
///
/// let mut world = World::new();
///
/// let entity_set: UniqueEntityVec = world.spawn_batch((0..3).map(A)).collect_set();
/// let entity_set: UniqueEntityArray<3> = entity_set.try_into().unwrap();
///
/// world.spawn(A(73));
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are spawning an entity here with component A(73) because you want to make sure that the query_state.get_many_unique_mut call only returns the entities in the UniqueEntityArray<Entity, 3> right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have copied and adjusted the existing doc tests for get_many and get_many_mut for this, so I am not the original writer, but that does seem to be the reason for that spawn.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cool beans! I think that an example that showcases this functionality would be good then cause you could comment why world.spawn(A(73)); was there.

Totally not blocking tho!

///
/// let mut query_state = world.query::<&mut A>();
///
/// let mut mutable_component_values = query_state.get_many_unique_mut(&mut world, entity_set).unwrap();
///
/// for mut a in &mut mutable_component_values {
/// a.0 += 5;
/// }
///
/// let component_values = query_state.get_many_unique(&world, entity_set).unwrap();
///
/// assert_eq!(component_values, [&A(5), &A(6), &A(7)]);
///
/// let wrong_entity = Entity::from_raw(57);
/// let invalid_entity = world.spawn_empty().id();
///
/// assert_eq!(match query_state.get_many_unique(&mut world, UniqueEntityArray::from([wrong_entity])).unwrap_err() {QueryEntityError::EntityDoesNotExist(error) => error.entity, _ => panic!()}, wrong_entity);
/// assert_eq!(match query_state.get_many_unique_mut(&mut world, UniqueEntityArray::from([invalid_entity])).unwrap_err() {QueryEntityError::QueryDoesNotMatch(entity, _) => entity, _ => panic!()}, invalid_entity);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is the reason this is returning an error because you expect every entity to have component 'A'? If so, why not just return an empty list of components?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These asserts test that the error cases are produced properly.
If an entity does not match the D in QueryData, that should return an error!
Same with an entity that does not exist in the first place.

/// ```
#[inline]
pub fn get_many_unique_mut<'w, const N: usize>(
&mut self,
world: &'w mut World,
entities: UniqueEntityArray<N, Entity>,
) -> Result<[D::Item<'w>; N], QueryEntityError> {
self.query_mut(world).get_many_unique_inner(entities)
}

/// Gets the query result for the given [`World`] and [`Entity`].
Expand Down
169 changes: 153 additions & 16 deletions crates/bevy_ecs/src/system/query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
use crate::{
batching::BatchingStrategy,
component::Tick,
entity::{Entity, EntityBorrow, EntityDoesNotExistError, EntitySet},
entity::{
unique_array::UniqueEntityArray, Entity, EntityBorrow, EntityDoesNotExistError, EntitySet,
},
query::{
DebugCheckedUnwrap, NopWorldQuery, QueryCombinationIter, QueryData, QueryEntityError,
QueryFilter, QueryIter, QueryManyIter, QueryManyUniqueIter, QueryParIter, QueryParManyIter,
Expand Down Expand Up @@ -1323,15 +1325,65 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
/// # See also
///
/// - [`get_many_mut`](Self::get_many_mut) to get mutable query items.
/// - [`get_many_unique`](Self::get_many_unique) to only handle unique inputs.
/// - [`many`](Self::many) for the panicking version.
#[inline]
pub fn get_many<const N: usize>(
&self,
entities: [Entity; N],
) -> Result<[ROQueryItem<'_, D>; N], QueryEntityError> {
// Note that this calls `get_many_readonly` instead of `get_many_inner`
// since we don't need to check for duplicates.
self.as_readonly().get_many_readonly(entities)
// Note that we call a separate `*_inner` method from `get_many_mut`
// because we don't need to check for duplicates.
self.as_readonly().get_many_inner(entities)
}

/// Returns the read-only query items for the given [`UniqueEntityArray`].
///
/// The returned query items are in the same order as the input.
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is returned instead.
///
/// # Examples
///
/// ```
/// use bevy_ecs::{prelude::*, query::QueryEntityError, entity::{EntitySetIterator, unique_array::UniqueEntityArray, unique_vec::UniqueEntityVec}};
///
/// #[derive(Component, PartialEq, Debug)]
/// struct A(usize);
///
/// let mut world = World::new();
/// let entity_set: UniqueEntityVec = world.spawn_batch((0..3).map(A)).collect_set();
/// let entity_set: UniqueEntityArray<3> = entity_set.try_into().unwrap();
///
/// world.spawn(A(73));
///
/// let mut query_state = world.query::<&A>();
/// let query = query_state.query(&world);
///
/// let component_values = query.get_many_unique(entity_set).unwrap();
///
/// assert_eq!(component_values, [&A(0), &A(1), &A(2)]);
///
/// let wrong_entity = Entity::from_raw(365);
///
/// assert_eq!(
/// match query.get_many_unique(UniqueEntityArray::from([wrong_entity])).unwrap_err() {
/// QueryEntityError::EntityDoesNotExist(error) => error.entity,
/// _ => panic!(),
/// },
/// wrong_entity
/// );
/// ```
///
/// # See also
///
/// - [`get_many_unique_mut`](Self::get_many_mut) to get mutable query items.
/// - [`get_many`](Self::get_many) to handle inputs with duplicates.
#[inline]
pub fn get_many_unique<const N: usize>(
&self,
entities: UniqueEntityArray<N>,
) -> Result<[ROQueryItem<'_, D>; N], QueryEntityError> {
self.as_readonly().get_many_unique_inner(entities)
}

/// Returns the read-only query items for the given array of [`Entity`].
Expand Down Expand Up @@ -1560,7 +1612,75 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
&mut self,
entities: [Entity; N],
) -> Result<[D::Item<'_>; N], QueryEntityError> {
self.reborrow().get_many_inner(entities)
self.reborrow().get_many_mut_inner(entities)
}

/// Returns the query items for the given [`UniqueEntityArray`].
///
/// The returned query items are in the same order as the input.
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is returned instead.
///
/// # Examples
///
/// ```
/// use bevy_ecs::{prelude::*, query::QueryEntityError, entity::{EntitySetIterator, unique_array::UniqueEntityArray, unique_vec::UniqueEntityVec}};
///
/// #[derive(Component, PartialEq, Debug)]
/// struct A(usize);
///
/// let mut world = World::new();
///
/// let entity_set: UniqueEntityVec<_> = world.spawn_batch((0..3).map(A)).collect_set();
/// let entity_set: UniqueEntityArray<3> = entity_set.try_into().unwrap();
///
/// world.spawn(A(73));
/// let wrong_entity = Entity::from_raw(57);
/// let invalid_entity = world.spawn_empty().id();
///
///
/// let mut query_state = world.query::<&mut A>();
/// let mut query = query_state.query_mut(&mut world);
///
/// let mut mutable_component_values = query.get_many_unique_mut(entity_set).unwrap();
///
/// for mut a in &mut mutable_component_values {
/// a.0 += 5;
/// }
///
/// let component_values = query.get_many_unique(entity_set).unwrap();
///
/// assert_eq!(component_values, [&A(5), &A(6), &A(7)]);
///
/// assert_eq!(
/// match query
/// .get_many_unique_mut(UniqueEntityArray::from([wrong_entity]))
/// .unwrap_err()
/// {
/// QueryEntityError::EntityDoesNotExist(error) => error.entity,
/// _ => panic!(),
/// },
/// wrong_entity
/// );
/// assert_eq!(
/// match query
/// .get_many_unique_mut(UniqueEntityArray::from([invalid_entity]))
/// .unwrap_err()
/// {
/// QueryEntityError::QueryDoesNotMatch(entity, _) => entity,
/// _ => panic!(),
/// },
/// invalid_entity
/// );
/// ```
/// # See also
///
/// - [`get_many_unique`](Self::get_many) to get read-only query items.
#[inline]
pub fn get_many_unique_mut<const N: usize>(
&mut self,
entities: UniqueEntityArray<N>,
) -> Result<[D::Item<'_>; N], QueryEntityError> {
self.reborrow().get_many_unique_inner(entities)
}

/// Returns the query items for the given array of [`Entity`].
Expand All @@ -1573,10 +1693,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
///
/// - [`get_many`](Self::get_many) to get read-only query items without checking for duplicate entities.
/// - [`get_many_mut`](Self::get_many_mut) to get items using a mutable reference.
/// - [`get_many_readonly`](Self::get_many_readonly) to get read-only query items without checking for duplicate entities
/// with the actual "inner" world lifetime.
/// - [`get_many_inner`](Self::get_many_mut_inner) to get read-only query items with the actual "inner" world lifetime.
#[inline]
pub fn get_many_inner<const N: usize>(
pub fn get_many_mut_inner<const N: usize>(
self,
entities: [Entity; N],
) -> Result<[D::Item<'w>; N], QueryEntityError> {
Expand All @@ -1588,7 +1707,6 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
}
}
}

// SAFETY: All entities are unique, so the results don't alias.
unsafe { self.get_many_impl(entities) }
}
Expand All @@ -1603,9 +1721,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
///
/// - [`get_many`](Self::get_many) to get read-only query items without checking for duplicate entities.
/// - [`get_many_mut`](Self::get_many_mut) to get items using a mutable reference.
/// - [`get_many_inner`](Self::get_many_readonly) to get mutable query items with the actual "inner" world lifetime.
/// - [`get_many_mut_inner`](Self::get_many_mut_inner) to get mutable query items with the actual "inner" world lifetime.
#[inline]
pub fn get_many_readonly<const N: usize>(
pub fn get_many_inner<const N: usize>(
self,
entities: [Entity; N],
) -> Result<[D::Item<'w>; N], QueryEntityError>
Expand All @@ -1616,6 +1734,25 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
unsafe { self.get_many_impl(entities) }
}

/// Returns the query items for the given [`UniqueEntityArray`].
/// This consumes the [`Query`] to return results with the actual "inner" world lifetime.
///
/// The returned query items are in the same order as the input.
/// In case of a nonexisting entity, duplicate entities or mismatched component, a [`QueryEntityError`] is returned instead.
///
/// # See also
///
/// - [`get_many_unique`](Self::get_many_unique) to get read-only query items without checking for duplicate entities.
/// - [`get_many_unique_mut`](Self::get_many_unique_mut) to get items using a mutable reference.
#[inline]
pub fn get_many_unique_inner<const N: usize>(
self,
entities: UniqueEntityArray<N>,
) -> Result<[D::Item<'w>; N], QueryEntityError> {
// SAFETY: All entities are unique, so the results don't alias.
unsafe { self.get_many_impl(entities.into_inner()) }
}

/// Returns the query items for the given array of [`Entity`].
/// This consumes the [`Query`] to return results with the actual "inner" world lifetime.
///
Expand Down Expand Up @@ -2518,35 +2655,35 @@ mod tests {

let mut query_state = world.query::<Entity>();

// It's best to test get_many_inner directly, as it is shared
// It's best to test get_many_mut_inner directly, as it is shared
// We don't care about aliased mutability for the read-only equivalent

// SAFETY: Query does not access world data.
assert!(query_state
.query_mut(&mut world)
.get_many_inner::<10>(entities.clone().try_into().unwrap())
.get_many_mut_inner::<10>(entities.clone().try_into().unwrap())
.is_ok());

assert_eq!(
query_state
.query_mut(&mut world)
.get_many_inner([entities[0], entities[0]])
.get_many_mut_inner([entities[0], entities[0]])
.unwrap_err(),
QueryEntityError::AliasedMutability(entities[0])
);

assert_eq!(
query_state
.query_mut(&mut world)
.get_many_inner([entities[0], entities[1], entities[0]])
.get_many_mut_inner([entities[0], entities[1], entities[0]])
.unwrap_err(),
QueryEntityError::AliasedMutability(entities[0])
);

assert_eq!(
query_state
.query_mut(&mut world)
.get_many_inner([entities[9], entities[9]])
.get_many_mut_inner([entities[9], entities[9]])
.unwrap_err(),
QueryEntityError::AliasedMutability(entities[9])
);
Expand Down