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

Add get_multiple and get_multiple_mut APIs for Query and QueryState (#4298)

Browse Source 
# Objective

- The inability to have multiple active mutable borrows into a query is a common source of borrow-checker pain for users.
- This is a pointless restriction if and only if we can guarantee that the entities they are accessing are unique.
- This could already by bypassed with get_unchecked, but that is an extremely unsafe API.
- Closes https://github.com/bevyengine/bevy/issues/2042.

## Solution

- Add `get_multiple`, `get_multiple_mut` and their unchecked equivalents (`multiple` and `multiple_mut`) to `Query` and `QueryState`.
- Improve the `QueryEntityError` type to provide more useful error information.

## Changelog

- Added `get_multiple`, `get_multiple_mut` and their unchecked equivalents (`multiple` and `multiple_mut`) to Query and QueryState.

## Migration Guide

- The `QueryEntityError` enum now has a `AliasedMutability variant, and returns the offending entity.

## Context

This is a fresh attempt at #3333; rebasing was behaving very badly and it was important to rebase on top of the recent query soundness fixes. Many thanks to all the reviewers in that thread, especially @BoxyUwU for the help with lifetimes.

## To-do

- [x] Add compile fail tests
- [x] Successfully deduplicate code
- [x] Decide what to do about failing doc tests
- [x] Get some reviews for lifetime soundness
This commit is contained in:
Alice Cecile 2022-03-30 19:16:48 +00:00
parent 63fee2572b
commit 509548190b
6 changed files with 489 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

311
crates/bevy_ecs/src/query/state.rs
View File

@ -153,6 +153,57 @@ where
}
}
/// Returns the read-only query results for the given array of [`Entity`].
///
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is
/// returned instead.
///
/// Note that the unlike [`QueryState::get_multiple_mut`], the entities passed in do not need to be unique.
///
/// # Examples
///
/// ```rust
/// use bevy_ecs::prelude::*;
/// use bevy_ecs::query::QueryEntityError;
///
/// #[derive(Component, PartialEq, Debug)]
/// struct A(usize);
///
/// let mut world = World::new();
/// let entity_vec: Vec<Entity> = (0..3).map(|i|world.spawn().insert(A(i)).id()).collect();
/// let entities: [Entity; 3] = entity_vec.try_into().unwrap();
///
/// world.spawn().insert(A(73));
///
/// let mut query_state = world.query::<&A>();
///
/// let component_values = query_state.get_multiple(&world, entities).unwrap();
///
/// assert_eq!(component_values, [&A(0), &A(1), &A(2)]);
///
/// let wrong_entity = Entity::from_raw(365);
///
/// assert_eq!(query_state.get_multiple(&world, [wrong_entity]), Err(QueryEntityError::NoSuchEntity(wrong_entity)));
/// ```
#[inline]
pub fn get_multiple<'w, 's, const N: usize>(
&'s mut self,
world: &'w World,
entities: [Entity; N],
) -> Result<[<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item; N], QueryEntityError> {
self.update_archetypes(world);
// SAFE: update_archetypes validates the `World` matches
unsafe {
self.get_multiple_read_only_manual(
world,
entities,
world.last_change_tick(),
world.read_change_tick(),
)
}
}
/// Gets the query result for the given [`World`] and [`Entity`].
#[inline]
pub fn get_mut<'w, 's>(
@ -172,6 +223,64 @@ where
}
}
/// Returns the query results for the given array of [`Entity`].
///
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is
/// returned instead.
///
/// ```rust
/// use bevy_ecs::prelude::*;
/// use bevy_ecs::query::QueryEntityError;
///
/// #[derive(Component, PartialEq, Debug)]
/// struct A(usize);
///
/// let mut world = World::new();
///
/// let entities: Vec<Entity> = (0..3).map(|i|world.spawn().insert(A(i)).id()).collect();
/// let entities: [Entity; 3] = entities.try_into().unwrap();
///
/// world.spawn().insert(A(73));
///
/// let mut query_state = world.query::<&mut A>();
///
/// let mut mutable_component_values = query_state.get_multiple_mut(&mut world, entities).unwrap();
///
/// for mut a in mutable_component_values.iter_mut(){
/// a.0 += 5;
/// }
///
/// let component_values = query_state.get_multiple(&world, entities).unwrap();
///
/// assert_eq!(component_values, [&A(5), &A(6), &A(7)]);
///
/// let wrong_entity = Entity::from_raw(57);
/// let invalid_entity = world.spawn().id();
///
/// assert_eq!(query_state.get_multiple_mut(&mut world, [wrong_entity]).unwrap_err(), QueryEntityError::NoSuchEntity(wrong_entity));
/// assert_eq!(query_state.get_multiple_mut(&mut world, [invalid_entity]).unwrap_err(), QueryEntityError::QueryDoesNotMatch(invalid_entity));
/// assert_eq!(query_state.get_multiple_mut(&mut world, [entities[0], entities[0]]).unwrap_err(), QueryEntityError::AliasedMutability(entities[0]));
/// ```
#[inline]
pub fn get_multiple_mut<'w, 's, const N: usize>(
&'s mut self,
world: &'w mut World,
entities: [Entity; N],
) -> Result<[<Q::Fetch as Fetch<'w, 's>>::Item; N], QueryEntityError> {
self.update_archetypes(world);
// SAFE: method requires exclusive world access
// and world has been validated via update_archetypes
unsafe {
self.get_multiple_unchecked_manual(
world,
entities,
world.last_change_tick(),
world.read_change_tick(),
)
}
}
#[inline]
pub fn get_manual<'w, 's>(
&'s self,
@ -218,6 +327,9 @@ where
///
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
/// have unique access to the components they query.
///
/// This must be called on the same `World` that the `Query` was generated from:
/// use `QueryState::validate_world` to verify this.
pub(crate) unsafe fn get_unchecked_manual<'w, 's, QF: Fetch<'w, 's, State = Q::State>>(
&'s self,
world: &'w World,
@ -228,12 +340,12 @@ where
let location = world
.entities
.get(entity)
.ok_or(QueryEntityError::NoSuchEntity)?;
.ok_or(QueryEntityError::NoSuchEntity(entity))?;
if !self
.matched_archetypes
.contains(location.archetype_id.index())
{
return Err(QueryEntityError::QueryDoesNotMatch);
return Err(QueryEntityError::QueryDoesNotMatch(entity));
}
let archetype = &world.archetypes[location.archetype_id];
let mut fetch = QF::init(world, &self.fetch_state, last_change_tick, change_tick);
@ -245,10 +357,90 @@ where
if filter.archetype_filter_fetch(location.index) {
Ok(fetch.archetype_fetch(location.index))
} else {
Err(QueryEntityError::QueryDoesNotMatch)
Err(QueryEntityError::QueryDoesNotMatch(entity))
}
}
/// Gets the read-only query results for the given [`World`] and array of [`Entity`], where the last change and
/// the current change tick are given.
///
/// # Safety
///
/// This must be called on the same `World` that the `Query` was generated from:
/// use `QueryState::validate_world` to verify this.
pub(crate) unsafe fn get_multiple_read_only_manual<'s, 'w, const N: usize>(
&'s self,
world: &'w World,
entities: [Entity; N],
last_change_tick: u32,
change_tick: u32,
) -> Result<[<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item; N], QueryEntityError> {
// SAFE: fetch is read-only
// and world must be validated
let array_of_results = entities.map(|entity| {
self.get_unchecked_manual::<Q::ReadOnlyFetch>(
world,
entity,
last_change_tick,
change_tick,
)
});
// TODO: Replace with TryMap once https://github.com/rust-lang/rust/issues/79711 is stabilized
// If any of the get calls failed, bubble up the error
for result in &array_of_results {
match result {
Ok(_) => (),
Err(error) => return Err(*error),
}
}
// Since we have verified that all entities are present, we can safely unwrap
Ok(array_of_results.map(|result| result.unwrap()))
}
/// Gets the query results for the given [`World`] and array of [`Entity`], where the last change and
/// the current change tick are given.
///
/// # Safety
///
/// This does not check for unique access to subsets of the entity-component data.
/// To be safe, make sure mutable queries have unique access to the components they query.
///
/// This must be called on the same `World` that the `Query` was generated from:
/// use `QueryState::validate_world` to verify this.
pub(crate) unsafe fn get_multiple_unchecked_manual<'s, 'w, const N: usize>(
&'s self,
world: &'w World,
entities: [Entity; N],
last_change_tick: u32,
change_tick: u32,
) -> Result<[<Q::Fetch as Fetch<'w, 's>>::Item; N], QueryEntityError> {
// Verify that all entities are unique
for i in 0..N {
for j in 0..i {
if entities[i] == entities[j] {
return Err(QueryEntityError::AliasedMutability(entities[i]));
}
}
}
let array_of_results = entities.map(|entity| {
self.get_unchecked_manual::<Q::Fetch>(world, entity, last_change_tick, change_tick)
});
// If any of the get calls failed, bubble up the error
for result in &array_of_results {
match result {
Ok(_) => (),
Err(error) => return Err(*error),
}
}
// Since we have verified that all entities are present, we can safely unwrap
Ok(array_of_results.map(|result| result.unwrap()))
}
/// Returns an [`Iterator`] over the query results for the given [`World`].
///
/// This can only be called for read-only queries, see [`Self::iter_mut`] for write-queries.
@ -733,10 +925,117 @@ where
}
/// An error that occurs when retrieving a specific [`Entity`]'s query result.
#[derive(Error, Debug)]
// TODO: return the type_name as part of this error
#[derive(Error, Debug, PartialEq, Clone, Copy)]
pub enum QueryEntityError {
#[error("The given entity does not have the requested component.")]
QueryDoesNotMatch,
QueryDoesNotMatch(Entity),
#[error("The requested entity does not exist.")]
NoSuchEntity,
NoSuchEntity(Entity),
#[error("The entity was requested mutably more than once.")]
AliasedMutability(Entity),
}
#[cfg(test)]
mod tests {
use crate::{prelude::*, query::QueryEntityError};
#[test]
fn get_multiple_unchecked_manual_uniqueness() {
let mut world = World::new();
let entities: Vec<Entity> = (0..10).map(|_| world.spawn().id()).collect();
let query_state = world.query::<Entity>();
// These don't matter for the test
let last_change_tick = world.last_change_tick();
let change_tick = world.read_change_tick();
// It's best to test get_multiple_unchecked_manual directly,
// as it is shared and unsafe
// We don't care about aliased mutabilty for the read-only equivalent
assert!(unsafe {
query_state
.get_multiple_unchecked_manual::<10>(
&world,
entities.clone().try_into().unwrap(),
last_change_tick,
change_tick,
)
.is_ok()
});
assert_eq!(
unsafe {
query_state
.get_multiple_unchecked_manual(
&world,
[entities[0], entities[0]],
last_change_tick,
change_tick,
)
.unwrap_err()
},
QueryEntityError::AliasedMutability(entities[0])
);
assert_eq!(
unsafe {
query_state
.get_multiple_unchecked_manual(
&world,
[entities[0], entities[1], entities[0]],
last_change_tick,
change_tick,
)
.unwrap_err()
},
QueryEntityError::AliasedMutability(entities[0])
);
assert_eq!(
unsafe {
query_state
.get_multiple_unchecked_manual(
&world,
[entities[9], entities[9]],
last_change_tick,
change_tick,
)
.unwrap_err()
},
QueryEntityError::AliasedMutability(entities[9])
);
}
#[test]
#[should_panic]
fn right_world_get() {
let mut world_1 = World::new();
let world_2 = World::new();
let mut query_state = world_1.query::<Entity>();
let _panics = query_state.get(&world_2, Entity::from_raw(0));
}
#[test]
#[should_panic]
fn right_world_get_multiple() {
let mut world_1 = World::new();
let world_2 = World::new();
let mut query_state = world_1.query::<Entity>();
let _panics = query_state.get_multiple(&world_2, []);
}
#[test]
#[should_panic]
fn right_world_get_multiple_mut() {
let mut world_1 = World::new();
let mut world_2 = World::new();
let mut query_state = world_1.query::<Entity>();
let _panics = query_state.get_multiple_mut(&mut world_2, []);
}
}

