From 9bc894d26bd11d7a7d82d6071e376bc64e5d71ad Mon Sep 17 00:00:00 2001 From: Boxy Date: Tue, 22 Mar 2022 02:49:41 +0000 Subject: [PATCH] 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) { dbg!("hi"); { let data: &Foo = query.get(*e).unwrap(); let data2: Mut = query.get_mut(*e).unwrap(); assert_eq!(data, &*data2); // oops UB } { let data: &Foo = query.single(); let data2: Mut = query.single_mut(); assert_eq!(data, &*data2); // oops UB } { let data: &Foo = query.get_single().unwrap(); let data2: Mut = query.get_single_mut().unwrap(); assert_eq!(data, &*data2); // oops UB } { let data: &Foo = query.iter().next().unwrap(); let data2: Mut = query.iter_mut().next().unwrap(); assert_eq!(data, &*data2); // oops UB } { let mut opt_data: Option<&Foo> = None; let mut opt_data_2: Option> = 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 --- crates/bevy_ecs/src/system/mod.rs | 4 +- crates/bevy_ecs/src/system/query.rs | 127 +++++++++++++++--- .../tests/ui/query_lifetime_safety.rs | 92 +++++++++++++ .../tests/ui/query_lifetime_safety.stderr | 119 ++++++++++++++++ crates/bevy_pbr/src/render/mesh.rs | 2 +- crates/bevy_sprite/src/mesh2d/mesh.rs | 2 +- examples/shader/shader_instancing.rs | 2 +- 7 files changed, 324 insertions(+), 24 deletions(-) create mode 100644 crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.rs create mode 100644 crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.stderr diff --git a/crates/bevy_ecs/src/system/mod.rs b/crates/bevy_ecs/src/system/mod.rs index f79ab002961720..f8a68bd7b7d4ec 100644 --- a/crates/bevy_ecs/src/system/mod.rs +++ b/crates/bevy_ecs/src/system/mod.rs @@ -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> { 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 diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index 8d65633de149da..349121b06bd87b 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -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>::Item)>(&'s self, f: FN) { + pub fn for_each<'this>( + &'this self, + f: impl FnMut(>::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::( - self.world, - f, - self.last_change_tick, - self.change_tick, - ); + self.state.for_each_unchecked_manual::( + 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>::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(>::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::( + .par_for_each_unchecked_manual::( self.world, task_pool, batch_size, @@ -601,9 +603,9 @@ where /// ``` #[inline] pub fn get( - &'s self, + &self, entity: Entity, - ) -> Result<>::Item, QueryEntityError> { + ) -> Result<>::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) -> >::Item { + pub fn single(&self) -> >::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<>::Item, QuerySingleError> { + &self, + ) -> Result<>::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 + /// ) + /// { + /// 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<>::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::( + 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) + } + } +} diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.rs new file mode 100644 index 00000000000000..bb9ff14111ff39 --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.rs @@ -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::>::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 = query.get_mut(e).unwrap(); + assert_eq!(data, &mut *data2); // oops UB + } + + { + let mut data2: Mut = 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::(e).unwrap(); + let mut data2: Mut = query.get_component_mut(e).unwrap(); + assert_eq!(data, &mut *data2); // oops UB + } + + { + let mut data2: Mut = query.get_component_mut(e).unwrap(); + let data: &Foo = query.get_component::(e).unwrap(); + assert_eq!(data, &mut *data2); // oops UB + } + + { + let data: &Foo = query.single(); + let mut data2: Mut = query.single_mut(); + assert_eq!(data, &mut *data2); // oops UB + } + + { + let mut data2: Mut = 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 = query.get_single_mut().unwrap(); + assert_eq!(data, &mut *data2); // oops UB + } + + { + let mut data2: Mut = 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 = query.iter_mut().next().unwrap(); + assert_eq!(data, &mut *data2); // oops UB + } + + { + let mut data2: Mut = 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> = 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> = 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"); + } +} diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.stderr new file mode 100644 index 00000000000000..7b1d0fe610e9c3 --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_lifetime_safety.stderr @@ -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 = 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 = 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::(e).unwrap(); + | ----------------------------- immutable borrow occurs here +29 | let mut data2: Mut = 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 = query.get_component_mut(e).unwrap(); + | -------------------------- mutable borrow occurs here +35 | let data: &Foo = query.get_component::(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 = 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 = 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 = 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 = 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 = 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 = 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 diff --git a/crates/bevy_pbr/src/render/mesh.rs b/crates/bevy_pbr/src/render/mesh.rs index 8d00086881215c..d7d01bad5a9972 100644 --- a/crates/bevy_pbr/src/render/mesh.rs +++ b/crates/bevy_pbr/src/render/mesh.rs @@ -625,7 +625,7 @@ impl EntityRenderCommand for SetMeshViewBindGroup { 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, diff --git a/crates/bevy_sprite/src/mesh2d/mesh.rs b/crates/bevy_sprite/src/mesh2d/mesh.rs index 2c6e31085130aa..4d054b28c94580 100644 --- a/crates/bevy_sprite/src/mesh2d/mesh.rs +++ b/crates/bevy_sprite/src/mesh2d/mesh.rs @@ -401,7 +401,7 @@ impl EntityRenderCommand for SetMesh2dViewBindGroup { 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 diff --git a/examples/shader/shader_instancing.rs b/examples/shader/shader_instancing.rs index 11d5b22ea9ae1d..75f5a90b11349a 100644 --- a/examples/shader/shader_instancing.rs +++ b/examples/shader/shader_instancing.rs @@ -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,