Rework DeferredEntity creation

I realized that it's possible to construct commands
from the entity's world. It lets me hide all unsafety
in the constructor and re-use it in other places.

src/client.rs

@@ -369,7 +369,8 @@ fn apply_removals(
         .entity_map
         .get_by_server_or_insert(server_entity, || world.spawn(Replicated).id());
 
-    let (mut client_entity, mut commands) = read_entity(world, params.queue, client_entity);
+    let mut client_entity = DeferredEntity::new(world, client_entity);
+    let mut commands = client_entity.commands(params.queue);
     params
         .entity_markers
         .read(params.command_markers, &*client_entity);
@@ -417,7 +418,8 @@ fn apply_changes(
         .entity_map
         .get_by_server_or_insert(server_entity, || world.spawn(Replicated).id());
 
-    let (mut client_entity, mut commands) = read_entity(world, params.queue, client_entity);
+    let mut client_entity = DeferredEntity::new(world, client_entity);
+    let mut commands = client_entity.commands(params.queue);
     params
         .entity_markers
         .read(params.command_markers, &*client_entity);
@@ -513,7 +515,8 @@ fn apply_mutations(
         return Ok(());
     };
 
-    let (mut client_entity, mut commands) = read_entity(world, params.queue, client_entity);
+    let mut client_entity = DeferredEntity::new(world, client_entity);
+    let mut commands = client_entity.commands(params.queue);
     params
         .entity_markers
         .read(params.command_markers, &*client_entity);
@@ -588,20 +591,6 @@ fn apply_mutations(
     Ok(())
 }
 
-/// Splits world access into entity that disallows structural ECS changes and commands.
-fn read_entity<'w, 's>(
-    world: &'w mut World,
-    queue: &'s mut CommandQueue,
-    client_entity: Entity,
-) -> (DeferredEntity<'w>, Commands<'w, 's>) {
-    let world_cell = world.as_unsafe_world_cell();
-    // SAFETY: have write access and the cell used only to get entities.
-    let client_entity = unsafe { DeferredEntity::new(world_cell, client_entity) };
-    let commands = Commands::new_from_entities(queue, world_cell.entities());
-
-    (client_entity, commands)
-}
-
 /// Deserializes `entity` from compressed index and generation.
 ///
 /// For details see

src/core/replication/deferred_entity.rs

@@ -1,4 +1,4 @@
-use bevy::{ecs::world::unsafe_world_cell::UnsafeWorldCell, prelude::*};
+use bevy::{ecs::world::CommandQueue, prelude::*};
 
 /// An entity reference that disallows structural ECS changes.
 ///
@@ -11,18 +11,19 @@ pub struct DeferredEntity<'w> {
 }
 
 impl<'w> DeferredEntity<'w> {
-    /// Creates a new instance from a world cell.
-    ///
-    /// # Safety
-    ///
-    /// - The cell must have been created using [`World::as_unsafe_world_cell`].
-    /// - No structural ECS changes can be done using the cell.
-    /// - No other mutable references to the entity's components should exist.
-    pub(crate) unsafe fn new(world_cell: UnsafeWorldCell<'w>, entity: Entity) -> Self {
-        // Split access, `EntityMut` can't make structural changes and they share the lifetime.
-        let entity: EntityMut = world_cell.world_mut().entity_mut(entity).into();
-        let world = world_cell.world();
-        Self { entity, world }
+    pub(crate) fn new(world: &'w mut World, entity: Entity) -> Self {
+        let world_cell = world.as_unsafe_world_cell();
+        // SAFETY: access split, `EntityMut` cannot make structural ECS changes,
+        // and the world cannot be accessed simultaneously with the entity.
+        unsafe {
+            let entity: EntityMut = world_cell.world_mut().entity_mut(entity).into();
+            let world = world_cell.world();
+            Self { entity, world }
+        }
+    }
+
+    pub(crate) fn commands<'s>(&self, queue: &'s mut CommandQueue) -> Commands<'w, 's> {
+        Commands::new_from_entities(queue, self.world.entities())
     }
 
     /// Gets read-only access to the world that the current entity belongs to.

src/core/replication/replication_registry/test_fns.rs

@@ -123,12 +123,9 @@ impl TestFnsEntityExt for EntityWorldMut<'_> {
         self.world_scope(|world| {
             world.resource_scope(|world, mut entity_map: Mut<ServerEntityMap>| {
                 world.resource_scope(|world, registry: Mut<ReplicationRegistry>| {
-                    let world_cell = world.as_unsafe_world_cell();
-                    // SAFETY: have write access and the cell used only to get entities.
-                    let mut entity = unsafe { DeferredEntity::new(world_cell, entity) };
                     let mut queue = CommandQueue::default();
-                    let mut commands =
-                        Commands::new_from_entities(&mut queue, world_cell.entities());
+                    let mut entity = DeferredEntity::new(world, entity);
+                    let mut commands = entity.commands(&mut queue);
 
                     let (component_id, component_fns, rule_fns) = registry.get(fns_id);
                     let mut cursor = Cursor::new(data);
@@ -163,11 +160,9 @@ impl TestFnsEntityExt for EntityWorldMut<'_> {
         let entity = self.id();
         self.world_scope(|world| {
             world.resource_scope(|world, registry: Mut<ReplicationRegistry>| {
-                let world_cell = world.as_unsafe_world_cell();
-                // SAFETY: have write access and the cell used only to get entities.
-                let mut entity = unsafe { DeferredEntity::new(world_cell, entity) };
                 let mut queue = CommandQueue::default();
-                let mut commands = Commands::new_from_entities(&mut queue, world_cell.entities());
+                let mut entity = DeferredEntity::new(world, entity);
+                let mut commands = entity.commands(&mut queue);
 
                 let (component_id, component_fns, _) = registry.get(fns_id);
                 let mut ctx = RemoveCtx {