Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make QueryState::transmute&co validate the world of the &Components used #14631

Merged
merged 4 commits into from
Aug 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions crates/bevy_ecs/src/observer/runner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -375,8 +375,7 @@ fn observer_system_runner<E: Event, B: Bundle>(
};

// TODO: Move this check into the observer cache to avoid dynamic dispatch
// SAFETY: We only access world metadata
let last_trigger = unsafe { world.world_metadata() }.last_trigger_id();
let last_trigger = world.last_trigger_id();
if state.last_trigger_id == last_trigger {
return;
}
Expand Down
28 changes: 7 additions & 21 deletions crates/bevy_ecs/src/query/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -362,9 +362,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {

let world = self.world;

let query_lens_state = self
.query_state
.transmute_filtered::<(L, Entity), F>(world.components());
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);

// SAFETY:
// `self.world` has permission to access the required components.
Expand Down Expand Up @@ -456,9 +454,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {

let world = self.world;

let query_lens_state = self
.query_state
.transmute_filtered::<(L, Entity), F>(world.components());
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);

// SAFETY:
// `self.world` has permission to access the required components.
Expand Down Expand Up @@ -558,9 +554,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {

let world = self.world;

let query_lens_state = self
.query_state
.transmute_filtered::<(L, Entity), F>(world.components());
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);

// SAFETY:
// `self.world` has permission to access the required components.
Expand Down Expand Up @@ -626,9 +620,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {

let world = self.world;

let query_lens_state = self
.query_state
.transmute_filtered::<(L, Entity), F>(world.components());
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);

// SAFETY:
// `self.world` has permission to access the required components.
Expand Down Expand Up @@ -757,9 +749,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {

let world = self.world;

let query_lens_state = self
.query_state
.transmute_filtered::<(L, Entity), F>(world.components());
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);

// SAFETY:
// `self.world` has permission to access the required components.
Expand Down Expand Up @@ -828,9 +818,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {

let world = self.world;

let query_lens_state = self
.query_state
.transmute_filtered::<(L, Entity), F>(world.components());
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);

// SAFETY:
// `self.world` has permission to access the required components.
Expand Down Expand Up @@ -900,9 +888,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {

let world = self.world;

let query_lens_state = self
.query_state
.transmute_filtered::<(L, Entity), F>(world.components());
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);

