Make SystemParam::new_archetype and QueryState::new_archetype unsafe functions (#13044)

# Objective
Fix #2128. Both `Query::new_archetype` and `SystemParam::new_archetype`
do not check if the `Archetype` comes from the same World the state is
initialized from. This could result in unsoundness via invalid accesses
if called incorrectly.

## Solution
Make them `unsafe` functions and lift the invariant to the caller. This
also caught one instance of us not validating the World in
`SystemState::update_archetypes_unsafe_world_cell`'s implementation.

---

## Changelog
Changed: `QueryState::new_archetype` is now an unsafe function.
Changed: `SystemParam::new_archetype` is now an unsafe function.

## Migration Guide
`QueryState::new_archetype` and `SystemParam::new_archetype` are now an
unsafe functions that must be sure that the provided `Archetype` is from
the same `World` that the state was initialized from. Callers may need
to add additional assertions or propagate the safety invariant upwards
through the callstack to ensure safety.

Author: james7132

crates/bevy_ecs/macros/src/lib.rs

@@ -239,8 +239,9 @@ pub fn impl_param_set(_input: TokenStream) -> TokenStream {
                     (#(#param,)*)
                 }
 
-                fn new_archetype(state: &mut Self::State, archetype: &Archetype, system_meta: &mut SystemMeta) {
-                    <(#(#param,)*) as SystemParam>::new_archetype(state, archetype, system_meta);
+                unsafe fn new_archetype(state: &mut Self::State, archetype: &Archetype, system_meta: &mut SystemMeta) {
+                    // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from in `init_state`.
+                    unsafe { <(#(#param,)*) as SystemParam>::new_archetype(state, archetype, system_meta); }
                 }
 
                 fn apply(state: &mut Self::State, system_meta: &SystemMeta, world: &mut World) {
@@ -425,8 +426,9 @@ pub fn derive_system_param(input: TokenStream) -> TokenStream {
                     }
                 }
 
-                fn new_archetype(state: &mut Self::State, archetype: &#path::archetype::Archetype, system_meta: &mut #path::system::SystemMeta) {
-                    <#fields_alias::<'_, '_, #punctuated_generic_idents> as #path::system::SystemParam>::new_archetype(&mut state.state, archetype, system_meta)
+                unsafe fn new_archetype(state: &mut Self::State, archetype: &#path::archetype::Archetype, system_meta: &mut #path::system::SystemMeta) {
+                    // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from in `init_state`.
+                    unsafe { <#fields_alias::<'_, '_, #punctuated_generic_idents> as #path::system::SystemParam>::new_archetype(&mut state.state, archetype, system_meta) }
                 }
 
                 fn apply(state: &mut Self::State, system_meta: &#path::system::SystemMeta, world: &mut #path::world::World) {

crates/bevy_ecs/src/query/state.rs

@@ -161,8 +161,12 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
     ) -> Self {
         let mut state = Self::new_uninitialized(world);
         for archetype in world.archetypes.iter() {
-            if state.new_archetype_internal(archetype) {
-                state.update_archetype_component_access(archetype, access);
+            // SAFETY: The state was just initialized from the `world` above, and the archetypes being added
+            // come directly from the same world.
+            unsafe {
+                if state.new_archetype_internal(archetype) {
+                    state.update_archetype_component_access(archetype, access);
+                }
             }
         }
         state.archetype_generation = world.archetypes.generation();
@@ -342,7 +346,11 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
             std::mem::replace(&mut self.archetype_generation, archetypes.generation());
 
         for archetype in &archetypes[old_generation..] {
-            self.new_archetype_internal(archetype);
+            // SAFETY: The validate_world call ensures that the world is the same the QueryState
+            // was initialized from.
+            unsafe {
+                self.new_archetype_internal(archetype);
+            }
         }
     }
 
@@ -371,13 +379,19 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
     /// (if applicable, i.e. if the archetype has any intersecting [`ComponentId`] with the current [`QueryState`]).
     ///
     /// The passed in `access` will be updated with any new accesses introduced by the new archetype.
-    pub fn new_archetype(
+    ///
+    /// # Safety
+    /// `archetype` must be from the `World` this state was initialized from.
+    pub unsafe fn new_archetype(
         &mut self,
         archetype: &Archetype,
         access: &mut Access<ArchetypeComponentId>,
     ) {
-        if self.new_archetype_internal(archetype) {
-            self.update_archetype_component_access(archetype, access);
+        // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from.
+        let matches = unsafe { self.new_archetype_internal(archetype) };
+        if matches {
+            // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from.
+            unsafe { self.update_archetype_component_access(archetype, access) };
         }
     }
 
@@ -386,7 +400,10 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
     ///
     /// Returns `true` if the given `archetype` matches the query. Otherwise, returns `false`.
     /// If there is no match, then there is no need to update the query's [`FilteredAccess`].
-    fn new_archetype_internal(&mut self, archetype: &Archetype) -> bool {
+    ///
+    /// # Safety
+    /// `archetype` must be from the `World` this state was initialized from.
+    unsafe fn new_archetype_internal(&mut self, archetype: &Archetype) -> bool {
         if D::matches_component_set(&self.fetch_state, &|id| archetype.contains(id))
             && F::matches_component_set(&self.filter_state, &|id| archetype.contains(id))
             && self.matches_component_set(&|id| archetype.contains(id))
@@ -431,7 +448,10 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
     /// For the given `archetype`, adds any component accessed used by this query's underlying [`FilteredAccess`] to `access`.
     ///
     /// The passed in `access` will be updated with any new accesses introduced by the new archetype.
-    pub fn update_archetype_component_access(
+    ///
+    /// # Safety
+    /// `archetype` must be from the `World` this state was initialized from.
+    pub unsafe fn update_archetype_component_access(
         &mut self,
         archetype: &Archetype,
         access: &mut Access<ArchetypeComponentId>,

crates/bevy_ecs/src/system/function_system.rs

@@ -284,12 +284,15 @@ impl<Param: SystemParam> SystemState<Param> {
     /// This method only accesses world metadata.
     #[inline]
     pub fn update_archetypes_unsafe_world_cell(&mut self, world: UnsafeWorldCell) {
+        assert_eq!(self.world_id, world.id(), "Encountered a mismatched World. A System cannot be used with Worlds other than the one it was initialized with.");
+
         let archetypes = world.archetypes();
         let old_generation =
             std::mem::replace(&mut self.archetype_generation, archetypes.generation());
 
         for archetype in &archetypes[old_generation..] {
-            Param::new_archetype(&mut self.param_state, archetype, &mut self.meta);
+            // SAFETY: The assertion above ensures that the param_state was initialized from `world`.
+            unsafe { Param::new_archetype(&mut self.param_state, archetype, &mut self.meta) };
         }
     }
 
@@ -527,7 +530,8 @@ where
 
         for archetype in &archetypes[old_generation..] {
             let param_state = self.param_state.as_mut().unwrap();
-            F::Param::new_archetype(param_state, archetype, &mut self.system_meta);
+            // SAFETY: The assertion above ensures that the param_state was initialized from `world`.
+            unsafe { F::Param::new_archetype(param_state, archetype, &mut self.system_meta) };
         }
     }
 

crates/bevy_ecs/src/system/system_param.rs

@@ -137,12 +137,16 @@ pub unsafe trait SystemParam: Sized {
     /// and creates a new instance of this param's [`State`](Self::State).
     fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self::State;
 
-    /// For the specified [`Archetype`], registers the components accessed by this [`SystemParam`] (if applicable).
+    /// For the specified [`Archetype`], registers the components accessed by this [`SystemParam`] (if applicable).a
+    ///
+    /// # Safety
+    /// `archetype` must be from the [`World`] used to initialize `state` in `init_state`.
     #[inline]
-    fn new_archetype(
-        _state: &mut Self::State,
-        _archetype: &Archetype,
-        _system_meta: &mut SystemMeta,
+    #[allow(unused_variables)]
+    unsafe fn new_archetype(
+        state: &mut Self::State,
+        archetype: &Archetype,
+        system_meta: &mut SystemMeta,
     ) {
     }
 
@@ -208,7 +212,11 @@ unsafe impl<D: QueryData + 'static, F: QueryFilter + 'static> SystemParam for Qu
         state
     }
 
-    fn new_archetype(state: &mut Self::State, archetype: &Archetype, system_meta: &mut SystemMeta) {
+    unsafe fn new_archetype(
+        state: &mut Self::State,
+        archetype: &Archetype,
+        system_meta: &mut SystemMeta,
+    ) {
         state.new_archetype(archetype, &mut system_meta.archetype_component_access);
     }
 
@@ -1343,8 +1351,10 @@ macro_rules! impl_system_param_tuple {
             }
 
             #[inline]
-            fn new_archetype(($($param,)*): &mut Self::State, _archetype: &Archetype, _system_meta: &mut SystemMeta) {
-                $($param::new_archetype($param, _archetype, _system_meta);)*
+            #[allow(unused_unsafe)]
+            unsafe fn new_archetype(($($param,)*): &mut Self::State, _archetype: &Archetype, _system_meta: &mut SystemMeta) {
+                // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from in `init_state`.
+                unsafe { $($param::new_archetype($param, _archetype, _system_meta);)* }
             }
 
             #[inline]
@@ -1487,8 +1497,13 @@ unsafe impl<P: SystemParam + 'static> SystemParam for StaticSystemParam<'_, '_,
         P::init_state(world, system_meta)
     }
 
-    fn new_archetype(state: &mut Self::State, archetype: &Archetype, system_meta: &mut SystemMeta) {
-        P::new_archetype(state, archetype, system_meta);
+    unsafe fn new_archetype(
+        state: &mut Self::State,
+        archetype: &Archetype,
+        system_meta: &mut SystemMeta,
+    ) {
+        // SAFETY: The caller guarantees that the provided `archetype` matches the World used to initialize `state`.
+        unsafe { P::new_archetype(state, archetype, system_meta) };
     }
 
     fn apply(state: &mut Self::State, system_meta: &SystemMeta, world: &mut World) {

crates/bevy_gizmos/src/gizmos.rs

@@ -55,21 +55,26 @@ pub struct GizmosFetchState<T: GizmoConfigGroup> {
 unsafe impl<T: GizmoConfigGroup> SystemParam for Gizmos<'_, '_, T> {
     type State = GizmosFetchState<T>;
     type Item<'w, 's> = Gizmos<'w, 's, T>;
+
     fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self::State {
         GizmosFetchState {
             state: GizmosState::<T>::init_state(world, system_meta),
         }
     }
-    fn new_archetype(
+
+    unsafe fn new_archetype(
         state: &mut Self::State,
         archetype: &bevy_ecs::archetype::Archetype,
         system_meta: &mut SystemMeta,
     ) {
-        GizmosState::<T>::new_archetype(&mut state.state, archetype, system_meta);
+        // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from in `init_state`.
+        unsafe { GizmosState::<T>::new_archetype(&mut state.state, archetype, system_meta) };
     }
+
     fn apply(state: &mut Self::State, system_meta: &SystemMeta, world: &mut World) {
         GizmosState::<T>::apply(&mut state.state, system_meta, world);
     }
+
     unsafe fn get_param<'w, 's>(
         state: &'s mut Self::State,
         system_meta: &SystemMeta,