Componentize system-internal QueryState.

crates/bevy_audio/src/audio_output.rs

@@ -54,10 +54,10 @@ pub struct PlaybackDespawnMarker;
 pub struct PlaybackRemoveMarker;
 
 #[derive(SystemParam)]
-pub(crate) struct EarPositions<'w, 's> {
-    pub(crate) query: Query<'w, 's, (Entity, &'static GlobalTransform, &'static SpatialListener)>,
+pub(crate) struct EarPositions<'w> {
+    pub(crate) query: Query<'w, 'w, (Entity, &'static GlobalTransform, &'static SpatialListener)>,
 }
-impl<'w, 's> EarPositions<'w, 's> {
+impl<'w> EarPositions<'w> {
     /// Gets a set of transformed ear positions.
     ///
     /// If there are no listeners, use the default values. If a user has added multiple

crates/bevy_ecs/compile_fail/tests/ui/system_param_derive_readonly.rs

@@ -5,12 +5,11 @@ use bevy_ecs::system::{ReadOnlySystemParam, SystemParam, SystemState};
 struct Foo;
 
 #[derive(SystemParam)]
-struct Mutable<'w, 's> {
-    a: Query<'w, 's, &'static mut Foo>,
+struct Mutable<'w> {
+    a: Query<'w, 'w, &'static mut Foo>,
 }
 
 fn main() {
-
     let mut world = World::default();
     let state = SystemState::<Mutable>::new(&mut world);
     state.get(&world);

crates/bevy_ecs/src/query/state.rs

@@ -3,7 +3,7 @@ use crate::{
     component::{ComponentId, Tick},
     entity::{Entity, EntityEquivalent, EntitySet, UniqueEntityArray},
     entity_disabling::DefaultQueryFilters,
-    prelude::FromWorld,
+    prelude::{Component, FromWorld},
     query::{FilteredAccess, QueryCombinationIter, QueryIter, QueryParIter, WorldQuery},
     storage::{SparseSetIndex, TableId},
     system::Query,
@@ -14,7 +14,8 @@ use crate::{
 use crate::entity::UniqueEntityEquivalentSlice;
 
 use alloc::vec::Vec;
-use core::{fmt, ptr};
+use core::{fmt, ops::DerefMut, ptr};
+use derive_more::derive::{Deref, DerefMut};
 use fixedbitset::FixedBitSet;
 use log::warn;
 #[cfg(feature = "trace")]
@@ -25,6 +26,12 @@ use super::{
     QueryManyUniqueIter, QuerySingleError, ROQueryItem, ReadOnlyQueryData,
 };
 
+#[repr(C)]
+#[derive(Component, Deref, DerefMut)]
+#[component(storage = "SparseSet", immutable)]
+/// A Internal Wrapper of [`QueryState`] for safety reasons.
+pub(crate) struct InternalQueryState<D: QueryData, F: QueryFilter>(QueryState<D, F>);
+
 /// An ID for either a table or an archetype. Used for Query iteration.
 ///
 /// Query iteration is exclusively dense (over tables) or archetypal (over archetypes) based on whether
@@ -1751,6 +1758,33 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
     }
 }
 
+impl<D: QueryData + 'static, F: QueryFilter + 'static> QueryState<D, F> {
+    /// cache a [`QueryState`] into world.
+    pub(crate) fn cached(self, world: &mut World) -> (Entity, ComponentId) {
+        let id = world.register_component::<InternalQueryState<D, F>>();
+        let e = world.spawn(InternalQueryState(self)).id();
+        (e, id)
+    }
+
+    /// fetch a cached [`QueryState`] from world
+    ///
+    /// Safety:
+    /// - Must not mutably alias the returned [`QueryState`].
+    pub(crate) unsafe fn fetch_mut_from_cached<'w>(
+        cached: (Entity, ComponentId),
+        w: UnsafeWorldCell<'w>,
+    ) -> Option<&'w mut QueryState<D, F>> {
+        w.storages()
+            .sparse_sets
+            .get(cached.1)?
+            .get(cached.0)
+            .map(|ptr| {
+                ptr.assert_unique()
+                    .deref_mut::<InternalQueryState<D, F>>()
+                    .deref_mut()
+            })
+    }
+}
 impl<D: QueryData, F: QueryFilter> From<QueryBuilder<'_, D, F>> for QueryState<D, F> {
     fn from(mut value: QueryBuilder<D, F>) -> Self {
         QueryState::from_builder(&mut value)

crates/bevy_ecs/src/system/builder.rs

@@ -3,6 +3,8 @@ use bevy_platform::cell::SyncCell;
 use variadics_please::all_tuples;
 
 use crate::{
+    component::ComponentId,
+    entity::Entity,
     prelude::QueryBuilder,
     query::{QueryData, QueryFilter, QueryState},
     resource::Resource,
@@ -212,10 +214,10 @@ impl ParamBuilder {
 unsafe impl<'w, 's, D: QueryData + 'static, F: QueryFilter + 'static>
     SystemParamBuilder<Query<'w, 's, D, F>> for QueryState<D, F>
 {
-    fn build(self, world: &mut World, system_meta: &mut SystemMeta) -> QueryState<D, F> {
+    fn build(self, world: &mut World, system_meta: &mut SystemMeta) -> (Entity, ComponentId) {
         self.validate_world(world.id());
         init_query_param(world, system_meta, &self);
-        self
+        self.cached(world)
     }
 }
 
@@ -291,12 +293,12 @@ unsafe impl<
         T: FnOnce(&mut QueryBuilder<D, F>),
     > SystemParamBuilder<Query<'w, 's, D, F>> for QueryParamBuilder<T>
 {
-    fn build(self, world: &mut World, system_meta: &mut SystemMeta) -> QueryState<D, F> {
+    fn build(self, world: &mut World, system_meta: &mut SystemMeta) -> (Entity, ComponentId) {
         let mut builder = QueryBuilder::new(world);
         (self.0)(&mut builder);
         let state = builder.build();
         init_query_param(world, system_meta, &state);
-        state
+        state.cached(world)
     }
 }
 
@@ -966,7 +968,7 @@ mod tests {
     #[derive(SystemParam)]
     #[system_param(builder)]
     struct CustomParam<'w, 's> {
-        query: Query<'w, 's, ()>,
+        query: Query<'w, 'w, ()>,
         local: Local<'s, usize>,
     }
 

crates/bevy_ecs/src/system/mod.rs

@@ -1329,10 +1329,10 @@ mod tests {
     #[test]
     fn system_state_spawned() {
         let mut world = World::default();
-        world.spawn_empty();
+        world.spawn(A);
         let spawn_tick = world.change_tick();
 
-        let mut system_state: SystemState<Option<Single<SpawnDetails, Spawned>>> =
+        let mut system_state: SystemState<Option<Single<SpawnDetails, (Spawned, With<A>)>>> =
             SystemState::new(&mut world);
         {
             let query = system_state.get(&world);

crates/bevy_ecs/src/system/system_param.rs

@@ -4,7 +4,7 @@ use crate::{
     bundle::Bundles,
     change_detection::{MaybeLocation, Ticks, TicksMut},
     component::{ComponentId, ComponentTicks, Components, Tick},
-    entity::Entities,
+    entity::{Entities, Entity},
     query::{
         Access, FilteredAccess, FilteredAccessSet, QueryData, QueryFilter, QuerySingleError,
         QueryState, ReadOnlyQueryData,
@@ -65,7 +65,7 @@ use variadics_please::{all_tuples, all_tuples_enumerated};
 /// # #[derive(SystemParam)]
 /// # struct ParamsExample<'w, 's> {
 /// #    query:
-/// Query<'w, 's, Entity>,
+/// Query<'w, 'w, Entity>,
 /// #    res:
 /// Res<'w, SomeResource>,
 /// #    res_mut:
@@ -171,7 +171,7 @@ use variadics_please::{all_tuples, all_tuples_enumerated};
 ///     #[derive(SystemParam)]
 ///     #[system_param(builder)]
 ///     pub struct CustomParam<'w, 's> {
-///         query: Query<'w, 's, ()>,
+///         query: Query<'w, 'w, ()>,
 ///         local: Local<'s, usize>,
 ///     }
 ///
@@ -321,13 +321,13 @@ unsafe impl<'w, 's, D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> Re
 // SAFETY: Relevant query ComponentId access is applied to SystemMeta. If
 // this Query conflicts with any prior access, a panic will occur.
 unsafe impl<D: QueryData + 'static, F: QueryFilter + 'static> SystemParam for Query<'_, '_, D, F> {
-    type State = QueryState<D, F>;
-    type Item<'w, 's> = Query<'w, 's, D, F>;
+    type State = (Entity, ComponentId);
+    type Item<'w, 's> = Query<'w, 'w, D, F>;
 
     fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self::State {
-        let state = QueryState::new(world);
+        let state: QueryState<D, F> = QueryState::new(world);
         init_query_param(world, system_meta, &state);
-        state
+        state.cached(world)
     }
 
     #[inline]
@@ -336,11 +336,14 @@ unsafe impl<D: QueryData + 'static, F: QueryFilter + 'static> SystemParam for Qu
         system_meta: &SystemMeta,
         world: UnsafeWorldCell<'w>,
         change_tick: Tick,
-    ) -> Self::Item<'w, 's> {
+    ) -> Self::Item<'w, 'w> {
+        let state = QueryState::fetch_mut_from_cached(*state, world)
+            .expect("cached querystate entity not found");
         // SAFETY: We have registered all of the query's world accesses,
         // so the caller ensures that `world` has permission to access any
         // world data that the query needs.
         // The caller ensures the world matches the one used in init_state.
+
         unsafe { state.query_unchecked_with_ticks(world, system_meta.last_run, change_tick) }
     }
 }
@@ -386,11 +389,11 @@ fn assert_component_access_compatibility(
 // SAFETY: Relevant query ComponentId access is applied to SystemMeta. If
 // this Query conflicts with any prior access, a panic will occur.
 unsafe impl<'a, D: QueryData + 'static, F: QueryFilter + 'static> SystemParam for Single<'a, D, F> {
-    type State = QueryState<D, F>;
+    type State = (Entity, ComponentId);
     type Item<'w, 's> = Single<'w, D, F>;
 
     fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self::State {
-        Query::init_state(world, system_meta)
+        Query::<D, F>::init_state(world, system_meta)
     }
 
     #[inline]
@@ -400,10 +403,8 @@ unsafe impl<'a, D: QueryData + 'static, F: QueryFilter + 'static> SystemParam fo
         world: UnsafeWorldCell<'w>,
         change_tick: Tick,
     ) -> Self::Item<'w, 's> {
-        // SAFETY: State ensures that the components it accesses are not accessible somewhere elsewhere.
-        // The caller ensures the world matches the one used in init_state.
-        let query =
-            unsafe { state.query_unchecked_with_ticks(world, system_meta.last_run, change_tick) };
+        let query: Query<'_, '_, D, F> = Query::get_param(state, system_meta, world, change_tick);
+
         let single = query
             .single_inner()
             .expect("The query was expected to contain exactly one matching entity.");
@@ -419,12 +420,9 @@ unsafe impl<'a, D: QueryData + 'static, F: QueryFilter + 'static> SystemParam fo
         system_meta: &SystemMeta,
         world: UnsafeWorldCell,
     ) -> Result<(), SystemParamValidationError> {
-        // SAFETY: State ensures that the components it accesses are not mutably accessible elsewhere
-        // and the query is read only.
-        // The caller ensures the world matches the one used in init_state.
-        let query = unsafe {
-            state.query_unchecked_with_ticks(world, system_meta.last_run, world.change_tick())
-        };
+        let query: Query<'_, '_, D, F> =
+            Query::get_param(state, system_meta, world, world.change_tick());
+
         match query.single_inner() {
             Ok(_) => Ok(()),
             Err(QuerySingleError::NoEntities(_)) => Err(
@@ -448,11 +446,11 @@ unsafe impl<'a, D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> ReadOn
 unsafe impl<D: QueryData + 'static, F: QueryFilter + 'static> SystemParam
     for Populated<'_, '_, D, F>
 {
-    type State = QueryState<D, F>;
-    type Item<'w, 's> = Populated<'w, 's, D, F>;
+    type State = (Entity, ComponentId);
+    type Item<'w, 's> = Populated<'w, 'w, D, F>;
 
     fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self::State {
-        Query::init_state(world, system_meta)
+        Query::<D, F>::init_state(world, system_meta)
     }
 
     #[inline]
@@ -461,7 +459,7 @@ unsafe impl<D: QueryData + 'static, F: QueryFilter + 'static> SystemParam
         system_meta: &SystemMeta,
         world: UnsafeWorldCell<'w>,
         change_tick: Tick,
-    ) -> Self::Item<'w, 's> {
+    ) -> Self::Item<'w, 'w> {
         // SAFETY: Delegate to existing `SystemParam` implementations.
         let query = unsafe { Query::get_param(state, system_meta, world, change_tick) };
         Populated(query)
@@ -473,12 +471,9 @@ unsafe impl<D: QueryData + 'static, F: QueryFilter + 'static> SystemParam
         system_meta: &SystemMeta,
         world: UnsafeWorldCell,
     ) -> Result<(), SystemParamValidationError> {
-        // SAFETY:
-        // - We have read-only access to the components accessed by query.
-        // - The caller ensures the world matches the one used in init_state.
-        let query = unsafe {
-            state.query_unchecked_with_ticks(world, system_meta.last_run, world.change_tick())
-        };
+        // SAFETY: Delegate to existing `SystemParam` implementations.
+        let query: Query<'_, '_, D, F> =
+            unsafe { Query::get_param(state, system_meta, world, world.change_tick()) };
         if query.is_empty() {
             Err(SystemParamValidationError::skipped::<Self>(
                 "No matching entities",
@@ -2577,11 +2572,10 @@ mod tests {
         #[derive(SystemParam)]
         pub struct SpecialQuery<
             'w,
-            's,
             D: QueryData + Send + Sync + 'static,
             F: QueryFilter + Send + Sync + 'static = (),
         > {
-            _query: Query<'w, 's, D, F>,
+            _query: Query<'w, 'w, D, F>,
         }
 
         fn my_system(_: SpecialQuery<(), ()>) {}
@@ -2710,11 +2704,11 @@ mod tests {
     #[test]
     fn system_param_where_clause() {
         #[derive(SystemParam)]
-        pub struct WhereParam<'w, 's, D>
+        pub struct WhereParam<'w, D>
         where
             D: 'static + QueryData,
         {
-            _q: Query<'w, 's, D, ()>,
+            _q: Query<'w, 'w, D, ()>,
         }
 
         fn my_system(_: WhereParam<()>) {}

crates/bevy_ecs/src/world/entity_ref.rs

@@ -5207,10 +5207,9 @@ mod tests {
 
         fn system(_: Query<&mut TestComponent>, query: Query<EntityRefExcept<TestComponent>>) {
             for entity_ref in query.iter() {
-                assert!(matches!(
-                    entity_ref.get::<TestComponent2>(),
-                    Some(TestComponent2(0))
-                ));
+                if let Some(c) = entity_ref.get::<TestComponent2>() {
+                    assert!(c.0 == 0);
+                }
             }
         }
     }
@@ -5230,10 +5229,9 @@ mod tests {
             assert!(entity_mut.get::<TestComponent>().is_none());
             assert!(entity_mut.get_ref::<TestComponent>().is_none());
             assert!(entity_mut.get_mut::<TestComponent>().is_none());
-            assert!(matches!(
-                entity_mut.get::<TestComponent2>(),
-                Some(TestComponent2(0))
-            ));
+            if let Some(c) = entity_mut.get::<TestComponent2>() {
+                assert!(c.0 == 0);
+            }
         }
 
         assert!(found);
@@ -5271,9 +5269,9 @@ mod tests {
 
         fn system(_: Query<&mut TestComponent>, mut query: Query<EntityMutExcept<TestComponent2>>) {
             for mut entity_mut in query.iter_mut() {
-                assert!(entity_mut
-                    .get_mut::<TestComponent2>()
-                    .is_some_and(|component| component.0 == 0));
+                if let Some(c) = entity_mut.get_mut::<TestComponent2>() {
+                    assert!(c.0 == 0);
+                }
             }
         }
     }
@@ -5289,9 +5287,9 @@ mod tests {
 
         fn system(_: Query<&mut TestComponent>, mut query: Query<EntityMutExcept<TestComponent>>) {
             for mut entity_mut in query.iter_mut() {
-                assert!(entity_mut
-                    .get_mut::<TestComponent2>()
-                    .is_some_and(|component| component.0 == 0));
+                if let Some(c) = entity_mut.get_mut::<TestComponent2>() {
+                    assert!(c.0 == 0);
+                }
             }
         }
     }

crates/bevy_input_focus/src/lib.rs

@@ -293,13 +293,13 @@ pub trait IsFocused {
 ///
 /// When working with the entire [`World`], consider using the [`IsFocused`] instead.
 #[derive(SystemParam)]
-pub struct IsFocusedHelper<'w, 's> {
-    parent_query: Query<'w, 's, &'static ChildOf>,
+pub struct IsFocusedHelper<'w> {
+    parent_query: Query<'w, 'w, &'static ChildOf>,
     input_focus: Option<Res<'w, InputFocus>>,
     input_focus_visible: Option<Res<'w, InputFocusVisible>>,
 }
 
-impl IsFocused for IsFocusedHelper<'_, '_> {
+impl IsFocused for IsFocusedHelper<'_> {
     fn is_focused(&self, entity: Entity) -> bool {
         self.input_focus
             .as_deref()