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>
This commit is contained in:
Christian Hughes 2025-03-17 13:05:50 -05:00 committed by GitHub
parent 770059c539
commit fecf2d2591
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 256 additions and 49 deletions

View File

@ -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
);
}
}
}

View File

@ -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`].
///

View File

@ -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:

View File

@ -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());
}
}