Skip to content

Commit 69b1fe9

Browse files
BoxyUwUcart
authored andcommitted
yeet unsound lifetime annotations on Query methods (bevyengine#4243)
# 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 <mcanders1@gmail.com>
1 parent 008bc87 commit 69b1fe9

File tree

7 files changed

+324
-24
lines changed

7 files changed

+324
-24
lines changed

crates/bevy_ecs/src/system/mod.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -804,13 +804,13 @@ mod tests {
804804
}
805805
fn hold_component<'w>(&mut self, world: &'w World, entity: Entity) -> Holder<'w> {
806806
let q = self.state_q.get(world);
807-
let a = q.get(entity).unwrap();
807+
let a = q.get_inner(entity).unwrap();
808808
Holder { value: a }
809809
}
810810
fn hold_components<'w>(&mut self, world: &'w World) -> Vec<Holder<'w>> {
811811
let mut components = Vec::new();
812812
let q = self.state_q.get(world);
813-
for a in q.iter() {
813+
for a in q.iter_inner() {
814814
components.push(Holder { value: a });
815815
}
816816
components

crates/bevy_ecs/src/system/query.rs

+108-19
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use crate::{
33
entity::Entity,
44
query::{
55
Fetch, FilterFetch, NopFetch, QueryCombinationIter, QueryEntityError, QueryIter,
6-
QueryState, WorldQuery,
6+
QueryState, ReadOnlyFetch, WorldQuery,
77
},
88
world::{Mut, World},
99
};
@@ -299,7 +299,7 @@ where
299299
/// # bevy_ecs::system::assert_is_system(report_names_system);
300300
/// ```
301301
#[inline]
302-
pub fn iter(&'s self) -> QueryIter<'w, 's, Q, Q::ReadOnlyFetch, F> {
302+
pub fn iter(&self) -> QueryIter<'_, 's, Q, Q::ReadOnlyFetch, F> {
303303
// SAFE: system runs without conflicts with other systems.
304304
// same-system queries have runtime borrow checks when they conflict
305305
unsafe {
@@ -454,17 +454,19 @@ where
454454
/// # bevy_ecs::system::assert_is_system(report_names_system);
455455
/// ```
456456
#[inline]
457-
pub fn for_each<FN: FnMut(<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item)>(&'s self, f: FN) {
457+
pub fn for_each<'this>(
458+
&'this self,
459+
f: impl FnMut(<Q::ReadOnlyFetch as Fetch<'this, 's>>::Item),
460+
) {
458461
// SAFE: system runs without conflicts with other systems.
459462
// same-system queries have runtime borrow checks when they conflict
460463
unsafe {
461-
self.state
462-
.for_each_unchecked_manual::<Q::ReadOnlyFetch, FN>(
463-
self.world,
464-
f,
465-
self.last_change_tick,
466-
self.change_tick,
467-
);
464+
self.state.for_each_unchecked_manual::<Q::ReadOnlyFetch, _>(
465+
self.world,
466+
f,
467+
self.last_change_tick,
468+
self.change_tick,
469+
);
468470
};
469471
}
470472

@@ -524,17 +526,17 @@ where
524526
///* `batch_size` - The number of batches to spawn
525527
///* `f` - The function to run on each item in the query
526528
#[inline]
527-
pub fn par_for_each<FN: Fn(<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item) + Send + Sync + Clone>(
528-
&'s self,
529+
pub fn par_for_each<'this>(
530+
&'this self,
529531
task_pool: &TaskPool,
530532
batch_size: usize,
531-
f: FN,
533+
f: impl Fn(<Q::ReadOnlyFetch as Fetch<'this, 's>>::Item) + Send + Sync + Clone,
532534
) {
533535
// SAFE: system runs without conflicts with other systems. same-system queries have runtime
534536
// borrow checks when they conflict
535537
unsafe {
536538
self.state
537-
.par_for_each_unchecked_manual::<Q::ReadOnlyFetch, FN>(
539+
.par_for_each_unchecked_manual::<Q::ReadOnlyFetch, _>(
538540
self.world,
539541
task_pool,
540542
batch_size,
@@ -601,9 +603,9 @@ where
601603
/// ```
602604
#[inline]
603605
pub fn get(
604-
&'s self,
606+
&self,
605607
entity: Entity,
606-
) -> Result<<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item, QueryEntityError> {
608+
) -> Result<<Q::ReadOnlyFetch as Fetch<'_, 's>>::Item, QueryEntityError> {
607609
// SAFE: system runs without conflicts with other systems.
608610
// same-system queries have runtime borrow checks when they conflict
609611
unsafe {
@@ -834,7 +836,7 @@ where
834836
/// Panics if the number of query results is not exactly one. Use
835837
/// [`get_single`](Self::get_single) to return a `Result` instead of panicking.
836838
#[track_caller]
837-
pub fn single(&'s self) -> <Q::ReadOnlyFetch as Fetch<'w, 's>>::Item {
839+
pub fn single(&self) -> <Q::ReadOnlyFetch as Fetch<'_, 's>>::Item {
838840
self.get_single().unwrap()
839841
}
840842

@@ -870,8 +872,8 @@ where
870872
/// # bevy_ecs::system::assert_is_system(player_scoring_system);
871873
/// ```
872874
pub fn get_single(
873-
&'s self,
874-
) -> Result<<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item, QuerySingleError> {
875+
&self,
876+
) -> Result<<Q::ReadOnlyFetch as Fetch<'_, 's>>::Item, QuerySingleError> {
875877
let mut query = self.iter();
876878
let first = query.next();
877879
let extra = query.next().is_some();
@@ -1038,3 +1040,90 @@ pub enum QuerySingleError {
10381040
#[error("Multiple entities fit the query {0}!")]
10391041
MultipleEntities(&'static str),
10401042
}
1043+
1044+
impl<'w, 's, Q: WorldQuery, F: WorldQuery> Query<'w, 's, Q, F>
1045+
where
1046+
F::Fetch: FilterFetch,
1047+
Q::Fetch: ReadOnlyFetch,
1048+
{
1049+
/// Returns the query result for the given [`Entity`], with the actual "inner" world lifetime.
1050+
///
1051+
/// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is
1052+
/// returned instead.
1053+
///
1054+
/// This can only return immutable data (mutable data will be cast to an immutable form).
1055+
/// See [`get_mut`](Self::get_mut) for queries that contain at least one mutable component.
1056+
///
1057+
/// # Example
1058+
///
1059+
/// Here, `get` is used to retrieve the exact query result of the entity specified by the
1060+
/// `SelectedCharacter` resource.
1061+
///
1062+
/// ```
1063+
/// # use bevy_ecs::prelude::*;
1064+
/// #
1065+
/// # struct SelectedCharacter { entity: Entity }
1066+
/// # #[derive(Component)]
1067+
/// # struct Character { name: String }
1068+
/// #
1069+
/// fn print_selected_character_name_system(
1070+
/// query: Query<&Character>,
1071+
/// selection: Res<SelectedCharacter>
1072+
/// )
1073+
/// {
1074+
/// if let Ok(selected_character) = query.get(selection.entity) {
1075+
/// println!("{}", selected_character.name);
1076+
/// }
1077+
/// }
1078+
/// # bevy_ecs::system::assert_is_system(print_selected_character_name_system);
1079+
/// ```
1080+
#[inline]
1081+
pub fn get_inner(
1082+
&'s self,
1083+
entity: Entity,
1084+
) -> Result<<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item, QueryEntityError> {
1085+
// SAFE: system runs without conflicts with other systems.
1086+
// same-system queries have runtime borrow checks when they conflict
1087+
unsafe {
1088+
self.state.get_unchecked_manual::<Q::ReadOnlyFetch>(
1089+
self.world,
1090+
entity,
1091+
self.last_change_tick,
1092+
self.change_tick,
1093+
)
1094+
}
1095+
}
1096+
1097+
/// Returns an [`Iterator`] over the query results, with the actual "inner" world lifetime.
1098+
///
1099+
/// This can only return immutable data (mutable data will be cast to an immutable form).
1100+
/// See [`Self::iter_mut`] for queries that contain at least one mutable component.
1101+
///
1102+
/// # Example
1103+
///
1104+
/// Here, the `report_names_system` iterates over the `Player` component of every entity
1105+
/// that contains it:
1106+
///
1107+
/// ```
1108+
/// # use bevy_ecs::prelude::*;
1109+
/// #
1110+
/// # #[derive(Component)]
1111+
/// # struct Player { name: String }
1112+
/// #
1113+
/// fn report_names_system(query: Query<&Player>) {
1114+
/// for player in query.iter() {
1115+
/// println!("Say hello to {}!", player.name);
1116+
/// }
1117+
/// }
1118+
/// # bevy_ecs::system::assert_is_system(report_names_system);
1119+
/// ```
1120+
#[inline]
1121+
pub fn iter_inner(&'s self) -> QueryIter<'w, 's, Q, Q::ReadOnlyFetch, F> {
1122+
// SAFE: system runs without conflicts with other systems.
1123+
// same-system queries have runtime borrow checks when they conflict
1124+
unsafe {
1125+
self.state
1126+
.iter_unchecked_manual(self.world, self.last_change_tick, self.change_tick)
1127+
}
1128+
}
1129+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
use bevy_ecs::prelude::*;
2+
use bevy_ecs::system::SystemState;
3+
4+
#[derive(Component, Eq, PartialEq, Debug)]
5+
struct Foo(u32);
6+
7+
fn main() {
8+
let mut world = World::default();
9+
let e = world.spawn().insert(Foo(10_u32)).id();
10+
11+
let mut system_state = SystemState::<Query<&mut Foo>>::new(&mut world);
12+
{
13+
let mut query = system_state.get_mut(&mut world);
14+
dbg!("hi");
15+
{
16+
let data: &Foo = query.get(e).unwrap();
17+
let mut data2: Mut<Foo> = query.get_mut(e).unwrap();
18+
assert_eq!(data, &mut *data2); // oops UB
19+
}
20+
21+
{
22+
let mut data2: Mut<Foo> = query.get_mut(e).unwrap();
23+
let data: &Foo = query.get(e).unwrap();
24+
assert_eq!(data, &mut *data2); // oops UB
25+
}
26+
27+
{
28+
let data: &Foo = query.get_component::<Foo>(e).unwrap();
29+
let mut data2: Mut<Foo> = query.get_component_mut(e).unwrap();
30+
assert_eq!(data, &mut *data2); // oops UB
31+
}
32+
33+
{
34+
let mut data2: Mut<Foo> = query.get_component_mut(e).unwrap();
35+
let data: &Foo = query.get_component::<Foo>(e).unwrap();
36+
assert_eq!(data, &mut *data2); // oops UB
37+
}
38+
39+
{
40+
let data: &Foo = query.single();
41+
let mut data2: Mut<Foo> = query.single_mut();
42+
assert_eq!(data, &mut *data2); // oops UB
43+
}
44+
45+
{
46+
let mut data2: Mut<Foo> = query.single_mut();
47+
let data: &Foo = query.single();
48+
assert_eq!(data, &mut *data2); // oops UB
49+
}
50+
51+
{
52+
let data: &Foo = query.get_single().unwrap();
53+
let mut data2: Mut<Foo> = query.get_single_mut().unwrap();
54+
assert_eq!(data, &mut *data2); // oops UB
55+
}
56+
57+
{
58+
let mut data2: Mut<Foo> = query.get_single_mut().unwrap();
59+
let data: &Foo = query.get_single().unwrap();
60+
assert_eq!(data, &mut *data2); // oops UB
61+
}
62+
63+
{
64+
let data: &Foo = query.iter().next().unwrap();
65+
let mut data2: Mut<Foo> = query.iter_mut().next().unwrap();
66+
assert_eq!(data, &mut *data2); // oops UB
67+
}
68+
69+
{
70+
let mut data2: Mut<Foo> = query.iter_mut().next().unwrap();
71+
let data: &Foo = query.iter().next().unwrap();
72+
assert_eq!(data, &mut *data2); // oops UB
73+
}
74+
75+
{
76+
let mut opt_data: Option<&Foo> = None;
77+
let mut opt_data_2: Option<Mut<Foo>> = None;
78+
query.for_each(|data| opt_data = Some(data));
79+
query.for_each_mut(|data| opt_data_2 = Some(data));
80+
assert_eq!(opt_data.unwrap(), &mut *opt_data_2.unwrap()); // oops UB
81+
}
82+
83+
{
84+
let mut opt_data_2: Option<Mut<Foo>> = None;
85+
let mut opt_data: Option<&Foo> = None;
86+
query.for_each_mut(|data| opt_data_2 = Some(data));
87+
query.for_each(|data| opt_data = Some(data));
88+
assert_eq!(opt_data.unwrap(), &mut *opt_data_2.unwrap()); // oops UB
89+
}
90+
dbg!("bye");
91+
}
92+
}

0 commit comments

Comments
 (0)