Provide a safe abstraction for split access to entities and commands (#18215)

# Objective

Currently experimenting with manually implementing
`Relationship`/`RelationshipTarget` to support associated edge data,
which means I need to replace the default hook implementations provided
by those traits. However, copying them over for editing revealed that
`UnsafeWorldCell::get_raw_command_queue` is `pub(crate)`, and I would
like to not have to clone the source collection, like the default impl.
So instead, I've taken to providing a safe abstraction for being able to
access entities and queue commands simultaneously.

## Solution

Added `World::entities_and_commands` and
`DeferredWorld::entities_and_commands`, which can be used like so:

```rust
let eid: Entity = /* ... */;
let (mut fetcher, mut commands) = world.entities_and_commands();
let emut = fetcher.get_mut(eid).unwrap();
commands.entity(eid).despawn();
```

## Testing

- Added a new test for each of the added functions.

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>

Author: ItsDoot

crates/bevy_ecs/src/relationship/mod.rs

@@ -207,29 +207,23 @@ pub trait RelationshipTarget: Component<Mutability = Mutable> + Sized {
     /// The `on_replace` component hook that maintains the [`Relationship`] / [`RelationshipTarget`] connection.
     // note: think of this as "on_drop"
     fn on_replace(mut world: DeferredWorld, HookContext { entity, caller, .. }: HookContext) {
-        // NOTE: this unsafe code is an optimization. We could make this safe, but it would require
-        // copying the RelationshipTarget collection
-        // SAFETY: This only reads the Self component and queues Remove commands
-        unsafe {
-            let world = world.as_unsafe_world_cell();
-            let relationship_target = world.get_entity(entity).unwrap().get::<Self>().unwrap();
-            let mut commands = world.get_raw_command_queue();
-            for source_entity in relationship_target.iter() {
-                if world.get_entity(source_entity).is_ok() {
-                    commands.push(
-                        entity_command::remove::<Self::Relationship>()
-                            .with_entity(source_entity)
-                            .handle_error_with(error_handler::silent()),
-                    );
-                } else {
-                    warn!(
-                        "{}Tried to despawn non-existent entity {}",
-                        caller
-                            .map(|location| format!("{location}: "))
-                            .unwrap_or_default(),
-                        source_entity
-                    );
-                }
+        let (entities, mut commands) = world.entities_and_commands();
+        let relationship_target = entities.get(entity).unwrap().get::<Self>().unwrap();
+        for source_entity in relationship_target.iter() {
+            if entities.get(source_entity).is_ok() {
+                commands.queue(
+                    entity_command::remove::<Self::Relationship>()
+                        .with_entity(source_entity)
+                        .handle_error_with(error_handler::silent()),
+                );
+            } else {
+                warn!(
+                    "{}Tried to despawn non-existent entity {}",
+                    caller
+                        .map(|location| format!("{location}: "))
+                        .unwrap_or_default(),
+                    source_entity
+                );
             }
         }
     }
@@ -238,29 +232,23 @@ pub trait RelationshipTarget: Component<Mutability = Mutable> + Sized {
     /// that entity is despawned.
     // note: think of this as "on_drop"
     fn on_despawn(mut world: DeferredWorld, HookContext { entity, caller, .. }: HookContext) {
-        // NOTE: this unsafe code is an optimization. We could make this safe, but it would require
-        // copying the RelationshipTarget collection
-        // SAFETY: This only reads the Self component and queues despawn commands
-        unsafe {
-            let world = world.as_unsafe_world_cell();
-            let relationship_target = world.get_entity(entity).unwrap().get::<Self>().unwrap();
-            let mut commands = world.get_raw_command_queue();
-            for source_entity in relationship_target.iter() {
-                if world.get_entity(source_entity).is_ok() {
-                    commands.push(
-                        entity_command::despawn()
-                            .with_entity(source_entity)
-                            .handle_error_with(error_handler::silent()),
-                    );
-                } else {
-                    warn!(
-                        "{}Tried to despawn non-existent entity {}",
-                        caller
-                            .map(|location| format!("{location}: "))
-                            .unwrap_or_default(),
-                        source_entity
-                    );
-                }
+        let (entities, mut commands) = world.entities_and_commands();
+        let relationship_target = entities.get(entity).unwrap().get::<Self>().unwrap();
+        for source_entity in relationship_target.iter() {
+            if entities.get(source_entity).is_ok() {
+                commands.queue(
+                    entity_command::despawn()
+                        .with_entity(source_entity)
+                        .handle_error_with(error_handler::silent()),
+                );
+            } else {
+                warn!(
+                    "{}Tried to despawn non-existent entity {}",
+                    caller
+                        .map(|location| format!("{location}: "))
+                        .unwrap_or_default(),
+                    source_entity
+                );
             }
         }
     }

crates/bevy_ecs/src/world/deferred_world.rs

@@ -13,7 +13,7 @@ use crate::{
     resource::Resource,
     system::{Commands, Query},
     traversal::Traversal,
-    world::{error::EntityMutableFetchError, WorldEntityFetch},
+    world::{error::EntityMutableFetchError, EntityFetcher, WorldEntityFetch},
 };
 
 use super::{unsafe_world_cell::UnsafeWorldCell, Mut, World, ON_INSERT, ON_REPLACE};
@@ -341,6 +341,53 @@ impl<'w> DeferredWorld<'w> {
         self.get_entity_mut(entities).unwrap()
     }
 
+    /// Simultaneously provides access to entity data and a command queue, which
+    /// will be applied when the [`World`] is next flushed.
+    ///
+    /// This allows using borrowed entity data to construct commands where the
+    /// borrow checker would otherwise prevent it.
+    ///
+    /// See [`World::entities_and_commands`] for the non-deferred version.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// # use bevy_ecs::{prelude::*, world::DeferredWorld};
+    /// #[derive(Component)]
+    /// struct Targets(Vec<Entity>);
+    /// #[derive(Component)]
+    /// struct TargetedBy(Entity);
+    ///
+    /// # let mut _world = World::new();
+    /// # let e1 = _world.spawn_empty().id();
+    /// # let e2 = _world.spawn_empty().id();
+    /// # let eid = _world.spawn(Targets(vec![e1, e2])).id();
+    /// let mut world: DeferredWorld = // ...
+    /// #   DeferredWorld::from(&mut _world);
+    /// let (entities, mut commands) = world.entities_and_commands();
+    ///
+    /// let entity = entities.get(eid).unwrap();
+    /// for &target in entity.get::<Targets>().unwrap().0.iter() {
+    ///     commands.entity(target).insert(TargetedBy(eid));
+    /// }
+    /// # _world.flush();
+    /// # assert_eq!(_world.get::<TargetedBy>(e1).unwrap().0, eid);
+    /// # assert_eq!(_world.get::<TargetedBy>(e2).unwrap().0, eid);
+    /// ```
+    pub fn entities_and_commands(&mut self) -> (EntityFetcher, Commands) {
+        let cell = self.as_unsafe_world_cell();
+        // SAFETY: `&mut self` gives mutable access to the entire world, and prevents simultaneous access.
+        let fetcher = unsafe { EntityFetcher::new(cell) };
+        // SAFETY:
+        // - `&mut self` gives mutable access to the entire world, and prevents simultaneous access.
+        // - Command queue access does not conflict with entity access.
+        let raw_queue = unsafe { cell.get_raw_command_queue() };
+        // SAFETY: `&mut self` ensures the commands does not outlive the world.
+        let commands = unsafe { Commands::new_raw_from_entities(raw_queue, cell.entities()) };
+
+        (fetcher, commands)
+    }
+
     /// Returns [`Query`] for the given [`QueryState`], which is used to efficiently
     /// run queries on the [`World`] by storing and reusing the [`QueryState`].
     ///

crates/bevy_ecs/src/world/entity_fetch.rs

@@ -3,12 +3,98 @@ use core::mem::MaybeUninit;
 
 use crate::{
     entity::{hash_map::EntityHashMap, hash_set::EntityHashSet, Entity, EntityDoesNotExistError},
+    error::Result,
     world::{
         error::EntityMutableFetchError, unsafe_world_cell::UnsafeWorldCell, EntityMut, EntityRef,
         EntityWorldMut,
     },
 };
 
+/// Provides a safe interface for non-structural access to the entities in a [`World`].
+///
+/// This cannot add or remove components, or spawn or despawn entities,
+/// making it relatively safe to access in concert with other ECS data.
+/// This type can be constructed via [`World::entities_and_commands`],
+/// or [`DeferredWorld::entities_and_commands`].
+///
+/// [`World`]: crate::world::World
+/// [`World::entities_and_commands`]: crate::world::World::entities_and_commands
+/// [`DeferredWorld::entities_and_commands`]: crate::world::DeferredWorld::entities_and_commands
+pub struct EntityFetcher<'w> {
+    cell: UnsafeWorldCell<'w>,
+}
+
+impl<'w> EntityFetcher<'w> {
+    // SAFETY:
+    // - The given `cell` has mutable access to all entities.
+    // - No other references to entities exist at the same time.
+    pub(crate) unsafe fn new(cell: UnsafeWorldCell<'w>) -> Self {
+        Self { cell }
+    }
+
+    /// Returns [`EntityRef`]s that expose read-only operations for the given
+    /// `entities`, returning [`Err`] if any of the given entities do not exist.
+    ///
+    /// This function supports fetching a single entity or multiple entities:
+    /// - Pass an [`Entity`] to receive a single [`EntityRef`].
+    /// - Pass a slice of [`Entity`]s to receive a [`Vec<EntityRef>`].
+    /// - Pass an array of [`Entity`]s to receive an equally-sized array of [`EntityRef`]s.
+    /// - Pass a reference to a [`EntityHashSet`](crate::entity::hash_map::EntityHashMap) to receive an
+    ///   [`EntityHashMap<EntityRef>`](crate::entity::hash_map::EntityHashMap).
+    ///
+    /// # Errors
+    ///
+    /// If any of the given `entities` do not exist in the world, the first
+    /// [`Entity`] found to be missing will return an [`EntityDoesNotExistError`].
+    ///
+    /// # Examples
+    ///
+    /// For examples, see [`World::entity`].
+    ///
+    /// [`World::entity`]: crate::world::World::entity
+    #[inline]
+    pub fn get<F: WorldEntityFetch>(
+        &self,
+        entities: F,
+    ) -> Result<F::Ref<'_>, EntityDoesNotExistError> {
+        // SAFETY: `&self` gives read access to all entities, and prevents mutable access.
+        unsafe { entities.fetch_ref(self.cell) }
+    }
+
+    /// Returns [`EntityMut`]s that expose read and write operations for the
+    /// given `entities`, returning [`Err`] if any of the given entities do not
+    /// exist.
+    ///
+    /// This function supports fetching a single entity or multiple entities:
+    /// - Pass an [`Entity`] to receive a single [`EntityMut`].
+    ///    - This reference type allows for structural changes to the entity,
+    ///      such as adding or removing components, or despawning the entity.
+    /// - Pass a slice of [`Entity`]s to receive a [`Vec<EntityMut>`].
+    /// - Pass an array of [`Entity`]s to receive an equally-sized array of [`EntityMut`]s.
+    /// - Pass a reference to a [`EntityHashSet`](crate::entity::hash_map::EntityHashMap) to receive an
+    ///   [`EntityHashMap<EntityMut>`](crate::entity::hash_map::EntityHashMap).
+    /// # Errors
+    ///
+    /// - Returns [`EntityMutableFetchError::EntityDoesNotExist`] if any of the given `entities` do not exist in the world.
+    ///     - Only the first entity found to be missing will be returned.
+    /// - Returns [`EntityMutableFetchError::AliasedMutability`] if the same entity is requested multiple times.
+    ///
+    /// # Examples
+    ///
+    /// For examples, see [`DeferredWorld::entity_mut`].
+    ///
+    /// [`DeferredWorld::entity_mut`]: crate::world::DeferredWorld::entity_mut
+    #[inline]
+    pub fn get_mut<F: WorldEntityFetch>(
+        &mut self,
+        entities: F,
+    ) -> Result<F::DeferredMut<'_>, EntityMutableFetchError> {
+        // SAFETY: `&mut self` gives mutable access to all entities,
+        // and prevents any other access to entities.
+        unsafe { entities.fetch_deferred_mut(self.cell) }
+    }
+}
+
 /// Types that can be used to fetch [`Entity`] references from a [`World`].
 ///
 /// Provided implementations are:

crates/bevy_ecs/src/world/mod.rs

@@ -21,7 +21,7 @@ pub use crate::{
 pub use bevy_ecs_macros::FromWorld;
 pub use component_constants::*;
 pub use deferred_world::DeferredWorld;
-pub use entity_fetch::WorldEntityFetch;
+pub use entity_fetch::{EntityFetcher, WorldEntityFetch};
 pub use entity_ref::{
     DynamicComponentFetch, EntityMut, EntityMutExcept, EntityRef, EntityRefExcept, EntityWorldMut,
     Entry, FilteredEntityMut, FilteredEntityRef, OccupiedEntry, TryFromFilteredError, VacantEntry,
@@ -1015,6 +1015,52 @@ impl World {
         })
     }
 
+    /// Simultaneously provides access to entity data and a command queue, which
+    /// will be applied when the world is next flushed.
+    ///
+    /// This allows using borrowed entity data to construct commands where the
+    /// borrow checker would otherwise prevent it.
+    ///
+    /// See [`DeferredWorld::entities_and_commands`] for the deferred version.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// # use bevy_ecs::{prelude::*, world::DeferredWorld};
+    /// #[derive(Component)]
+    /// struct Targets(Vec<Entity>);
+    /// #[derive(Component)]
+    /// struct TargetedBy(Entity);
+    ///
+    /// let mut world: World = // ...
+    /// #    World::new();
+    /// # let e1 = world.spawn_empty().id();
+    /// # let e2 = world.spawn_empty().id();
+    /// # let eid = world.spawn(Targets(vec![e1, e2])).id();
+    /// let (entities, mut commands) = world.entities_and_commands();
+    ///
+    /// let entity = entities.get(eid).unwrap();
+    /// for &target in entity.get::<Targets>().unwrap().0.iter() {
+    ///     commands.entity(target).insert(TargetedBy(eid));
+    /// }
+    /// # world.flush();
+    /// # assert_eq!(world.get::<TargetedBy>(e1).unwrap().0, eid);
+    /// # assert_eq!(world.get::<TargetedBy>(e2).unwrap().0, eid);
+    /// ```
+    pub fn entities_and_commands(&mut self) -> (EntityFetcher, Commands) {
+        let cell = self.as_unsafe_world_cell();
+        // SAFETY: `&mut self` gives mutable access to the entire world, and prevents simultaneous access.
+        let fetcher = unsafe { EntityFetcher::new(cell) };
+        // SAFETY:
+        // - `&mut self` gives mutable access to the entire world, and prevents simultaneous access.
+        // - Command queue access does not conflict with entity access.
+        let raw_queue = unsafe { cell.get_raw_command_queue() };
+        // SAFETY: `&mut self` ensures the commands does not outlive the world.
+        let commands = unsafe { Commands::new_raw_from_entities(raw_queue, cell.entities()) };
+
+        (fetcher, commands)
+    }
+
     /// Spawns a new [`Entity`] and returns a corresponding [`EntityWorldMut`], which can be used
     /// to add components to the entity or retrieve its id.
     ///
@@ -3692,7 +3738,7 @@ mod tests {
         entity_disabling::{DefaultQueryFilters, Disabled},
         ptr::OwningPtr,
         resource::Resource,
-        world::error::EntityMutableFetchError,
+        world::{error::EntityMutableFetchError, DeferredWorld},
     };
     use alloc::{
         borrow::ToOwned,
@@ -4388,4 +4434,44 @@ mod tests {
         world.remove_resource::<DefaultQueryFilters>();
         assert_eq!(2, world.query::<&Foo>().iter(&world).count());
     }
+
+    #[test]
+    fn entities_and_commands() {
+        #[derive(Component, PartialEq, Debug)]
+        struct Foo(u32);
+
+        let mut world = World::new();
+
+        let eid = world.spawn(Foo(35)).id();
+
+        let (mut fetcher, mut commands) = world.entities_and_commands();
+        let emut = fetcher.get_mut(eid).unwrap();
+        commands.entity(eid).despawn();
+        assert_eq!(emut.get::<Foo>().unwrap(), &Foo(35));
+
+        world.flush();
+
+        assert!(world.get_entity(eid).is_err());
+    }
+
+    #[test]
+    fn entities_and_commands_deferred() {
+        #[derive(Component, PartialEq, Debug)]
+        struct Foo(u32);
+
+        let mut world = World::new();
+
+        let eid = world.spawn(Foo(1)).id();
+
+        let mut dworld = DeferredWorld::from(&mut world);
+
+        let (mut fetcher, mut commands) = dworld.entities_and_commands();
+        let emut = fetcher.get_mut(eid).unwrap();
+        commands.entity(eid).despawn();
+        assert_eq!(emut.get::<Foo>().unwrap(), &Foo(1));
+
+        world.flush();
+
+        assert!(world.get_entity(eid).is_err());
+    }
 }