Skip to content

Commit b7d5254

Browse files
VictoronzVictoronz
authored
implement get_many_unique (#18315)
# Objective Continuation to #16547 and #17954. The `get_many` family are the last methods on `Query`/`QueryState` for which we're still missing a `unique` version. ## Solution Offer `get_many_unique`/`get_many_unique_mut` and `get_many_unique_inner`! Their implementation is the same as `get_many`, the difference lies in their guaranteed-to-be unique inputs, meaning we never do any aliasing checks. To reduce confusion, we also rename `get_many_readonly` into `get_many_inner` and the current `get_many_inner` into `get_many_mut_inner` to clarify their purposes. ## Testing Doc examples. ## Migration Guide `get_many_inner` is now called `get_many_mut_inner`. `get_many_readonly` is now called `get_many_inner`.
1 parent 9bcaad7 commit b7d5254

File tree

2 files changed

+238
-18
lines changed

2 files changed

+238
-18
lines changed

‎crates/bevy_ecs/src/query/state.rs

Lines changed: 85 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use crate::{
22
archetype::{Archetype, ArchetypeComponentId, ArchetypeGeneration, ArchetypeId},
33
component::{ComponentId, Tick},
4-
entity::{Entity, EntityBorrow, EntitySet},
4+
entity::{unique_array::UniqueEntityArray, Entity, EntityBorrow, EntitySet},
55
entity_disabling::DefaultQueryFilters,
66
prelude::FromWorld,
77
query::{Access, FilteredAccess, QueryCombinationIter, QueryIter, QueryParIter, WorldQuery},
@@ -997,6 +997,44 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
997997
self.query(world).get_many_inner(entities)
998998
}
999999

1000+
/// Returns the read-only query results for the given [`UniqueEntityArray`].
1001+
///
1002+
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is
1003+
/// returned instead.
1004+
///
1005+
/// # Examples
1006+
///
1007+
/// ```
1008+
/// use bevy_ecs::{prelude::*, query::QueryEntityError, entity::{EntitySetIterator, unique_array::UniqueEntityArray, unique_vec::UniqueEntityVec}};
1009+
///
1010+
/// #[derive(Component, PartialEq, Debug)]
1011+
/// struct A(usize);
1012+
///
1013+
/// let mut world = World::new();
1014+
/// let entity_set: UniqueEntityVec = world.spawn_batch((0..3).map(A)).collect_set();
1015+
/// let entity_set: UniqueEntityArray<3> = entity_set.try_into().unwrap();
1016+
///
1017+
/// world.spawn(A(73));
1018+
///
1019+
/// let mut query_state = world.query::<&A>();
1020+
///
1021+
/// let component_values = query_state.get_many_unique(&world, entity_set).unwrap();
1022+
///
1023+
/// assert_eq!(component_values, [&A(0), &A(1), &A(2)]);
1024+
///
1025+
/// let wrong_entity = Entity::from_raw(365);
1026+
///
1027+
/// assert_eq!(match query_state.get_many_unique(&mut world, UniqueEntityArray::from([wrong_entity])).unwrap_err() {QueryEntityError::EntityDoesNotExist(error) => error.entity, _ => panic!()}, wrong_entity);
1028+
/// ```
1029+
#[inline]
1030+
pub fn get_many_unique<'w, const N: usize>(
1031+
&mut self,
1032+
world: &'w World,
1033+
entities: UniqueEntityArray<N>,
1034+
) -> Result<[ROQueryItem<'w, D>; N], QueryEntityError> {
1035+
self.query(world).get_many_unique_inner(entities)
1036+
}
1037+
10001038
/// Gets the query result for the given [`World`] and [`Entity`].
10011039
///
10021040
/// This is always guaranteed to run in `O(1)` time.
@@ -1053,7 +1091,52 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
10531091
world: &'w mut World,
10541092
entities: [Entity; N],
10551093
) -> Result<[D::Item<'w>; N], QueryEntityError> {
1056-
self.query_mut(world).get_many_inner(entities)
1094+
self.query_mut(world).get_many_mut_inner(entities)
1095+
}
1096+
1097+
/// Returns the query results for the given [`UniqueEntityArray`].
1098+
///
1099+
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is
1100+
/// returned instead.
1101+
///
1102+
/// ```
1103+
/// use bevy_ecs::{prelude::*, query::QueryEntityError, entity::{EntitySetIterator, unique_array::UniqueEntityArray, unique_vec::UniqueEntityVec}};
1104+
///
1105+
/// #[derive(Component, PartialEq, Debug)]
1106+
/// struct A(usize);
1107+
///
1108+
/// let mut world = World::new();
1109+
///
1110+
/// let entity_set: UniqueEntityVec = world.spawn_batch((0..3).map(A)).collect_set();
1111+
/// let entity_set: UniqueEntityArray<3> = entity_set.try_into().unwrap();
1112+
///
1113+
/// world.spawn(A(73));
1114+
///
1115+
/// let mut query_state = world.query::<&mut A>();
1116+
///
1117+
/// let mut mutable_component_values = query_state.get_many_unique_mut(&mut world, entity_set).unwrap();
1118+
///
1119+
/// for mut a in &mut mutable_component_values {
1120+
/// a.0 += 5;
1121+
/// }
1122+
///
1123+
/// let component_values = query_state.get_many_unique(&world, entity_set).unwrap();
1124+
///
1125+
/// assert_eq!(component_values, [&A(5), &A(6), &A(7)]);
1126+
///
1127+
/// let wrong_entity = Entity::from_raw(57);
1128+
/// let invalid_entity = world.spawn_empty().id();
1129+
///
1130+
/// assert_eq!(match query_state.get_many_unique(&mut world, UniqueEntityArray::from([wrong_entity])).unwrap_err() {QueryEntityError::EntityDoesNotExist(error) => error.entity, _ => panic!()}, wrong_entity);
1131+
/// assert_eq!(match query_state.get_many_unique_mut(&mut world, UniqueEntityArray::from([invalid_entity])).unwrap_err() {QueryEntityError::QueryDoesNotMatch(entity, _) => entity, _ => panic!()}, invalid_entity);
1132+
/// ```
1133+
#[inline]
1134+
pub fn get_many_unique_mut<'w, const N: usize>(
1135+
&mut self,
1136+
world: &'w mut World,
1137+
entities: UniqueEntityArray<N>,
1138+
) -> Result<[D::Item<'w>; N], QueryEntityError> {
1139+
self.query_mut(world).get_many_unique_inner(entities)
10571140
}
10581141