// SAFETY:
// `self.world` has permission to access the required components.
Expand Down
98 changes: 60 additions & 38 deletions crates/bevy_ecs/src/query/state.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use crate::{
archetype::{Archetype, ArchetypeComponentId, ArchetypeGeneration, ArchetypeId},
batching::BatchingStrategy,
component::{ComponentId, Components, Tick},
component::{ComponentId, Tick},
entity::Entity,
prelude::FromWorld,
query::{
Expand Down Expand Up @@ -476,21 +476,27 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
/// You might end up with a mix of archetypes that only matched the original query + archetypes that only match
/// the new [`QueryState`]. Most of the safe methods on [`QueryState`] call [`QueryState::update_archetypes`] internally, so this
/// best used through a [`Query`](crate::system::Query).
pub fn transmute<NewD: QueryData>(&self, components: &Components) -> QueryState<NewD> {
self.transmute_filtered::<NewD, ()>(components)
pub fn transmute<'a, NewD: QueryData>(
&self,
world: impl Into<UnsafeWorldCell<'a>>,
) -> QueryState<NewD> {
self.transmute_filtered::<NewD, ()>(world.into())
}

/// Creates a new [`QueryState`] with the same underlying [`FilteredAccess`], matched tables and archetypes
/// as self but with a new type signature.
///
/// Panics if `NewD` or `NewF` require accesses that this query does not have.
pub fn transmute_filtered<NewD: QueryData, NewF: QueryFilter>(
pub fn transmute_filtered<'a, NewD: QueryData, NewF: QueryFilter>(
&self,
components: &Components,
world: impl Into<UnsafeWorldCell<'a>>,
) -> QueryState<NewD, NewF> {
let world = world.into();
self.validate_world(world.id());

let mut component_access = FilteredAccess::default();
let mut fetch_state = NewD::get_state(components).expect("Could not create fetch_state, Please initialize all referenced components before transmuting.");
let filter_state = NewF::get_state(components).expect("Could not create filter_state, Please initialize all referenced components before transmuting.");
let mut fetch_state = NewD::get_state(world.components()).expect("Could not create fetch_state, Please initialize all referenced components before transmuting.");
let filter_state = NewF::get_state(world.components()).expect("Could not create filter_state, Please initialize all referenced components before transmuting.");

NewD::set_access(&mut fetch_state, &self.component_access);
NewD::update_component_access(&fetch_state, &mut component_access);
Expand Down Expand Up @@ -542,12 +548,12 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
/// ## Panics
///
/// Will panic if `NewD` contains accesses not in `Q` or `OtherQ`.
pub fn join<OtherD: QueryData, NewD: QueryData>(
pub fn join<'a, OtherD: QueryData, NewD: QueryData>(
&self,
components: &Components,
world: impl Into<UnsafeWorldCell<'a>>,
other: &QueryState<OtherD>,
) -> QueryState<NewD, ()> {
self.join_filtered::<_, (), NewD, ()>(components, other)
self.join_filtered::<_, (), NewD, ()>(world, other)
}

/// Use this to combine two queries. The data accessed will be the intersection
Expand All @@ -557,23 +563,28 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
///
/// Will panic if `NewD` or `NewF` requires accesses not in `Q` or `OtherQ`.
pub fn join_filtered<
'a,
OtherD: QueryData,
OtherF: QueryFilter,
NewD: QueryData,
NewF: QueryFilter,
>(
&self,
components: &Components,
world: impl Into<UnsafeWorldCell<'a>>,
other: &QueryState<OtherD, OtherF>,
) -> QueryState<NewD, NewF> {
if self.world_id != other.world_id {
panic!("Joining queries initialized on different worlds is not allowed.");
}

let world = world.into();

self.validate_world(world.id());

let mut component_access = FilteredAccess::default();
let mut new_fetch_state = NewD::get_state(components)
let mut new_fetch_state = NewD::get_state(world.components())
.expect("Could not create fetch_state, Please initialize all referenced components before transmuting.");
let new_filter_state = NewF::get_state(components)
let new_filter_state = NewF::get_state(world.components())
.expect("Could not create filter_state, Please initialize all referenced components before transmuting.");

NewD::set_access(&mut new_fetch_state, &self.component_access);
Expand Down Expand Up @@ -1779,7 +1790,7 @@ mod tests {
world.spawn((A(1), B(0)));

let query_state = world.query::<(&A, &B)>();
let mut new_query_state = query_state.transmute::<&A>(world.components());
let mut new_query_state = query_state.transmute::<&A>(&world);
assert_eq!(new_query_state.iter(&world).len(), 1);
let a = new_query_state.single(&world);

Expand All @@ -1793,7 +1804,7 @@ mod tests {
world.spawn((A(1), B(0), C(0)));

let query_state = world.query_filtered::<(&A, &B), Without<C>>();
let mut new_query_state = query_state.transmute::<&A>(world.components());
let mut new_query_state = query_state.transmute::<&A>(&world);
// even though we change the query to not have Without<C>, we do not get the component with C.
let a = new_query_state.single(&world);

Expand All @@ -1807,7 +1818,7 @@ mod tests {
let entity = world.spawn(A(10)).id();

let q = world.query::<()>();
let mut q = q.transmute::<Entity>(world.components());
let mut q = q.transmute::<Entity>(&world);
assert_eq!(q.single(&world), entity);
}

Expand All @@ -1817,11 +1828,11 @@ mod tests {
world.spawn(A(10));

let q = world.query::<&A>();
let mut new_q = q.transmute::<Ref<A>>(world.components());
let mut new_q = q.transmute::<Ref<A>>(&world);
assert!(new_q.single(&world).is_added());

let q = world.query::<Ref<A>>();
let _ = q.transmute::<&A>(world.components());
let _ = q.transmute::<&A>(&world);
}

#[test]
Expand All @@ -1830,8 +1841,8 @@ mod tests {
world.spawn(A(0));

let q = world.query::<&mut A>();
let _ = q.transmute::<Ref<A>>(world.components());
let _ = q.transmute::<&A>(world.components());
let _ = q.transmute::<Ref<A>>(&world);
let _ = q.transmute::<&A>(&world);
}

#[test]
Expand All @@ -1840,7 +1851,7 @@ mod tests {
world.spawn(A(0));

let q: QueryState<EntityMut<'_>> = world.query::<EntityMut>();
let _ = q.transmute::<EntityRef>(world.components());
let _ = q.transmute::<EntityRef>(&world);
}

#[test]
Expand All @@ -1849,8 +1860,8 @@ mod tests {
world.spawn((A(0), B(0)));

let query_state = world.query::<(Option<&A>, &B)>();
let _ = query_state.transmute::<Option<&A>>(world.components());
let _ = query_state.transmute::<&B>(world.components());
let _ = query_state.transmute::<Option<&A>>(&world);
let _ = query_state.transmute::<&B>(&world);
}

#[test]
Expand All @@ -1864,7 +1875,7 @@ mod tests {
world.spawn(A(0));

let query_state = world.query::<&A>();
let mut _new_query_state = query_state.transmute::<(&A, &B)>(world.components());
let mut _new_query_state = query_state.transmute::<(&A, &B)>(&world);
}

#[test]
Expand All @@ -1876,7 +1887,7 @@ mod tests {
world.spawn(A(0));

let query_state = world.query::<&A>();
let mut _new_query_state = query_state.transmute::<&mut A>(world.components());
let mut _new_query_state = query_state.transmute::<&mut A>(&world);
}

#[test]
Expand All @@ -1888,7 +1899,7 @@ mod tests {
world.spawn(C(0));

let query_state = world.query::<Option<&A>>();
let mut new_query_state = query_state.transmute::<&A>(world.components());
let mut new_query_state = query_state.transmute::<&A>(&world);
let x = new_query_state.single(&world);
assert_eq!(x.0, 1234);
}
Expand All @@ -1902,15 +1913,15 @@ mod tests {
world.init_component::<A>();

let q = world.query::<EntityRef>();
let _ = q.transmute::<&A>(world.components());
let _ = q.transmute::<&A>(&world);
}

#[test]
fn can_transmute_filtered_entity() {
let mut world = World::new();
let entity = world.spawn((A(0), B(1))).id();
let query = QueryState::<(Entity, &A, &B)>::new(&mut world)
.transmute::<FilteredEntityRef>(world.components());
let query =
QueryState::<(Entity, &A, &B)>::new(&mut world).transmute::<FilteredEntityRef>(&world);

let mut query = query;
// Our result is completely untyped
Expand All @@ -1927,7 +1938,7 @@ mod tests {
let entity_a = world.spawn(A(0)).id();

let mut query = QueryState::<(Entity, &A, Has<B>)>::new(&mut world)
.transmute_filtered::<(Entity, Has<B>), Added<A>>(world.components());
.transmute_filtered::<(Entity, Has<B>), Added<A>>(&world);

assert_eq!((entity_a, false), query.single(&world));

Expand All @@ -1947,7 +1958,7 @@ mod tests {
let entity_a = world.spawn(A(0)).id();

let mut detection_query = QueryState::<(Entity, &A)>::new(&mut world)
.transmute_filtered::<Entity, Changed<A>>(world.components());
.transmute_filtered::<Entity, Changed<A>>(&world);

let mut change_query = QueryState::<&mut A>::new(&mut world);
assert_eq!(entity_a, detection_query.single(&world));
Expand All @@ -1970,7 +1981,20 @@ mod tests {
world.init_component::<A>();
world.init_component::<B>();
let query = QueryState::<&A>::new(&mut world);
let _new_query = query.transmute_filtered::<Entity, Changed<B>>(world.components());
let _new_query = query.transmute_filtered::<Entity, Changed<B>>(&world);
}

// Regression test for #14629
#[test]
#[should_panic]
fn transmute_with_different_world() {
let mut world = World::new();
world.spawn((A(1), B(2)));

let mut world2 = World::new();
world2.init_component::<B>();

world.query::<(&A, &B)>().transmute::<&B>(&world2);
}

#[test]
Expand All @@ -1983,8 +2007,7 @@ mod tests {

let query_1 = QueryState::<&A, Without<C>>::new(&mut world);
let query_2 = QueryState::<&B, Without<C>>::new(&mut world);
let mut new_query: QueryState<Entity, ()> =
query_1.join_filtered(world.components(), &query_2);
let mut new_query: QueryState<Entity, ()> = query_1.join_filtered(&world, &query_2);

assert_eq!(new_query.single(&world), entity_ab);
}
Expand All @@ -1999,8 +2022,7 @@ mod tests {

let query_1 = QueryState::<&A>::new(&mut world);
let query_2 = QueryState::<&B, Without<C>>::new(&mut world);
let mut new_query: QueryState<Entity, ()> =
query_1.join_filtered(world.components(), &query_2);
let mut new_query: QueryState<Entity, ()> = query_1.join_filtered(&world, &query_2);

assert!(new_query.get(&world, entity_ab).is_ok());
// should not be able to get entity with c.
Expand All @@ -2016,7 +2038,7 @@ mod tests {
world.init_component::<C>();
let query_1 = QueryState::<&A>::new(&mut world);
let query_2 = QueryState::<&B>::new(&mut world);
let _query: QueryState<&C> = query_1.join(world.components(), &query_2);
let _query: QueryState<&C> = query_1.join(&world, &query_2);
}

#[test]
Expand All @@ -2030,6 +2052,6 @@ mod tests {
let mut world = World::new();
let query_1 = QueryState::<&A, Without<C>>::new(&mut world);
let query_2 = QueryState::<&B, Without<C>>::new(&mut world);
let _: QueryState<Entity, Changed<C>> = query_1.join_filtered(world.components(), &query_2);
let _: QueryState<Entity, Changed<C>> = query_1.join_filtered(&world, &query_2);
}
}
6 changes: 2 additions & 4 deletions crates/bevy_ecs/src/system/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1368,8 +1368,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
pub fn transmute_lens_filtered<NewD: QueryData, NewF: QueryFilter>(
&mut self,
) -> QueryLens<'_, NewD, NewF> {
let components = self.world.components();
let state = self.state.transmute_filtered::<NewD, NewF>(components);
let state = self.state.transmute_filtered::<NewD, NewF>(self.world);
QueryLens {
world: self.world,
state,
Expand Down Expand Up @@ -1460,10 +1459,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
&mut self,
other: &mut Query<OtherD, OtherF>,
) -> QueryLens<'_, NewD, NewF> {
let components = self.world.components();
let state = self
.state
.join_filtered::<OtherD, OtherF, NewD, NewF>(components, other.state);
.join_filtered::<OtherD, OtherF, NewD, NewF>(self.world, other.state);
QueryLens {
world: self.world,
state,
Expand Down
Loading