138
crates/bevy_ecs/src/system/query.rs
View File

@ -618,6 +618,73 @@ where
}
}
/// Returns the read-only query results for the given array of [`Entity`].
///
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is
/// returned instead.
///
/// Note that the unlike [`Query::get_multiple_mut`], the entities passed in do not need to be unique.
///
/// See [`Query::multiple`] for the infallible equivalent.
#[inline]
pub fn get_multiple<const N: usize>(
&self,
entities: [Entity; N],
) -> Result<[<Q::ReadOnlyFetch as Fetch<'_, 's>>::Item; N], QueryEntityError> {
// SAFE: it is the scheduler's responsibility to ensure that `Query` is never handed out on the wrong `World`.
unsafe {
self.state.get_multiple_read_only_manual(
self.world,
entities,
self.last_change_tick,
self.change_tick,
)
}
}
/// Returns the read-only query items for the provided array of [`Entity`]
///
/// See [`Query::get_multiple`] for the [`Result`]-returning equivalent.
///
/// # Examples
/// ```rust, no_run
/// use bevy_ecs::prelude::*;
///
/// #[derive(Component)]
/// struct Targets([Entity; 3]);
///
/// #[derive(Component)]
/// struct Position{
/// x: i8,
/// y: i8
/// };
///
/// impl Position {
/// fn distance(&self, other: &Position) -> i8 {
/// // Manhattan distance is way easier to compute!
/// (self.x - other.x).abs() + (self.y - other.y).abs()
/// }
/// }
///
/// fn check_all_targets_in_range(targeting_query: Query<(Entity, &Targets, &Position)>, targets_query: Query<&Position>){
/// for (targeting_entity, targets, origin) in targeting_query.iter(){
/// // We can use "destructuring" to unpack the results nicely
/// let [target_1, target_2, target_3] = targets_query.multiple(targets.0);
///
/// assert!(target_1.distance(origin) <= 5);
/// assert!(target_2.distance(origin) <= 5);
/// assert!(target_3.distance(origin) <= 5);
/// }
/// }
/// ```
#[inline]
pub fn multiple<const N: usize>(
&self,
entities: [Entity; N],
) -> [<Q::ReadOnlyFetch as Fetch<'_, 's>>::Item; N] {
self.get_multiple(entities).unwrap()
}
/// Returns the query result for the given [`Entity`].
///
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is
@ -659,6 +726,77 @@ where
}
}
/// Returns the query results for the given array of [`Entity`].
///
/// In case of a nonexisting entity, duplicate entities or mismatched component, a [`QueryEntityError`] is
/// returned instead.
///
/// See [`Query::multiple_mut`] for the infallible equivalent.
#[inline]
pub fn get_multiple_mut<const N: usize>(
&mut self,
entities: [Entity; N],
) -> Result<[<Q::Fetch as Fetch<'_, 's>>::Item; N], QueryEntityError> {
// SAFE: scheduler ensures safe Query world access
unsafe {
self.state.get_multiple_unchecked_manual(
self.world,
entities,
self.last_change_tick,
self.change_tick,
)
}
}
/// Returns the query items for the provided array of [`Entity`]
///
/// See [`Query::get_multiple_mut`] for the [`Result`]-returning equivalent.
///
/// # Examples
///
/// ```rust, no_run
/// use bevy_ecs::prelude::*;
///
/// #[derive(Component)]
/// struct Spring{
/// connected_entities: [Entity; 2],
/// strength: f32,
/// }
///
/// #[derive(Component)]
/// struct Position {
/// x: f32,
/// y: f32,
/// }
///
/// #[derive(Component)]
/// struct Force {
/// x: f32,
/// y: f32,
/// }
///
/// fn spring_forces(spring_query: Query<&Spring>, mut mass_query: Query<(&Position, &mut Force)>){
/// for spring in spring_query.iter(){
/// // We can use "destructuring" to unpack our query items nicely
/// let [(position_1, mut force_1), (position_2, mut force_2)] = mass_query.multiple_mut(spring.connected_entities);
///
/// force_1.x += spring.strength * (position_1.x - position_2.x);
/// force_1.y += spring.strength * (position_1.y - position_2.y);
///
/// // Silence borrow-checker: I have split your mutable borrow!
/// force_2.x += spring.strength * (position_2.x - position_1.x);
/// force_2.y += spring.strength * (position_2.y - position_1.y);
/// }
/// }
/// ```
#[inline]
pub fn multiple_mut<const N: usize>(
&mut self,
entities: [Entity; N],
) -> [<Q::Fetch as Fetch<'_, 's>>::Item; N] {
self.get_multiple_mut(entities).unwrap()
}
/// Returns the query result for the given [`Entity`].
///
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is

