constrain WorldQuery::init_state argument to ComponentInitializer (#13442)

# Objective

In #13343, `WorldQuery::get_state` was constrained from `&World` as the
argument to `&Components`, but `WorldQuery::init_state` hasn't yet been
changed from `&mut World` to match.

Fixes #13358

## Solution

Create a wrapper around `&mut Components` and `&mut Storages` that can
be obtained from `&mut World` with a `component_initializer` method.
This new `ComponentInitializer` re-exposes the API on `&mut Components`
minus the `&mut Storages` parameter where it was present. For the
`&Components` API, it simply derefs to its `components` field.

## Changelog

### Added
The `World::component_initializer` method.
The `ComponentInitializer` struct that re-exposes `Components` API.
### Changed
`WorldQuery::init_state` now takes `&mut ComponentInitializer` instead
of `&mut World`.

## Migration Guide
Instead of passing `&mut World` to `WorldQuery::init_state` directly,
pass in a mutable reference to the struct returned from
`World::component_initializer`.

crates/bevy_ecs/macros/src/world_query.rs

@@ -164,9 +164,9 @@ pub(crate) fn world_query_impl(
                 #( <#field_types>::update_component_access(&state.#named_field_idents, _access); )*
             }
 
-            fn init_state(world: &mut #path::world::World) -> #state_struct_name #user_ty_generics {
+            fn init_state(initializer: &mut #path::component::ComponentInitializer) -> #state_struct_name #user_ty_generics {
                 #state_struct_name {
-                    #(#named_field_idents: <#field_types>::init_state(world),)*
+                    #(#named_field_idents: <#field_types>::init_state(initializer),)*
                 }
             }
 

crates/bevy_ecs/src/component.rs

@@ -21,6 +21,7 @@ use std::{
     borrow::Cow,
     marker::PhantomData,
     mem::needs_drop,
+    ops::Deref,
 };
 
 /// A data type that can be used to store data for an [entity].
@@ -832,6 +833,75 @@ impl Components {
     }
 }
 
+/// A wrapper over a mutable [`Components`] reference that allows for state initialization.
+/// This can be obtained with [`World::component_initializer`].
+pub struct ComponentInitializer<'w> {
+    pub(crate) components: &'w mut Components,
+    pub(crate) storages: &'w mut Storages,
+}
+
+impl<'w> Deref for ComponentInitializer<'w> {
+    type Target = Components;
+
+    fn deref(&self) -> &Components {
+        self.components
+    }
+}
+
+impl<'w> ComponentInitializer<'w> {
+    /// Initializes a component of type `T` with this instance.
+    /// If a component of this type has already been initialized, this will return
+    /// the ID of the pre-existing component.
+    ///
+    /// # See also
+    ///
+    /// * [`Components::component_id()`]
+    /// * [`Components::init_component_with_descriptor()`]
+    #[inline]
+    pub fn init_component<T: Component>(&mut self) -> ComponentId {
+        self.components.init_component::<T>(self.storages)
+    }
+
+    /// Initializes a component described by `descriptor`.
+    ///
+    /// ## Note
+    ///
+    /// If this method is called multiple times with identical descriptors, a distinct `ComponentId`
+    /// will be created for each one.
+    ///
+    /// # See also
+    ///
+    /// * [`Components::component_id()`]
+    /// * [`Components::init_component()`]
+    pub fn init_component_with_descriptor(
+        &mut self,
+        descriptor: ComponentDescriptor,
+    ) -> ComponentId {
+        self.components
+            .init_component_with_descriptor(self.storages, descriptor)
+    }
+
+    /// Initializes a [`Resource`] of type `T` with this instance.
+    /// If a resource of this type has already been initialized, this will return
+    /// the ID of the pre-existing resource.
+    ///
+    /// # See also
+    ///
+    /// * [`Components::resource_id()`]
+    #[inline]
+    pub fn init_resource<T: Resource>(&mut self) -> ComponentId {
+        self.components.init_resource::<T>()
+    }
+
+    /// Initializes a [non-send resource](crate::system::NonSend) of type `T` with this instance.
+    /// If a resource of this type has already been initialized, this will return
+    /// the ID of the pre-existing resource.
+    #[inline]
+    pub fn init_non_send<T: Any>(&mut self) -> ComponentId {
+        self.components.init_non_send::<T>()
+    }
+}
+
 /// A value that tracks when a system ran relative to other systems.
 /// This is used to power change detection.
 ///

crates/bevy_ecs/src/query/builder.rs

