Skip to content

Commit

Permalink
yeet unsound lifetime annotations on Query methods (bevyengine#4243)
Browse files Browse the repository at this point in the history 
# Objective
Continuation of bevyengine#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 <[email protected]>
  • Loading branch information
@BoxyUwU @cart @ItsDoot
2 people authored and ItsDoot committed Feb 1, 2023
1 parent 0540eb3 commit 9bc894d
Show file tree
Hide file tree
Showing 7 changed files with 324 additions and 24 deletions.
4 changes: 2 additions & 2 deletions crates/bevy_ecs/src/system/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
127 changes: 108 additions & 19 deletions crates/bevy_ecs/src/system/query.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use crate::{
entity::Entity,
query::{
Fetch, FilterFetch, NopFetch, QueryCombinationIter, QueryEntityError, QueryIter,
QueryState, WorldQuery,
QueryState, ReadOnlyFetch, WorldQuery,
},
world::{Mut, World},
};
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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,
);
};
}

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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()
}

Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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)
}
}
}
92 changes: 92 additions & 0 deletions crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.rs
View file
Original file line number Diff line number Diff line change
@@ -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");
}
}
Loading

0 comments on commit 9bc894d

Please sign in to comment.