10591142
/// Gets the query result for the given [`World`] and [`Entity`].

‎crates/bevy_ecs/src/system/query.rs

Lines changed: 153 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
use crate::{
22
batching::BatchingStrategy,
33
component::Tick,
4-
entity::{Entity, EntityBorrow, EntityDoesNotExistError, EntitySet},
4+
entity::{
5+
unique_array::UniqueEntityArray, Entity, EntityBorrow, EntityDoesNotExistError, EntitySet,
6+
},
57
query::{
68
DebugCheckedUnwrap, NopWorldQuery, QueryCombinationIter, QueryData, QueryEntityError,
79
QueryFilter, QueryIter, QueryManyIter, QueryManyUniqueIter, QueryParIter, QueryParManyIter,
@@ -1323,15 +1325,65 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
13231325
/// # See also
13241326
///
13251327
/// - [`get_many_mut`](Self::get_many_mut) to get mutable query items.
1328+
/// - [`get_many_unique`](Self::get_many_unique) to only handle unique inputs.
13261329
/// - [`many`](Self::many) for the panicking version.
13271330
#[inline]
13281331
pub fn get_many<const N: usize>(
13291332
&self,
13301333
entities: [Entity; N],
13311334
) -> Result<[ROQueryItem<'_, D>; N], QueryEntityError> {
1332-
// Note that this calls `get_many_readonly` instead of `get_many_inner`
1333-
// since we don't need to check for duplicates.
1334-
self.as_readonly().get_many_readonly(entities)
1335+
// Note that we call a separate `*_inner` method from `get_many_mut`
1336+
// because we don't need to check for duplicates.
1337+
self.as_readonly().get_many_inner(entities)
1338+
}
1339+
1340+
/// Returns the read-only query items for the given [`UniqueEntityArray`].
1341+
///
1342+
/// The returned query items are in the same order as the input.
1343+
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is returned instead.
1344+
///
1345+
/// # Examples
1346+
///
1347+
/// ```
1348+
/// use bevy_ecs::{prelude::*, query::QueryEntityError, entity::{EntitySetIterator, unique_array::UniqueEntityArray, unique_vec::UniqueEntityVec}};
1349+
///
1350+
/// #[derive(Component, PartialEq, Debug)]
1351+
/// struct A(usize);
1352+
///
1353+
/// let mut world = World::new();
1354+
/// let entity_set: UniqueEntityVec = world.spawn_batch((0..3).map(A)).collect_set();
1355+
/// let entity_set: UniqueEntityArray<3> = entity_set.try_into().unwrap();
1356+
///
1357+
/// world.spawn(A(73));
1358+
///
1359+
/// let mut query_state = world.query::<&A>();
1360+
/// let query = query_state.query(&world);
1361+
///
1362+
/// let component_values = query.get_many_unique(entity_set).unwrap();
1363+
///
1364+
/// assert_eq!(component_values, [&A(0), &A(1), &A(2)]);
1365+
///
1366+
/// let wrong_entity = Entity::from_raw(365);
1367+
///
1368+
/// assert_eq!(
1369+
/// match query.get_many_unique(UniqueEntityArray::from([wrong_entity])).unwrap_err() {
1370+
/// QueryEntityError::EntityDoesNotExist(error) => error.entity,
1371+
/// _ => panic!(),
1372+
/// },
1373+
/// wrong_entity
1374+
/// );
1375+
/// ```
1376+
///
1377+
/// # See also
1378+
///
1379+
/// - [`get_many_unique_mut`](Self::get_many_mut) to get mutable query items.
1380+
/// - [`get_many`](Self::get_many) to handle inputs with duplicates.
1381+
#[inline]
1382+
pub fn get_many_unique<const N: usize>(
1383+
&self,
1384+
entities: UniqueEntityArray<N>,
1385+
) -> Result<[ROQueryItem<'_, D>; N], QueryEntityError> {
1386+
self.as_readonly().get_many_unique_inner(entities)
13351387
}
13361388

13371389
/// Returns the read-only query items for the given array of [`Entity`].
@@ -1560,7 +1612,75 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
15601612
&mut self,
15611613
entities: [Entity; N],
15621614
) -> Result<[D::Item<'_>; N], QueryEntityError> {
1563-
self.reborrow().get_many_inner(entities)
1615+
self.reborrow().get_many_mut_inner(entities)
1616+
}
1617+
1618+
/// Returns the query items for the given [`UniqueEntityArray`].
1619+
///
1620+
/// The returned query items are in the same order as the input.
1621+
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is returned instead.
1622+
///
1623+
/// # Examples
1624+
///
1625+
/// ```
1626+
/// use bevy_ecs::{prelude::*, query::QueryEntityError, entity::{EntitySetIterator, unique_array::UniqueEntityArray, unique_vec::UniqueEntityVec}};
1627+
///
1628+
/// #[derive(Component, PartialEq, Debug)]
1629+
/// struct A(usize);
1630+
///
1631+
/// let mut world = World::new();
1632+
///
1633+
/// let entity_set: UniqueEntityVec<_> = world.spawn_batch((0..3).map(A)).collect_set();
1634+
/// let entity_set: UniqueEntityArray<3> = entity_set.try_into().unwrap();
1635+
///
1636+
/// world.spawn(A(73));
1637+
/// let wrong_entity = Entity::from_raw(57);
1638+
/// let invalid_entity = world.spawn_empty().id();
1639+
///
1640+
///
1641+
/// let mut query_state = world.query::<&mut A>();
1642+
/// let mut query = query_state.query_mut(&mut world);
1643+
///
1644+
/// let mut mutable_component_values = query.get_many_unique_mut(entity_set).unwrap();
1645+
///
1646+
/// for mut a in &mut mutable_component_values {
1647+
/// a.0 += 5;
1648+
/// }
1649+
///
1650+
/// let component_values = query.get_many_unique(entity_set).unwrap();
1651+
///
1652+
/// assert_eq!(component_values, [&A(5), &A(6), &A(7)]);
1653+
///
1654+
/// assert_eq!(
1655+
/// match query
1656+
/// .get_many_unique_mut(UniqueEntityArray::from([wrong_entity]))
1657+
/// .unwrap_err()
1658+
/// {
1659+
/// QueryEntityError::EntityDoesNotExist(error) => error.entity,
1660+
/// _ => panic!(),
1661+
/// },
1662+
/// wrong_entity
1663+
/// );
1664+
/// assert_eq!(
1665+
/// match query
1666+
/// .get_many_unique_mut(UniqueEntityArray::from([invalid_entity]))
1667+
/// .unwrap_err()
1668+
/// {
1669+
/// QueryEntityError::QueryDoesNotMatch(entity, _) => entity,
1670+
/// _ => panic!(),
1671+
/// },
1672+
/// invalid_entity
1673+
/// );
1674+
/// ```
1675+
/// # See also
1676+
///
1677+
/// - [`get_many_unique`](Self::get_many) to get read-only query items.
1678+
#[inline]
1679+
pub fn get_many_unique_mut<const N: usize>(
1680+
&mut self,
1681+
entities: UniqueEntityArray<N>,
1682+
) -> Result<[D::Item<'_>; N], QueryEntityError> {
1683+
self.reborrow().get_many_unique_inner(entities)
15641684
}
15651685

15661686
/// Returns the query items for the given array of [`Entity`].
@@ -1573,10 +1693,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
15731693
///
15741694
/// - [`get_many`](Self::get_many) to get read-only query items without checking for duplicate entities.
15751695
/// - [`get_many_mut`](Self::get_many_mut) to get items using a mutable reference.
1576-
/// - [`get_many_readonly`](Self::get_many_readonly) to get read-only query items without checking for duplicate entities
1577-
/// with the actual "inner" world lifetime.
1696+
/// - [`get_many_inner`](Self::get_many_mut_inner) to get read-only query items with the actual "inner" world lifetime.
15781697
#[inline]
1579-
pub fn get_many_inner<const N: usize>(
1698+
pub fn get_many_mut_inner<const N: usize>(
15801699
self,
15811700
entities: [Entity; N],
15821701
) -> Result<[D::Item<'w>; N], QueryEntityError> {
@@ -1588,7 +1707,6 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
15881707
}
15891708
}
15901709
}
1591-
15921710
// SAFETY: All entities are unique, so the results don't alias.
15931711
unsafe { self.get_many_impl(entities) }
15941712
}
@@ -1603,9 +1721,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
16031721
///
16041722
/// - [`get_many`](Self::get_many) to get read-only query items without checking for duplicate entities.
16051723
/// - [`get_many_mut`](Self::get_many_mut) to get items using a mutable reference.
1606-
/// - [`get_many_inner`](Self::get_many_readonly) to get mutable query items with the actual "inner" world lifetime.
1724+
/// - [`get_many_mut_inner`](Self::get_many_mut_inner) to get mutable query items with the actual "inner" world lifetime.
16071725
#[inline]
1608-
pub fn get_many_readonly<const N: usize>(
1726+
pub fn get_many_inner<const N: usize>(
16091727
self,
16101728
entities: [Entity; N],
16111729
) -> Result<[D::Item<'w>; N], QueryEntityError>
@@ -1616,6 +1734,25 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
16161734
unsafe { self.get_many_impl(entities) }
16171735
}
16181736

1737+
/// Returns the query items for the given [`UniqueEntityArray`].
1738+
/// This consumes the [`Query`] to return results with the actual "inner" world lifetime.
1739+
///
1740+
/// The returned query items are in the same order as the input.
1741+
/// In case of a nonexisting entity, duplicate entities or mismatched component, a [`QueryEntityError`] is returned instead.
1742+
///
1743+
/// # See also
1744+
///
1745+
/// - [`get_many_unique`](Self::get_many_unique) to get read-only query items without checking for duplicate entities.
1746+
/// - [`get_many_unique_mut`](Self::get_many_unique_mut) to get items using a mutable reference.
1747+
#[inline]
1748+
pub fn get_many_unique_inner<const N: usize>(
1749+
self,
1750+
entities: UniqueEntityArray<N>,
1751+
) -> Result<[D::Item<'w>; N], QueryEntityError> {
1752+
// SAFETY: All entities are unique, so the results don't alias.
1753+
unsafe { self.get_many_impl(entities.into_inner()) }
1754+
}
1755+
16191756
/// Returns the query items for the given array of [`Entity`].
16201757
/// This consumes the [`Query`] to return results with the actual "inner" world lifetime.
16211758
///
@@ -2518,35 +2655,35 @@ mod tests {
25182655

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

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

25242661
// SAFETY: Query does not access world data.
25252662
assert!(query_state
25262663
.query_mut(&mut world)
2527-
.get_many_inner::<10>(entities.clone().try_into().unwrap())
2664+
.get_many_mut_inner::<10>(entities.clone().try_into().unwrap())
25282665
.is_ok());
25292666

25302667
assert_eq!(
25312668
query_state
25322669
.query_mut(&mut world)
2533-
.get_many_inner([entities[0], entities[0]])
2670+
.get_many_mut_inner([entities[0], entities[0]])
25342671
.unwrap_err(),
25352672
QueryEntityError::AliasedMutability(entities[0])
25362673
);
25372674

25382675
assert_eq!(
25392676
query_state
25402677
.query_mut(&mut world)
2541-
.get_many_inner([entities[0], entities[1], entities[0]])
2678+
.get_many_mut_inner([entities[0], entities[1], entities[0]])
25422679
.unwrap_err(),
25432680
QueryEntityError::AliasedMutability(entities[0])
25442681
);
25452682

25462683
assert_eq!(
25472684
query_state
25482685
.query_mut(&mut world)
2549-
.get_many_inner([entities[9], entities[9]])
2686+
.get_many_mut_inner([entities[9], entities[9]])
25502687
.unwrap_err(),
25512688
QueryEntityError::AliasedMutability(entities[9])
25522689
);

0 commit comments

Comments
 (0)