@@ -43,8 +43,9 @@ pub struct QueryBuilder<'w, D: QueryData = (), F: QueryFilter = ()> {
 impl<'w, D: QueryData, F: QueryFilter> QueryBuilder<'w, D, F> {
     /// Creates a new builder with the accesses required for `Q` and `F`
     pub fn new(world: &'w mut World) -> Self {
-        let fetch_state = D::init_state(world);
+        let initializer = &mut world.component_initializer();
-        let filter_state = F::init_state(world);
+        let fetch_state = D::init_state(initializer);
+        let filter_state = F::init_state(initializer);
 
         let mut access = FilteredAccess::default();
         D::update_component_access(&fetch_state, &mut access);
@@ -95,7 +96,7 @@ impl<'w, D: QueryData, F: QueryFilter> QueryBuilder<'w, D, F> {
 
     /// Adds accesses required for `T` to self.
     pub fn data<T: QueryData>(&mut self) -> &mut Self {
-        let state = T::init_state(self.world);
+        let state = T::init_state(&mut self.world.component_initializer());
         let mut access = FilteredAccess::default();
         T::update_component_access(&state, &mut access);
         self.extend_access(access);
@@ -104,7 +105,7 @@ impl<'w, D: QueryData, F: QueryFilter> QueryBuilder<'w, D, F> {
 
     /// Adds filter from `T` to self.
     pub fn filter<T: QueryFilter>(&mut self) -> &mut Self {
-        let state = T::init_state(self.world);
+        let state = T::init_state(&mut self.world.component_initializer());
         let mut access = FilteredAccess::default();
         T::update_component_access(&state, &mut access);
         self.extend_access(access);
@@ -222,8 +223,9 @@ impl<'w, D: QueryData, F: QueryFilter> QueryBuilder<'w, D, F> {
     pub fn transmute_filtered<NewD: QueryData, NewF: QueryFilter>(
         &mut self,
     ) -> &mut QueryBuilder<'w, NewD, NewF> {
-        let mut fetch_state = NewD::init_state(self.world);
+        let initializer = &mut self.world.component_initializer();
-        let filter_state = NewF::init_state(self.world);
+        let mut fetch_state = NewD::init_state(initializer);
+        let filter_state = NewF::init_state(initializer);
 
         NewD::set_access(&mut fetch_state, &self.access);
 

crates/bevy_ecs/src/query/fetch.rs

@@ -1,20 +1,20 @@
 use crate::{
     archetype::{Archetype, Archetypes},
     change_detection::{Ticks, TicksMut},
-    component::{Component, ComponentId, Components, StorageType, Tick},
+    component::{Component, ComponentId, ComponentInitializer, Components, StorageType, Tick},
     entity::{Entities, Entity, EntityLocation},
     query::{Access, DebugCheckedUnwrap, FilteredAccess, WorldQuery},
     storage::{ComponentSparseSet, Table, TableRow},
     world::{
         unsafe_world_cell::UnsafeWorldCell, EntityMut, EntityRef, FilteredEntityMut,
-        FilteredEntityRef, Mut, Ref, World,
+        FilteredEntityRef, Mut, Ref,
     },
 };
 use bevy_ptr::{ThinSlicePtr, UnsafeCellDeref};
 use bevy_utils::all_tuples;
 use std::{cell::UnsafeCell, marker::PhantomData};
 
-/// Types that can be fetched from a [`World`] using a [`Query`].
+/// Types that can be fetched from a [`World`](crate::world::World) using a [`Query`].
 ///
 /// There are many types that natively implement this trait:
 ///
@@ -331,7 +331,7 @@ unsafe impl WorldQuery for Entity {
 
     fn update_component_access(_state: &Self::State, _access: &mut FilteredAccess<ComponentId>) {}
 
-    fn init_state(_world: &mut World) {}
+    fn init_state(_initializer: &mut ComponentInitializer) {}
 
     fn get_state(_components: &Components) -> Option<()> {
         Some(())
@@ -403,7 +403,7 @@ unsafe impl WorldQuery for EntityLocation {
 
     fn update_component_access(_state: &Self::State, _access: &mut FilteredAccess<ComponentId>) {}
 
-    fn init_state(_world: &mut World) {}
+    fn init_state(_initializer: &mut ComponentInitializer) {}
 
     fn get_state(_components: &Components) -> Option<()> {
         Some(())
@@ -482,7 +482,7 @@ unsafe impl<'a> WorldQuery for EntityRef<'a> {
         access.read_all();
     }
 
-    fn init_state(_world: &mut World) {}
+    fn init_state(_initializer: &mut ComponentInitializer) {}
 
     fn get_state(_components: &Components) -> Option<()> {
         Some(())
@@ -558,7 +558,7 @@ unsafe impl<'a> WorldQuery for EntityMut<'a> {
         access.write_all();
     }
 
-    fn init_state(_world: &mut World) {}
+    fn init_state(_initializer: &mut ComponentInitializer) {}
 
     fn get_state(_components: &Components) -> Option<()> {
         Some(())
@@ -656,7 +656,7 @@ unsafe impl<'a> WorldQuery for FilteredEntityRef<'a> {
         filtered_access.access.extend(&state.access);
     }
 
-    fn init_state(_world: &mut World) -> Self::State {
+    fn init_state(_initializer: &mut ComponentInitializer) -> Self::State {
         FilteredAccess::default()
     }
 
@@ -768,7 +768,7 @@ unsafe impl<'a> WorldQuery for FilteredEntityMut<'a> {
         filtered_access.access.extend(&state.access);
     }
 
-    fn init_state(_world: &mut World) -> Self::State {
+    fn init_state(_initializer: &mut ComponentInitializer) -> Self::State {
         FilteredAccess::default()
     }
 
@@ -842,7 +842,7 @@ unsafe impl WorldQuery for &Archetype {
 
     fn update_component_access(_state: &Self::State, _access: &mut FilteredAccess<ComponentId>) {}
 
-    fn init_state(_world: &mut World) {}
+    fn init_state(_initializer: &mut ComponentInitializer) {}
 
     fn get_state(_components: &Components) -> Option<()> {
         Some(())
@@ -991,8 +991,8 @@ unsafe impl<T: Component> WorldQuery for &T {
         access.add_read(component_id);
     }
 
-    fn init_state(world: &mut World) -> ComponentId {
+    fn init_state(initializer: &mut ComponentInitializer) -> ComponentId {
-        world.init_component::<T>()
+        initializer.init_component::<T>()
     }
 
     fn get_state(components: &Components) -> Option<Self::State> {
@@ -1174,8 +1174,8 @@ unsafe impl<'__w, T: Component> WorldQuery for Ref<'__w, T> {
         access.add_read(component_id);
     }
 
-    fn init_state(world: &mut World) -> ComponentId {
+    fn init_state(initializer: &mut ComponentInitializer<'_>) -> ComponentId {
-        world.init_component::<T>()
+        initializer.init_component::<T>()
     }
 
     fn get_state(components: &Components) -> Option<Self::State> {
@@ -1357,8 +1357,8 @@ unsafe impl<'__w, T: Component> WorldQuery for &'__w mut T {
         access.add_write(component_id);
     }
 
-    fn init_state(world: &mut World) -> ComponentId {
+    fn init_state(initializer: &mut ComponentInitializer<'_>) -> ComponentId {
-        world.init_component::<T>()
+        initializer.init_component::<T>()
     }
 
     fn get_state(components: &Components) -> Option<Self::State> {
@@ -1456,8 +1456,8 @@ unsafe impl<'__w, T: Component> WorldQuery for Mut<'__w, T> {
     }
 
     // Forwarded to `&mut T`
-    fn init_state(world: &mut World) -> ComponentId {
+    fn init_state(initializer: &mut ComponentInitializer) -> ComponentId {
-        <&mut T as WorldQuery>::init_state(world)
+        <&mut T as WorldQuery>::init_state(initializer)
     }
 
     // Forwarded to `&mut T`
@@ -1577,8 +1577,8 @@ unsafe impl<T: WorldQuery> WorldQuery for Option<T> {
         access.extend_access(&intermediate);
     }
 
-    fn init_state(world: &mut World) -> T::State {
+    fn init_state(initializer: &mut ComponentInitializer) -> T::State {
-        T::init_state(world)
+        T::init_state(initializer)
     }
 
     fn get_state(components: &Components) -> Option<Self::State> {
@@ -1732,8 +1732,8 @@ unsafe impl<T: Component> WorldQuery for Has<T> {
         access.access_mut().add_archetypal(component_id);
     }
 
-    fn init_state(world: &mut World) -> ComponentId {
+    fn init_state(initializer: &mut ComponentInitializer) -> ComponentId {
-        world.init_component::<T>()
+        initializer.init_component::<T>()
     }
 
     fn get_state(components: &Components) -> Option<Self::State> {
@@ -1878,8 +1878,8 @@ macro_rules! impl_anytuple_fetch {
                 *_access = _new_access;
             }
             #[allow(unused_variables)]
-            fn init_state(world: &mut World) -> Self::State {
+            fn init_state(initializer: &mut ComponentInitializer) -> Self::State {
-                ($($name::init_state(world),)*)
+                ($($name::init_state(initializer),)*)
             }
             #[allow(unused_variables)]
             fn get_state(components: &Components) -> Option<Self::State> {
@@ -1955,8 +1955,8 @@ unsafe impl<D: QueryData> WorldQuery for NopWorldQuery<D> {
 
     fn update_component_access(_state: &D::State, _access: &mut FilteredAccess<ComponentId>) {}
 
-    fn init_state(world: &mut World) -> Self::State {
+    fn init_state(initializer: &mut ComponentInitializer) -> Self::State {
-        D::init_state(world)
+        D::init_state(initializer)
     }
 
     fn get_state(components: &Components) -> Option<Self::State> {
@@ -2022,7 +2022,7 @@ unsafe impl<T: ?Sized> WorldQuery for PhantomData<T> {
 
     fn update_component_access(_state: &Self::State, _access: &mut FilteredAccess<ComponentId>) {}
 
-    fn init_state(_world: &mut World) -> Self::State {}
+    fn init_state(_initializer: &mut ComponentInitializer) -> Self::State {}
 
     fn get_state(_components: &Components) -> Option<Self::State> {
         Some(())

crates/bevy_ecs/src/query/filter.rs

@@ -1,10 +1,10 @@
 use crate::{
     archetype::Archetype,
-    component::{Component, ComponentId, Components, StorageType, Tick},
+    component::{Component, ComponentId, ComponentInitializer, Components, StorageType, Tick},
     entity::Entity,
     query::{DebugCheckedUnwrap, FilteredAccess, WorldQuery},
     storage::{Column, ComponentSparseSet, Table, TableRow},
-    world::{unsafe_world_cell::UnsafeWorldCell, World},
+    world::unsafe_world_cell::UnsafeWorldCell,
 };
 use bevy_ptr::{ThinSlicePtr, UnsafeCellDeref};
 use bevy_utils::all_tuples;
@@ -179,8 +179,8 @@ unsafe impl<T: Component> WorldQuery for With<T> {
         access.and_with(id);
     }
 
-    fn init_state(world: &mut World) -> ComponentId {
+    fn init_state(initializer: &mut ComponentInitializer) -> ComponentId {
-        world.init_component::<T>()
+        initializer.init_component::<T>()
     }
 
     fn get_state(components: &Components) -> Option<Self::State> {
@@ -287,8 +287,8 @@ unsafe impl<T: Component> WorldQuery for Without<T> {
         access.and_without(id);
     }
 
-    fn init_state(world: &mut World) -> ComponentId {
+    fn init_state(initializer: &mut ComponentInitializer) -> ComponentId {
-        world.init_component::<T>()
+        initializer.init_component::<T>()
     }
 
     fn get_state(components: &Components) -> Option<Self::State> {
@@ -457,8 +457,8 @@ macro_rules! impl_or_query_filter {
                 *access = _new_access;
             }
 
-            fn init_state(world: &mut World) -> Self::State {
+            fn init_state(initializer: &mut ComponentInitializer) -> Self::State {
-                ($($filter::init_state(world),)*)
+                ($($filter::init_state(initializer),)*)
             }
 
             fn get_state(components: &Components) -> Option<Self::State> {
@@ -689,8 +689,8 @@ unsafe impl<T: Component> WorldQuery for Added<T> {
         access.add_read(id);
     }
 
-    fn init_state(world: &mut World) -> ComponentId {
+    fn init_state(initializer: &mut ComponentInitializer) -> ComponentId {
-        world.init_component::<T>()
+        initializer.init_component::<T>()
     }
 
     fn get_state(components: &Components) -> Option<ComponentId> {
@@ -900,8 +900,8 @@ unsafe impl<T: Component> WorldQuery for Changed<T> {
         access.add_read(id);
     }
 
-    fn init_state(world: &mut World) -> ComponentId {
+    fn init_state(initializer: &mut ComponentInitializer) -> ComponentId {
-        world.init_component::<T>()
+        initializer.init_component::<T>()
     }
 
     fn get_state(components: &Components) -> Option<ComponentId> {

crates/bevy_ecs/src/query/state.rs

@@ -178,8 +178,9 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
     /// `new_archetype` and its variants must be called on all of the World's archetypes before the
     /// state can return valid query results.
     fn new_uninitialized(world: &mut World) -> Self {
-        let fetch_state = D::init_state(world);
+        let initializer = &mut world.component_initializer();
-        let filter_state = F::init_state(world);
+        let fetch_state = D::init_state(initializer);
+        let filter_state = F::init_state(initializer);
 
         let mut component_access = FilteredAccess::default();
         D::update_component_access(&fetch_state, &mut component_access);
@@ -214,8 +215,9 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
 
     /// Creates a new [`QueryState`] from a given [`QueryBuilder`] and inherits its [`FilteredAccess`].
     pub fn from_builder(builder: &mut QueryBuilder<D, F>) -> Self {
-        let mut fetch_state = D::init_state(builder.world_mut());
+        let initializer = &mut builder.world_mut().component_initializer();
-        let filter_state = F::init_state(builder.world_mut());
+        let mut fetch_state = D::init_state(initializer);
+        let filter_state = F::init_state(initializer);
         D::set_access(&mut fetch_state, builder.access());
 
         let mut state = Self {

crates/bevy_ecs/src/query/world_query.rs

@@ -1,10 +1,10 @@
 use crate::{
     archetype::Archetype,
-    component::{ComponentId, Components, Tick},
+    component::{ComponentId, ComponentInitializer, Components, Tick},
     entity::Entity,
     query::FilteredAccess,
     storage::{Table, TableRow},
-    world::{unsafe_world_cell::UnsafeWorldCell, World},
+    world::unsafe_world_cell::UnsafeWorldCell,
 };
 use bevy_utils::all_tuples;
 
@@ -79,7 +79,7 @@ pub unsafe trait WorldQuery {
     ///
     /// # Safety
     ///
-    /// - `archetype` and `tables` must be from the same [`World`] that [`WorldQuery::init_state`] was called on.
+    /// - `archetype` and `tables` must be from the same [`World`](crate::world::World) that [`WorldQuery::init_state`] was called on.
     /// - `table` must correspond to `archetype`.
     /// - `state` must be the [`State`](Self::State) that `fetch` was initialized with.
     unsafe fn set_archetype<'w>(
@@ -94,7 +94,7 @@ pub unsafe trait WorldQuery {
     ///
     /// # Safety
     ///
-    /// - `table` must be from the same [`World`] that [`WorldQuery::init_state`] was called on.
+    /// - `table` must be from the same [`World`](crate::world::World) that [`WorldQuery::init_state`] was called on.
     /// - `state` must be the [`State`](Self::State) that `fetch` was initialized with.
     unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, table: &'w Table);
 
@@ -127,7 +127,7 @@ pub unsafe trait WorldQuery {
     fn update_component_access(state: &Self::State, access: &mut FilteredAccess<ComponentId>);
 
     /// Creates and initializes a [`State`](WorldQuery::State) for this [`WorldQuery`] type.
-    fn init_state(world: &mut World) -> Self::State;
+    fn init_state(initializer: &mut ComponentInitializer) -> Self::State;
 
     /// Attempts to initialize a [`State`](WorldQuery::State) for this [`WorldQuery`] type using read-only
     /// access to [`Components`].
@@ -213,8 +213,8 @@ macro_rules! impl_tuple_world_query {
                 $($name::update_component_access($name, _access);)*
             }
             #[allow(unused_variables)]
-            fn init_state(world: &mut World) -> Self::State {
+            fn init_state(initializer: &mut ComponentInitializer) -> Self::State {
-                ($($name::init_state(world),)*)
+                ($($name::init_state(initializer),)*)
             }
             #[allow(unused_variables)]
             fn get_state(components: &Components) -> Option<Self::State> {

crates/bevy_ecs/src/world/mod.rs

@@ -8,6 +8,7 @@ mod spawn_batch;
 pub mod unsafe_world_cell;
 
 pub use crate::change_detection::{Mut, Ref, CHECK_TICK_THRESHOLD};
+use crate::component::ComponentInitializer;
 pub use crate::world::command_queue::CommandQueue;
 pub use deferred_world::DeferredWorld;
 pub use entity_ref::{
@@ -218,6 +219,15 @@ impl World {
         &self.bundles
     }
 
+    /// Creates a [`ComponentInitializer`] for this world.
+    #[inline]
+    pub fn component_initializer(&mut self) -> ComponentInitializer {
+        ComponentInitializer {
+            components: &mut self.components,
+            storages: &mut self.storages,
+        }
+    }
+
     /// Retrieves this world's [`RemovedComponentEvents`] collection
     #[inline]
     pub fn removed_components(&self) -> &RemovedComponentEvents {