cleanup

Author: re0312

crates/bevy_ecs/src/query/state.rs

@@ -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")]
+/// 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
@@ -60,8 +67,6 @@ pub(super) union StorageId {
 /// [`Fetch`]: crate::query::world_query::WorldQuery::Fetch
 /// [`Table`]: crate::storage::Table
 #[repr(C)]
-#[derive(Component)]
-#[component(storage = "SparseSet")]
 // SAFETY NOTE:
 // Do not add any new fields that use the `D` or `F` generic parameters as this may
 // make `QueryState::as_transmuted_state` unsound if not done with care.
@@ -1753,6 +1758,31 @@ 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
+    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)
+            .map(|s| s.get(cached.0))
+            .flatten()
+            .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

@@ -217,9 +217,7 @@ unsafe impl<'w, 's, D: QueryData + 'static, F: QueryFilter + 'static>
     fn build(self, world: &mut World, system_meta: &mut SystemMeta) -> (Entity, ComponentId) {
         self.validate_world(world.id());
         init_query_param(world, system_meta, &self);
-        let id = world.register_component::<QueryState<D, F>>();
-        let e = world.spawn(self).id();
-        (e, id)
+        self.cached(world)
     }
 }
 
@@ -300,9 +298,7 @@ unsafe impl<
         (self.0)(&mut builder);
         let state = builder.build();
         init_query_param(world, system_meta, &state);
-        let id = world.register_component::<QueryState<D, F>>();
-        let e = world.spawn(state).id();
-        (e, id)
+        state.cached(world)
     }
 }
 
@@ -971,9 +967,9 @@ mod tests {
 
     #[derive(SystemParam)]
     #[system_param(builder)]
-    struct CustomParam<'w> {
+    struct CustomParam<'w, 's> {
         query: Query<'w, 'w, ()>,
-        local: Local<usize>,
+        local: Local<'s, usize>,
     }
 
     #[test]

crates/bevy_ecs/src/system/system_param.rs

@@ -6,8 +6,8 @@ use crate::{
     component::{ComponentId, ComponentTicks, Components, Tick},
     entity::{Entities, Entity},
     query::{
-        Access, DebugCheckedUnwrap, FilteredAccess, FilteredAccessSet, QueryData, QueryFilter,
-        QuerySingleError, QueryState, ReadOnlyQueryData,
+        Access, FilteredAccess, FilteredAccessSet, QueryData, QueryFilter, QuerySingleError,
+        QueryState, ReadOnlyQueryData,
     },
     resource::Resource,
     storage::ResourceData,
@@ -327,9 +327,7 @@ unsafe impl<D: QueryData + 'static, F: QueryFilter + 'static> SystemParam for Qu
     fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self::State {
         let state: QueryState<D, F> = QueryState::new(world);
         init_query_param(world, system_meta, &state);
-        let e = world.spawn(state).id();
-        let id = world.register_component::<QueryState<D, F>>();
-        (e, id)
+        state.cached(world)
     }
 
     #[inline]
@@ -339,15 +337,8 @@ unsafe impl<D: QueryData + 'static, F: QueryFilter + 'static> SystemParam for Qu
         world: UnsafeWorldCell<'w>,
         change_tick: Tick,
     ) -> Self::Item<'w, 'w> {
-        let state = world
-            .storages()
-            .sparse_sets
-            .get(state.1)
-            .debug_checked_unwrap()
-            .get(state.0)
-            .debug_checked_unwrap()
-            .assert_unique()
-            .deref_mut::<QueryState<D, F>>();
+        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.