13
crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_lifetime_safety.rs Normal file
View File

@ -0,0 +1,13 @@
use bevy_ecs::prelude::*;
#[derive(Component)]
struct A(usize);
fn system(mut query: Query<&mut A>, e: Res<Entity>) {
let a1 = query.get_multiple([*e, *e]).unwrap();
let a2 = query.get_mut(*e).unwrap();
// this should fail to compile
println!("{} {}", a1[0].0, a2.0);
}
fn main() {}

10
crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_lifetime_safety.stderr Normal file
View File

@ -0,0 +1,10 @@
error[E0502]: cannot borrow `query` as mutable because it is also borrowed as immutable
--> tests/ui/system_query_get_multiple_lifetime_safety.rs:8:14
|
7 | let a1 = query.get_multiple([*e, *e]).unwrap();
| ---------------------------- immutable borrow occurs here
8 | let a2 = query.get_mut(*e).unwrap();
| ^^^^^^^^^^^^^^^^^ mutable borrow occurs here
9 | // this should fail to compile
10 | println!("{} {}", a1[0].0, a2.0);
| ----- immutable borrow later used here

13
crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_mut_lifetime_safety.rs Normal file
View File

@ -0,0 +1,13 @@
use bevy_ecs::prelude::*;
#[derive(Component)]
struct A(usize);
fn system(mut query: Query<&mut A>, e: Res<Entity>) {
let a1 = query.get_multiple_mut([*e, *e]).unwrap();
let a2 = query.get_mut(*e).unwrap();
// this should fail to compile
println!("{} {}", a1[0].0, a2.0);
}
fn main() {}

10
crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_mut_lifetime_safety.stderr Normal file
View File

@ -0,0 +1,10 @@
error[E0499]: cannot borrow `query` as mutable more than once at a time
--> tests/ui/system_query_get_multiple_mut_lifetime_safety.rs:8:14
|
7 | let a1 = query.get_multiple_mut([*e, *e]).unwrap();
| -------------------------------- first mutable borrow occurs here
8 | let a2 = query.get_mut(*e).unwrap();
| ^^^^^^^^^^^^^^^^^ second mutable borrow occurs here
9 | // this should fail to compile
10 | println!("{} {}", a1[0].0, a2.0);
| ----- first borrow later used here