yeet unsound lifetime annotations on Query methods (#4243)

# Objective
Continuation of #2964 (I really should have checked other methods when I made that PR)

yeet unsound lifetime annotations on `Query` methods.
Example unsoundness:
```rust
use bevy::prelude::*;

fn main() {
    App::new().add_startup_system(bar).add_system(foo).run();
}

pub fn bar(mut cmds: Commands) {
    let e = cmds.spawn().insert(Foo { a: 10 }).id();
    cmds.insert_resource(e);
}

#[derive(Component, Debug, PartialEq, Eq)]
pub struct Foo {
    a: u32,
}
pub fn foo(mut query: Query<&mut Foo>, e: Res<Entity>) {
    dbg!("hi");
    {
        let data: &Foo = query.get(*e).unwrap();
        let data2: Mut<Foo> = query.get_mut(*e).unwrap();
        assert_eq!(data, &*data2); // oops UB
    }

    {
        let data: &Foo = query.single();
        let data2: Mut<Foo> = query.single_mut();
        assert_eq!(data, &*data2); // oops UB
    }

    {
        let data: &Foo = query.get_single().unwrap();
        let data2: Mut<Foo> = query.get_single_mut().unwrap();
        assert_eq!(data, &*data2); // oops UB
    }

    {
        let data: &Foo = query.iter().next().unwrap();
        let data2: Mut<Foo> = query.iter_mut().next().unwrap();
        assert_eq!(data, &*data2); // oops UB
    }

    {
        let mut opt_data: Option<&Foo> = None;
        let mut opt_data_2: Option<Mut<Foo>> = None;
        query.for_each(|data| opt_data = Some(data));
        query.for_each_mut(|data| opt_data_2 = Some(data));
        assert_eq!(opt_data.unwrap(), &*opt_data_2.unwrap()); // oops UB
    }
    dbg!("bye");
}

```

## Solution
yeet unsound lifetime annotations on `Query` methods

Co-authored-by: Carter Anderson <mcanders1@gmail.com>
This commit is contained in:
Boxy 2022-03-22 02:49:41 +00:00
parent fbe7a49d5b
commit 024d98457c
7 changed files with 324 additions and 24 deletions

View file

@ -804,13 +804,13 @@ mod tests {
}
fn hold_component<'w>(&mut self, world: &'w World, entity: Entity) -> Holder<'w> {
let q = self.state_q.get(world);
let a = q.get(entity).unwrap();
let a = q.get_inner(entity).unwrap();
Holder { value: a }
}
fn hold_components<'w>(&mut self, world: &'w World) -> Vec<Holder<'w>> {
let mut components = Vec::new();
let q = self.state_q.get(world);
for a in q.iter() {
for a in q.iter_inner() {
components.push(Holder { value: a });
}
components

View file

@ -3,7 +3,7 @@ use crate::{
entity::Entity,
query::{
Fetch, FilterFetch, NopFetch, QueryCombinationIter, QueryEntityError, QueryIter,
QueryState, WorldQuery,
QueryState, ReadOnlyFetch, WorldQuery,
},
world::{Mut, World},
};
@ -299,7 +299,7 @@ where
/// # bevy_ecs::system::assert_is_system(report_names_system);
/// ```
#[inline]
pub fn iter(&'s self) -> QueryIter<'w, 's, Q, Q::ReadOnlyFetch, F> {
pub fn iter(&self) -> QueryIter<'_, 's, Q, Q::ReadOnlyFetch, F> {
// SAFE: system runs without conflicts with other systems.
// same-system queries have runtime borrow checks when they conflict
unsafe {
@ -454,17 +454,19 @@ where
/// # bevy_ecs::system::assert_is_system(report_names_system);
/// ```
#[inline]
pub fn for_each<FN: FnMut(<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item)>(&'s self, f: FN) {
pub fn for_each<'this>(
&'this self,
f: impl FnMut(<Q::ReadOnlyFetch as Fetch<'this, 's>>::Item),
) {
// SAFE: system runs without conflicts with other systems.
// same-system queries have runtime borrow checks when they conflict
unsafe {
self.state
.for_each_unchecked_manual::<Q::ReadOnlyFetch, FN>(
self.world,
f,
self.last_change_tick,
self.change_tick,
);
self.state.for_each_unchecked_manual::<Q::ReadOnlyFetch, _>(
self.world,
f,
self.last_change_tick,
self.change_tick,
);
};
}
@ -524,17 +526,17 @@ where
///* `batch_size` - The number of batches to spawn
///* `f` - The function to run on each item in the query
#[inline]
pub fn par_for_each<FN: Fn(<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item) + Send + Sync + Clone>(
&'s self,
pub fn par_for_each<'this>(
&'this self,
task_pool: &TaskPool,
batch_size: usize,
f: FN,
f: impl Fn(<Q::ReadOnlyFetch as Fetch<'this, 's>>::Item) + Send + Sync + Clone,
) {
// SAFE: system runs without conflicts with other systems. same-system queries have runtime
// borrow checks when they conflict
unsafe {
self.state
.par_for_each_unchecked_manual::<Q::ReadOnlyFetch, FN>(
.par_for_each_unchecked_manual::<Q::ReadOnlyFetch, _>(
self.world,
task_pool,
batch_size,
@ -601,9 +603,9 @@ where
/// ```
#[inline]
pub fn get(
&'s self,
&self,
entity: Entity,
) -> Result<<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item, QueryEntityError> {
) -> Result<<Q::ReadOnlyFetch as Fetch<'_, 's>>::Item, QueryEntityError> {
// SAFE: system runs without conflicts with other systems.
// same-system queries have runtime borrow checks when they conflict
unsafe {
@ -834,7 +836,7 @@ where
/// Panics if the number of query results is not exactly one. Use
/// [`get_single`](Self::get_single) to return a `Result` instead of panicking.
#[track_caller]
pub fn single(&'s self) -> <Q::ReadOnlyFetch as Fetch<'w, 's>>::Item {
pub fn single(&self) -> <Q::ReadOnlyFetch as Fetch<'_, 's>>::Item {
self.get_single().unwrap()
}
@ -870,8 +872,8 @@ where
/// # bevy_ecs::system::assert_is_system(player_scoring_system);
/// ```
pub fn get_single(
&'s self,
) -> Result<<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item, QuerySingleError> {
&self,
) -> Result<<Q::ReadOnlyFetch as Fetch<'_, 's>>::Item, QuerySingleError> {
let mut query = self.iter();
let first = query.next();
let extra = query.next().is_some();
@ -1038,3 +1040,90 @@ pub enum QuerySingleError {
#[error("Multiple entities fit the query {0}!")]
MultipleEntities(&'static str),
}
impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F>
where
F::Fetch: FilterFetch,
Q::Fetch: ReadOnlyFetch,
{
/// Returns the query result for the given [`Entity`], with the actual "inner" world lifetime.
///
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is
/// returned instead.
///
/// This can only return immutable data (mutable data will be cast to an immutable form).
/// See [`get_mut`](Self::get_mut) for queries that contain at least one mutable component.
///
/// # Example
///
/// Here, `get` is used to retrieve the exact query result of the entity specified by the
/// `SelectedCharacter` resource.
///
/// ```
/// # use bevy_ecs::prelude::*;
/// #
/// # struct SelectedCharacter { entity: Entity }
/// # #[derive(Component)]
/// # struct Character { name: String }
/// #
/// fn print_selected_character_name_system(
/// query: Query<&Character>,
/// selection: Res<SelectedCharacter>
/// )
/// {
/// if let Ok(selected_character) = query.get(selection.entity) {
/// println!("{}", selected_character.name);
/// }
/// }
/// # bevy_ecs::system::assert_is_system(print_selected_character_name_system);
/// ```
#[inline]
pub fn get_inner(
&'s self,
entity: Entity,
) -> Result<<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item, QueryEntityError> {
// SAFE: system runs without conflicts with other systems.
// same-system queries have runtime borrow checks when they conflict
unsafe {
self.state.get_unchecked_manual::<Q::ReadOnlyFetch>(
self.world,
entity,
self.last_change_tick,
self.change_tick,
)
}
}
/// Returns an [`Iterator`] over the query results, with the actual "inner" world lifetime.
///
/// This can only return immutable data (mutable data will be cast to an immutable form).
/// See [`Self::iter_mut`] for queries that contain at least one mutable component.
///
/// # Example
///
/// Here, the `report_names_system` iterates over the `Player` component of every entity
/// that contains it:
///
/// ```
/// # use bevy_ecs::prelude::*;
/// #
/// # #[derive(Component)]
/// # struct Player { name: String }
/// #
/// fn report_names_system(query: Query<&Player>) {
/// for player in query.iter() {
/// println!("Say hello to {}!", player.name);
/// }
/// }
/// # bevy_ecs::system::assert_is_system(report_names_system);
/// ```
#[inline]
pub fn iter_inner(&'s self) -> QueryIter<'w, 's, Q, Q::ReadOnlyFetch, F> {
// SAFE: system runs without conflicts with other systems.
// same-system queries have runtime borrow checks when they conflict
unsafe {
self.state
.iter_unchecked_manual(self.world, self.last_change_tick, self.change_tick)
}
}
}

View file

@ -0,0 +1,92 @@
use bevy_ecs::prelude::*;
use bevy_ecs::system::SystemState;
#[derive(Component, Eq, PartialEq, Debug)]
struct Foo(u32);
fn main() {
let mut world = World::default();
let e = world.spawn().insert(Foo(10_u32)).id();
let mut system_state = SystemState::<Query<&mut Foo>>::new(&mut world);
{
let mut query = system_state.get_mut(&mut world);
dbg!("hi");
{
let data: &Foo = query.get(e).unwrap();
let mut data2: Mut<Foo> = query.get_mut(e).unwrap();
assert_eq!(data, &mut *data2); // oops UB
}
{
let mut data2: Mut<Foo> = query.get_mut(e).unwrap();
let data: &Foo = query.get(e).unwrap();
assert_eq!(data, &mut *data2); // oops UB
}
{
let data: &Foo = query.get_component::<Foo>(e).unwrap();
let mut data2: Mut<Foo> = query.get_component_mut(e).unwrap();
assert_eq!(data, &mut *data2); // oops UB
}
{
let mut data2: Mut<Foo> = query.get_component_mut(e).unwrap();
let data: &Foo = query.get_component::<Foo>(e).unwrap();
assert_eq!(data, &mut *data2); // oops UB
}
{
let data: &Foo = query.single();
let mut data2: Mut<Foo> = query.single_mut();
assert_eq!(data, &mut *data2); // oops UB
}
{
let mut data2: Mut<Foo> = query.single_mut();
let data: &Foo = query.single();
assert_eq!(data, &mut *data2); // oops UB
}
{
let data: &Foo = query.get_single().unwrap();
let mut data2: Mut<Foo> = query.get_single_mut().unwrap();
assert_eq!(data, &mut *data2); // oops UB
}
{
let mut data2: Mut<Foo> = query.get_single_mut().unwrap();
let data: &Foo = query.get_single().unwrap();
assert_eq!(data, &mut *data2); // oops UB
}
{
let data: &Foo = query.iter().next().unwrap();
let mut data2: Mut<Foo> = query.iter_mut().next().unwrap();
assert_eq!(data, &mut *data2); // oops UB
}
{
let mut data2: Mut<Foo> = query.iter_mut().next().unwrap();
let data: &Foo = query.iter().next().unwrap();
assert_eq!(data, &mut *data2); // oops UB
}
{
let mut opt_data: Option<&Foo> = None;
let mut opt_data_2: Option<Mut<Foo>> = None;
query.for_each(|data| opt_data = Some(data));
query.for_each_mut(|data| opt_data_2 = Some(data));
assert_eq!(opt_data.unwrap(), &mut *opt_data_2.unwrap()); // oops UB
}
{
let mut opt_data_2: Option<Mut<Foo>> = None;
let mut opt_data: Option<&Foo> = None;
query.for_each_mut(|data| opt_data_2 = Some(data));
query.for_each(|data| opt_data = Some(data));
assert_eq!(opt_data.unwrap(), &mut *opt_data_2.unwrap()); // oops UB
}
dbg!("bye");
}
}

View file

@ -0,0 +1,119 @@
error[E0502]: cannot borrow `query` as mutable because it is also borrowed as immutable
--> tests/ui/query_lifetime_safety.rs:17:39
|
16 | let data: &Foo = query.get(e).unwrap();
| ------------ immutable borrow occurs here
17 | let mut data2: Mut<Foo> = query.get_mut(e).unwrap();
| ^^^^^^^^^^^^^^^^ mutable borrow occurs here
18 | assert_eq!(data, &mut *data2); // oops UB
| ----------------------------- immutable borrow later used here
error[E0502]: cannot borrow `query` as immutable because it is also borrowed as mutable
--> tests/ui/query_lifetime_safety.rs:23:30
|
22 | let mut data2: Mut<Foo> = query.get_mut(e).unwrap();
| ---------------- mutable borrow occurs here
23 | let data: &Foo = query.get(e).unwrap();
| ^^^^^^^^^^^^ immutable borrow occurs here
24 | assert_eq!(data, &mut *data2); // oops UB
| ----- mutable borrow later used here
error[E0502]: cannot borrow `query` as mutable because it is also borrowed as immutable
--> tests/ui/query_lifetime_safety.rs:29:39
|
28 | let data: &Foo = query.get_component::<Foo>(e).unwrap();
| ----------------------------- immutable borrow occurs here
29 | let mut data2: Mut<Foo> = query.get_component_mut(e).unwrap();
| ^^^^^^^^^^^^^^^^^^^^^^^^^^ mutable borrow occurs here
30 | assert_eq!(data, &mut *data2); // oops UB
| ----------------------------- immutable borrow later used here
error[E0502]: cannot borrow `query` as immutable because it is also borrowed as mutable
--> tests/ui/query_lifetime_safety.rs:35:30
|
34 | let mut data2: Mut<Foo> = query.get_component_mut(e).unwrap();
| -------------------------- mutable borrow occurs here
35 | let data: &Foo = query.get_component::<Foo>(e).unwrap();
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ immutable borrow occurs here
36 | assert_eq!(data, &mut *data2); // oops UB
| ----- mutable borrow later used here
error[E0502]: cannot borrow `query` as mutable because it is also borrowed as immutable
--> tests/ui/query_lifetime_safety.rs:41:39
|
40 | let data: &Foo = query.single();
| -------------- immutable borrow occurs here
41 | let mut data2: Mut<Foo> = query.single_mut();
| ^^^^^^^^^^^^^^^^^^ mutable borrow occurs here
42 | assert_eq!(data, &mut *data2); // oops UB
| ----------------------------- immutable borrow later used here
error[E0502]: cannot borrow `query` as immutable because it is also borrowed as mutable
--> tests/ui/query_lifetime_safety.rs:47:30
|
46 | let mut data2: Mut<Foo> = query.single_mut();
| ------------------ mutable borrow occurs here
47 | let data: &Foo = query.single();
| ^^^^^^^^^^^^^^ immutable borrow occurs here
48 | assert_eq!(data, &mut *data2); // oops UB
| ----- mutable borrow later used here
error[E0502]: cannot borrow `query` as mutable because it is also borrowed as immutable
--> tests/ui/query_lifetime_safety.rs:53:39
|
52 | let data: &Foo = query.get_single().unwrap();
| ------------------ immutable borrow occurs here
53 | let mut data2: Mut<Foo> = query.get_single_mut().unwrap();
| ^^^^^^^^^^^^^^^^^^^^^^ mutable borrow occurs here
54 | assert_eq!(data, &mut *data2); // oops UB
| ----------------------------- immutable borrow later used here
error[E0502]: cannot borrow `query` as immutable because it is also borrowed as mutable
--> tests/ui/query_lifetime_safety.rs:59:30
|
58 | let mut data2: Mut<Foo> = query.get_single_mut().unwrap();
| ---------------------- mutable borrow occurs here
59 | let data: &Foo = query.get_single().unwrap();
| ^^^^^^^^^^^^^^^^^^ immutable borrow occurs here
60 | assert_eq!(data, &mut *data2); // oops UB
| ----- mutable borrow later used here
error[E0502]: cannot borrow `query` as mutable because it is also borrowed as immutable
--> tests/ui/query_lifetime_safety.rs:65:39
|
64 | let data: &Foo = query.iter().next().unwrap();
| ------------ immutable borrow occurs here
65 | let mut data2: Mut<Foo> = query.iter_mut().next().unwrap();
| ^^^^^^^^^^^^^^^^ mutable borrow occurs here
66 | assert_eq!(data, &mut *data2); // oops UB
| ----------------------------- immutable borrow later used here
error[E0502]: cannot borrow `query` as immutable because it is also borrowed as mutable
--> tests/ui/query_lifetime_safety.rs:71:30
|
70 | let mut data2: Mut<Foo> = query.iter_mut().next().unwrap();
| ---------------- mutable borrow occurs here
71 | let data: &Foo = query.iter().next().unwrap();
| ^^^^^^^^^^^^ immutable borrow occurs here
72 | assert_eq!(data, &mut *data2); // oops UB
| ----- mutable borrow later used here
error[E0502]: cannot borrow `query` as mutable because it is also borrowed as immutable
--> tests/ui/query_lifetime_safety.rs:79:13
|
78 | query.for_each(|data| opt_data = Some(data));
| -------------------------------------------- immutable borrow occurs here
79 | query.for_each_mut(|data| opt_data_2 = Some(data));
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ mutable borrow occurs here
80 | assert_eq!(opt_data.unwrap(), &mut *opt_data_2.unwrap()); // oops UB
| -------- immutable borrow later used here
error[E0502]: cannot borrow `query` as immutable because it is also borrowed as mutable
--> tests/ui/query_lifetime_safety.rs:87:13
|
86 | query.for_each_mut(|data| opt_data_2 = Some(data));
| -------------------------------------------------- mutable borrow occurs here
87 | query.for_each(|data| opt_data = Some(data));
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ immutable borrow occurs here
88 | assert_eq!(opt_data.unwrap(), &mut *opt_data_2.unwrap()); // oops UB
| ---------- mutable borrow later used here

View file

@ -625,7 +625,7 @@ impl<const I: usize> EntityRenderCommand for SetMeshViewBindGroup<I> {
view_query: SystemParamItem<'w, '_, Self::Param>,
pass: &mut TrackedRenderPass<'w>,
) -> RenderCommandResult {
let (view_uniform, view_lights, mesh_view_bind_group) = view_query.get(view).unwrap();
let (view_uniform, view_lights, mesh_view_bind_group) = view_query.get_inner(view).unwrap();
pass.set_bind_group(
I,
&mesh_view_bind_group.value,

View file

@ -401,7 +401,7 @@ impl<const I: usize> EntityRenderCommand for SetMesh2dViewBindGroup<I> {
view_query: SystemParamItem<'w, '_, Self::Param>,
pass: &mut TrackedRenderPass<'w>,
) -> RenderCommandResult {
let (view_uniform, mesh2d_view_bind_group) = view_query.get(view).unwrap();
let (view_uniform, mesh2d_view_bind_group) = view_query.get_inner(view).unwrap();
pass.set_bind_group(I, &mesh2d_view_bind_group.value, &[view_uniform.offset]);
RenderCommandResult::Success

View file

@ -240,7 +240,7 @@ impl EntityRenderCommand for DrawMeshInstanced {
pass: &mut TrackedRenderPass<'w>,
) -> RenderCommandResult {
let mesh_handle = mesh_query.get(item).unwrap();
let instance_buffer = instance_buffer_query.get(item).unwrap();
let instance_buffer = instance_buffer_query.get_inner(item).unwrap();
let gpu_mesh = match meshes.into_inner().get(mesh_handle) {
Some(gpu_mesh) => gpu_mesh,