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

Introduce methods on QueryState to obtain a Query #15858

Merged
merged 7 commits into from
Feb 5, 2025
20 changes: 10 additions & 10 deletions crates/bevy_ecs/compile_fail/tests/ui/query_to_readonly.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,29 +6,29 @@ struct Foo;
fn for_loops(mut query: Query<&mut Foo>) {
// this should fail to compile
for _ in query.iter_mut() {
for _ in query.to_readonly().iter() {}
for _ in query.as_readonly().iter() {}
//~^ E0502
}

// this should fail to compile
for _ in query.to_readonly().iter() {
for _ in query.as_readonly().iter() {
for _ in query.iter_mut() {}
//~^ E0502
}

// this should *not* fail to compile
for _ in query.to_readonly().iter() {
for _ in query.to_readonly().iter() {}
for _ in query.as_readonly().iter() {
for _ in query.as_readonly().iter() {}
}

// this should *not* fail to compile
for _ in query.to_readonly().iter() {
for _ in query.as_readonly().iter() {
for _ in query.iter() {}
}

// this should *not* fail to compile
for _ in query.iter() {
for _ in query.to_readonly().iter() {}
for _ in query.as_readonly().iter() {}
}
}

Expand All @@ -38,11 +38,11 @@ fn single_mut_query(mut query: Query<&mut Foo>) {
let mut mut_foo = query.single_mut();

// This solves "temporary value dropped while borrowed"
let readonly_query = query.to_readonly();
let readonly_query = query.as_readonly();
//~^ E0502

let ref_foo = readonly_query.single();

*mut_foo = Foo;

println!("{ref_foo:?}");
Expand All @@ -51,7 +51,7 @@ fn single_mut_query(mut query: Query<&mut Foo>) {
// this should fail to compile
{
// This solves "temporary value dropped while borrowed"
let readonly_query = query.to_readonly();
let readonly_query = query.as_readonly();

let ref_foo = readonly_query.single();

Expand All @@ -66,7 +66,7 @@ fn single_mut_query(mut query: Query<&mut Foo>) {
// this should *not* fail to compile
{
// This solves "temporary value dropped while borrowed"
let readonly_query = query.to_readonly();
let readonly_query = query.as_readonly();

let readonly_foo = readonly_query.single();

Expand Down
18 changes: 3 additions & 15 deletions crates/bevy_ecs/src/query/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -849,13 +849,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
// SAFETY:
// `self.world` has permission to access the required components.
// The original query iter has not been iterated on, so no items are aliased from it.
let query_lens = unsafe {
query_lens_state.iter_unchecked_manual(
world,
world.last_change_tick(),
world.change_tick(),
)
};
let query_lens = unsafe { query_lens_state.query_unchecked_manual(world) }.into_iter();
let mut keyed_query: Vec<_> = query_lens
.map(|(key, entity)| (key, NeutralOrd(entity)))
.collect();
Expand Down Expand Up @@ -1689,14 +1683,8 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
// SAFETY:
// `self.world` has permission to access the required components.
// The original query iter has not been iterated on, so no items are aliased from it.
let query_lens = unsafe {
query_lens_state.iter_many_unchecked_manual(
self.entity_iter,
world,
world.last_change_tick(),
world.change_tick(),
)
};
let query_lens = unsafe { query_lens_state.query_unchecked_manual(world) }
.iter_many_inner(self.entity_iter);
let mut keyed_query: Vec<_> = query_lens
.map(|(key, entity)| (key, NeutralOrd(entity)))
.collect();
Expand Down
122 changes: 119 additions & 3 deletions crates/bevy_ecs/src/query/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ use crate::{
WorldQuery,
},
storage::{SparseSetIndex, TableId},
system::Query,
world::{unsafe_world_cell::UnsafeWorldCell, World, WorldId},
};

Expand Down Expand Up @@ -154,9 +155,7 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
pub fn matched_archetypes(&self) -> impl Iterator<Item = ArchetypeId> + '_ {
self.matched_archetypes.ones().map(ArchetypeId::new)
}
}

impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
/// Creates a new [`QueryState`] from a given [`World`] and inherits the result of `world.id()`.
pub fn new(world: &mut World) -> Self {
let mut state = Self::new_uninitialized(world);
Expand Down Expand Up @@ -319,6 +318,123 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
state
}

/// Creates a [`Query`] from the given [`QueryState`] and [`World`].
///
/// This can only be called for read-only queries, see [`Self::query_mut`] for write-queries.
chescock marked this conversation as resolved.
Show resolved Hide resolved
pub fn query<'w, 's>(&'s mut self, world: &'w World) -> Query<'w, 's, D::ReadOnly, F> {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh my gosh finally.

self.update_archetypes(world);
self.query_manual(world)
}

/// Creates a [`Query`] from the given [`QueryState`] and [`World`].
///
/// This method is slightly more efficient than [`QueryState::query`] in some situations, since
/// it does not update this instance's internal cache. The resulting query may skip an entity that
/// belongs to an archetype that has not been cached.
///
/// To ensure that the cache is up to date, call [`QueryState::update_archetypes`] before this method.
/// The cache is also updated in [`QueryState::new`], `QueryState::get`, or any method with mutable
chescock marked this conversation as resolved.
Show resolved Hide resolved
/// access to `self`.
///
/// This can only be called for read-only queries, see [`Self::query_mut`] for mutable queries.
pub fn query_manual<'w, 's>(&'s self, world: &'w World) -> Query<'w, 's, D::ReadOnly, F> {
// SAFETY: We have read access to the entire world, and we call `as_readonly()` so the query only performs read access.
unsafe {
self.as_readonly()
.query_unchecked_manual(world.as_unsafe_world_cell_readonly())
}
}

/// Creates a [`Query`] from the given [`QueryState`] and [`World`].
pub fn query_mut<'w, 's>(&'s mut self, world: &'w mut World) -> Query<'w, 's, D, F> {
let last_run = world.last_change_tick();
let this_run = world.change_tick();
// SAFETY: We have exclusive access to the entire world.
unsafe { self.query_unchecked_with_ticks(world.as_unsafe_world_cell(), last_run, this_run) }
}

/// Creates a [`Query`] from the given [`QueryState`] and [`World`].
///
/// # Safety
///
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
/// have unique access to the components they query.
pub unsafe fn query_unchecked<'w, 's>(
&'s mut self,
world: UnsafeWorldCell<'w>,
) -> Query<'w, 's, D, F> {
self.update_archetypes_unsafe_world_cell(world);
self.query_unchecked_manual(world)
chescock marked this conversation as resolved.
Show resolved Hide resolved
}

/// Creates a [`Query`] from the given [`QueryState`] and [`World`].
///
/// This method is slightly more efficient than [`QueryState::query_unchecked`] in some situations, since
/// it does not update this instance's internal cache. The resulting query may skip an entity that
/// belongs to an archetype that has not been cached.
///
/// To ensure that the cache is up to date, call [`QueryState::update_archetypes`] before this method.
/// The cache is also updated in [`QueryState::new`], `QueryState::get`, or any method with mutable
/// access to `self`.
///
/// # Safety
///
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
/// have unique access to the components they query.
pub unsafe fn query_unchecked_manual<'w, 's>(
&'s self,
world: UnsafeWorldCell<'w>,
) -> Query<'w, 's, D, F> {
let last_run = world.last_change_tick();
let this_run = world.change_tick();
// SAFETY: The caller ensured we have the correct access to the world.
unsafe { self.query_unchecked_manual_with_ticks(world, last_run, this_run) }
}

/// Creates a [`Query`] from the given [`QueryState`] and [`World`].
///
/// # Safety
///
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
/// have unique access to the components they query.
pub unsafe fn query_unchecked_with_ticks<'w, 's>(
&'s mut self,
world: UnsafeWorldCell<'w>,
last_run: Tick,
this_run: Tick,
) -> Query<'w, 's, D, F> {
self.update_archetypes_unsafe_world_cell(world);
// SAFETY: The caller ensured we have the correct access to the world.
unsafe { self.query_unchecked_manual_with_ticks(world, last_run, this_run) }
}

/// Creates a [`Query`] from the given [`QueryState`] and [`World`].
///
/// This method is slightly more efficient than [`QueryState::query_unchecked_with_ticks`] in some situations, since
/// it does not update this instance's internal cache. The resulting query may skip an entity that
/// belongs to an archetype that has not been cached.
///
/// To ensure that the cache is up to date, call [`QueryState::update_archetypes`] before this method.
/// The cache is also updated in [`QueryState::new`], `QueryState::get`, or any method with mutable
/// access to `self`.
///
/// # Safety
///
/// This does not check for mutable query correctness. To be safe, make sure mutable queries
/// have unique access to the components they query.
pub unsafe fn query_unchecked_manual_with_ticks<'w, 's>(
&'s self,
world: UnsafeWorldCell<'w>,
last_run: Tick,
this_run: Tick,
) -> Query<'w, 's, D, F> {
self.validate_world(world.id());
// SAFETY:
// - The caller ensured we have the correct access to the world.
// - `validate_world` did not panic, so the world matches.
unsafe { Query::new(world, self, last_run, this_run) }
}

/// Checks if the query is empty for the given [`World`], where the last change and current tick are given.
///
/// This is equivalent to `self.iter().next().is_none()`, and thus the worst case runtime will be `O(n)`
Expand Down Expand Up @@ -624,7 +740,7 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
/// You should not call [`update_archetypes`](Self::update_archetypes) on the returned [`QueryState`] as the result will be unpredictable.
/// 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).
/// best used through a [`Query`]
pub fn transmute<'a, NewD: QueryData>(
&self,
world: impl Into<UnsafeWorldCell<'a>>,
Expand Down
36 changes: 9 additions & 27 deletions crates/bevy_ecs/src/system/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -334,7 +334,7 @@ mod tests {
archetype::{ArchetypeComponentId, Archetypes},
bundle::Bundles,
change_detection::DetectChanges,
component::{Component, Components, Tick},
component::{Component, Components},
entity::{Entities, Entity},
prelude::{AnyOf, EntityRef},
query::{Added, Changed, Or, With, Without},
Expand Down Expand Up @@ -1361,7 +1361,7 @@ mod tests {
fn mutable_query(mut query: Query<&mut A>) {
for _ in &mut query {}

immutable_query(query.to_readonly());
immutable_query(query.as_readonly());
}

fn immutable_query(_: Query<&A>) {}
Expand All @@ -1376,7 +1376,7 @@ mod tests {
fn mutable_query(mut query: Query<Option<&mut A>>) {
for _ in &mut query {}

immutable_query(query.to_readonly());
immutable_query(query.as_readonly());
}

fn immutable_query(_: Query<Option<&A>>) {}
Expand All @@ -1391,7 +1391,7 @@ mod tests {
fn mutable_query(mut query: Query<(&mut A, &B)>) {
for _ in &mut query {}

immutable_query(query.to_readonly());
immutable_query(query.as_readonly());
}

fn immutable_query(_: Query<(&A, &B)>) {}
Expand All @@ -1406,7 +1406,7 @@ mod tests {
fn mutable_query(mut query: Query<(&mut A, &mut B)>) {
for _ in &mut query {}

immutable_query(query.to_readonly());
immutable_query(query.as_readonly());
}

fn immutable_query(_: Query<(&A, &B)>) {}
Expand All @@ -1421,7 +1421,7 @@ mod tests {
fn mutable_query(mut query: Query<(&mut A, &mut B), With<C>>) {
for _ in &mut query {}

immutable_query(query.to_readonly());
immutable_query(query.as_readonly());
}

fn immutable_query(_: Query<(&A, &B), With<C>>) {}
Expand All @@ -1436,7 +1436,7 @@ mod tests {
fn mutable_query(mut query: Query<(&mut A, &mut B), Without<C>>) {
for _ in &mut query {}

immutable_query(query.to_readonly());
immutable_query(query.as_readonly());
}

fn immutable_query(_: Query<(&A, &B), Without<C>>) {}
Expand All @@ -1451,7 +1451,7 @@ mod tests {
fn mutable_query(mut query: Query<(&mut A, &mut B), Added<C>>) {
for _ in &mut query {}

immutable_query(query.to_readonly());
immutable_query(query.as_readonly());
}

fn immutable_query(_: Query<(&A, &B), Added<C>>) {}
Expand All @@ -1466,7 +1466,7 @@ mod tests {
fn mutable_query(mut query: Query<(&mut A, &mut B), Changed<C>>) {
for _ in &mut query {}

immutable_query(query.to_readonly());
immutable_query(query.as_readonly());
}

fn immutable_query(_: Query<(&A, &B), Changed<C>>) {}
Expand Down Expand Up @@ -1572,24 +1572,6 @@ mod tests {
});
}

#[test]
#[should_panic = "Encountered a mismatched World."]
fn query_validates_world_id() {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why remove this test? Is it no longer applicable?

let mut world1 = World::new();
let world2 = World::new();
let qstate = world1.query::<()>();
// SAFETY: doesnt access anything
let query = unsafe {
Query::new(
world2.as_unsafe_world_cell_readonly(),
&qstate,
Tick::new(0),
Tick::new(0),
)
};
query.iter();
}

#[test]
#[should_panic]
fn assert_system_does_not_conflict() {
Expand Down
Loading