Add EntityWorldMut::(try_)resource_scope (#20162)

# Objective

Fixes #20139

## Solution

Implement the methods, and leverage them where applicable

## Testing

Added unit tests

---

## Showcase

```rust
let id = entity_world_mut.id();
let world = entity_world_mut.into_world_mut();
world.resource_scope::<_, _>(|world, res| {
    let entity_world_mut = world.entity_mut(id);
    /* ... */
});
```

becomes

```rust
entity_world_mut.resource_scope::<_, _>(|entity, res| {
    /* ... */
});
```
This commit is contained in:
BigWingBeat 2025-07-16 18:37:25 +01:00 committed by GitHub
parent d28d18e2ff
commit e7c14bd06b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 174 additions and 111 deletions

View File

@ -31,20 +31,16 @@ pub fn signal_event_update_system(signal: Option<ResMut<EventRegistry>>) {
/// A system that calls [`Events::update`](super::Events::update) on all registered [`Events`][super::Events] in the world.
pub fn event_update_system(world: &mut World, mut last_change_tick: Local<Tick>) {
if world.contains_resource::<EventRegistry>() {
world.resource_scope(|world, mut registry: Mut<EventRegistry>| {
registry.run_updates(world, *last_change_tick);
world.try_resource_scope(|world, mut registry: Mut<EventRegistry>| {
registry.run_updates(world, *last_change_tick);
registry.should_update = match registry.should_update {
// If we're always updating, keep doing so.
ShouldUpdateEvents::Always => ShouldUpdateEvents::Always,
// Disable the system until signal_event_update_system runs again.
ShouldUpdateEvents::Waiting | ShouldUpdateEvents::Ready => {
ShouldUpdateEvents::Waiting
}
};
});
}
registry.should_update = match registry.should_update {
// If we're always updating, keep doing so.
ShouldUpdateEvents::Always => ShouldUpdateEvents::Always,
// Disable the system until signal_event_update_system runs again.
ShouldUpdateEvents::Waiting | ShouldUpdateEvents::Ready => ShouldUpdateEvents::Waiting,
};
});
*last_change_tick = world.change_tick();
}

View File

@ -1,10 +1,9 @@
use crate::{
entity::Entity,
prelude::Mut,
reflect::{AppTypeRegistry, ReflectBundle, ReflectComponent},
resource::Resource,
system::EntityCommands,
world::{EntityWorldMut, World},
world::EntityWorldMut,
};
use alloc::{borrow::Cow, boxed::Box};
use bevy_reflect::{PartialReflect, TypeRegistry};
@ -22,7 +21,7 @@ pub trait ReflectCommandExt {
/// - If [`AppTypeRegistry`] does not have the reflection data for the given
/// [`Component`](crate::component::Component) or [`Bundle`](crate::bundle::Bundle).
/// - If the component or bundle data is invalid. See [`PartialReflect::apply`] for further details.
/// - If [`AppTypeRegistry`] is not present in the [`World`].
/// - If [`AppTypeRegistry`] is not present in the [`World`](crate::world::World).
///
/// # Note
///
@ -92,11 +91,11 @@ pub trait ReflectCommandExt {
///
/// # Panics
///
/// - If the given [`Resource`] is not present in the [`World`].
/// - If the given [`Resource`] is not present in the [`World`](crate::world::World).
///
/// # Note
///
/// - The given [`Resource`] is removed from the [`World`] before the command is applied.
/// - The given [`Resource`] is removed from the [`World`](crate::world::World) before the command is applied.
fn insert_reflect_with_registry<T: Resource + AsRef<TypeRegistry>>(
&mut self,
component: Box<dyn PartialReflect>,
@ -214,7 +213,7 @@ impl<'w> EntityWorldMut<'w> {
/// - If [`AppTypeRegistry`] does not have the reflection data for the given
/// [`Component`](crate::component::Component) or [`Bundle`](crate::bundle::Bundle).
/// - If the component or bundle data is invalid. See [`PartialReflect::apply`] for further details.
/// - If [`AppTypeRegistry`] is not present in the [`World`].
/// - If [`AppTypeRegistry`] is not present in the [`World`](crate::world::World).
///
/// # Note
///
@ -222,15 +221,10 @@ impl<'w> EntityWorldMut<'w> {
/// is much slower.
pub fn insert_reflect(&mut self, component: Box<dyn PartialReflect>) -> &mut Self {
self.assert_not_despawned();
let entity_id = self.id();
self.world_scope(|world| {
world.resource_scope(|world, registry: Mut<AppTypeRegistry>| {
let type_registry = &registry.as_ref().read();
insert_reflect_with_registry_ref(world, entity_id, type_registry, component);
});
world.flush();
self.resource_scope(|entity, registry: Mut<AppTypeRegistry>| {
let type_registry = &registry.as_ref().read();
insert_reflect_with_registry_ref(entity, type_registry, component);
});
self.update_location();
self
}
@ -245,21 +239,16 @@ impl<'w> EntityWorldMut<'w> {
/// - If the given [`Resource`] does not have the reflection data for the given
/// [`Component`](crate::component::Component) or [`Bundle`](crate::bundle::Bundle).
/// - If the component or bundle data is invalid. See [`PartialReflect::apply`] for further details.
/// - If the given [`Resource`] is not present in the [`World`].
/// - If the given [`Resource`] is not present in the [`World`](crate::world::World).
pub fn insert_reflect_with_registry<T: Resource + AsRef<TypeRegistry>>(
&mut self,
component: Box<dyn PartialReflect>,
) -> &mut Self {
self.assert_not_despawned();
let entity_id = self.id();
self.world_scope(|world| {
world.resource_scope(|world, registry: Mut<T>| {
let type_registry = registry.as_ref().as_ref();
insert_reflect_with_registry_ref(world, entity_id, type_registry, component);
});
world.flush();
self.resource_scope(|entity, registry: Mut<T>| {
let type_registry = registry.as_ref().as_ref();
insert_reflect_with_registry_ref(entity, type_registry, component);
});
self.update_location();
self
}
@ -275,7 +264,7 @@ impl<'w> EntityWorldMut<'w> {
/// # Panics
///
/// - If the entity has been despawned while this `EntityWorldMut` is still alive.
/// - If [`AppTypeRegistry`] is not present in the [`World`].
/// - If [`AppTypeRegistry`] is not present in the [`World`](crate::world::World).
///
/// # Note
///
@ -283,20 +272,10 @@ impl<'w> EntityWorldMut<'w> {
/// is much slower.
pub fn remove_reflect(&mut self, component_type_path: Cow<'static, str>) -> &mut Self {
self.assert_not_despawned();
let entity_id = self.id();
self.world_scope(|world| {
world.resource_scope(|world, registry: Mut<AppTypeRegistry>| {
let type_registry = &registry.as_ref().read();
remove_reflect_with_registry_ref(
world,
entity_id,
type_registry,
component_type_path,
);
});
world.flush();
self.resource_scope(|entity, registry: Mut<AppTypeRegistry>| {
let type_registry = &registry.as_ref().read();
remove_reflect_with_registry_ref(entity, type_registry, component_type_path);
});
self.update_location();
self
}
@ -313,34 +292,23 @@ impl<'w> EntityWorldMut<'w> {
/// # Panics
///
/// - If the entity has been despawned while this `EntityWorldMut` is still alive.
/// - If [`AppTypeRegistry`] is not present in the [`World`].
/// - If [`AppTypeRegistry`] is not present in the [`World`](crate::world::World).
pub fn remove_reflect_with_registry<T: Resource + AsRef<TypeRegistry>>(
&mut self,
component_type_path: Cow<'static, str>,
) -> &mut Self {
self.assert_not_despawned();
let entity_id = self.id();
self.world_scope(|world| {
world.resource_scope(|world, registry: Mut<T>| {
let type_registry = registry.as_ref().as_ref();
remove_reflect_with_registry_ref(
world,
entity_id,
type_registry,
component_type_path,
);
});
world.flush();
self.resource_scope(|entity, registry: Mut<T>| {
let type_registry = registry.as_ref().as_ref();
remove_reflect_with_registry_ref(entity, type_registry, component_type_path);
});
self.update_location();
self
}
}
/// Helper function to add a reflect component or bundle to a given entity
fn insert_reflect_with_registry_ref(
world: &mut World,
entity: Entity,
entity: &mut EntityWorldMut,
type_registry: &TypeRegistry,
component: Box<dyn PartialReflect>,
) {
@ -348,18 +316,14 @@ fn insert_reflect_with_registry_ref(
.get_represented_type_info()
.expect("component should represent a type.");
let type_path = type_info.type_path();
let Ok(mut entity) = world.get_entity_mut(entity) else {
panic!("error[B0003]: Could not insert a reflected component (of type {type_path}) for entity {entity}, which {}. See: https://bevy.org/learn/errors/b0003",
world.entities().entity_does_not_exist_error_details(entity));
};
let Some(type_registration) = type_registry.get(type_info.type_id()) else {
panic!("`{type_path}` should be registered in type registry via `App::register_type<{type_path}>`");
};
if let Some(reflect_component) = type_registration.data::<ReflectComponent>() {
reflect_component.insert(&mut entity, component.as_partial_reflect(), type_registry);
reflect_component.insert(entity, component.as_partial_reflect(), type_registry);
} else if let Some(reflect_bundle) = type_registration.data::<ReflectBundle>() {
reflect_bundle.insert(&mut entity, component.as_partial_reflect(), type_registry);
reflect_bundle.insert(entity, component.as_partial_reflect(), type_registry);
} else {
panic!("`{type_path}` should have #[reflect(Component)] or #[reflect(Bundle)]");
}
@ -367,21 +331,17 @@ fn insert_reflect_with_registry_ref(
/// Helper function to remove a reflect component or bundle from a given entity
fn remove_reflect_with_registry_ref(
world: &mut World,
entity: Entity,
entity: &mut EntityWorldMut,
type_registry: &TypeRegistry,
component_type_path: Cow<'static, str>,
) {
let Ok(mut entity) = world.get_entity_mut(entity) else {
return;
};
let Some(type_registration) = type_registry.get_with_type_path(&component_type_path) else {
return;
};
if let Some(reflect_component) = type_registration.data::<ReflectComponent>() {
reflect_component.remove(&mut entity);
reflect_component.remove(entity);
} else if let Some(reflect_bundle) = type_registration.data::<ReflectBundle>() {
reflect_bundle.remove(&mut entity);
reflect_bundle.remove(entity);
}
}

View File

@ -1547,6 +1547,50 @@ impl<'w> EntityWorldMut<'w> {
self.world.get_resource_mut()
}
/// Temporarily removes the requested resource from the [`World`], runs custom user code,
/// then re-adds the resource before returning.
///
/// # Panics
///
/// Panics if the resource does not exist.
/// Use [`try_resource_scope`](Self::try_resource_scope) instead if you want to handle this case.
///
/// See [`World::resource_scope`] for further details.
#[track_caller]
pub fn resource_scope<R: Resource, U>(
&mut self,
f: impl FnOnce(&mut EntityWorldMut, Mut<R>) -> U,
) -> U {
let id = self.id();
self.world_scope(|world| {
world.resource_scope(|world, res| {
// Acquiring a new EntityWorldMut here and using that instead of `self` is fine because
// the outer `world_scope` will handle updating our location if it gets changed by the user code
let mut this = world.entity_mut(id);
f(&mut this, res)
})
})
}
/// Temporarily removes the requested resource from the [`World`] if it exists, runs custom user code,
/// then re-adds the resource before returning. Returns `None` if the resource does not exist in the [`World`].
///
/// See [`World::try_resource_scope`] for further details.
pub fn try_resource_scope<R: Resource, U>(
&mut self,
f: impl FnOnce(&mut EntityWorldMut, Mut<R>) -> U,
) -> Option<U> {
let id = self.id();
self.world_scope(|world| {
world.try_resource_scope(|world, res| {
// Acquiring a new EntityWorldMut here and using that instead of `self` is fine because
// the outer `world_scope` will handle updating our location if it gets changed by the user code
let mut this = world.entity_mut(id);
f(&mut this, res)
})
})
}
/// Retrieves the change ticks for the given component. This can be useful for implementing change
/// detection in custom runtimes.
///
@ -4959,6 +5003,46 @@ mod tests {
assert!(entity.get_mut_by_id(invalid_component_id).is_err());
}
#[derive(Resource)]
struct R(usize);
#[test]
fn entity_mut_resource_scope() {
// Keep in sync with the `resource_scope` test in lib.rs
let mut world = World::new();
let mut entity = world.spawn_empty();
assert!(entity.try_resource_scope::<R, _>(|_, _| {}).is_none());
entity.world_scope(|world| world.insert_resource(R(0)));
entity.resource_scope(|entity: &mut EntityWorldMut, mut value: Mut<R>| {
value.0 += 1;
assert!(!entity.world().contains_resource::<R>());
});
assert_eq!(entity.resource::<R>().0, 1);
}
#[test]
fn entity_mut_resource_scope_panic() {
let mut world = World::new();
world.insert_resource(R(0));
let mut entity = world.spawn_empty();
let old_location = entity.location();
let result = std::panic::catch_unwind(AssertUnwindSafe(|| {
entity.resource_scope(|entity: &mut EntityWorldMut, _: Mut<R>| {
// Change the entity's `EntityLocation`.
entity.insert(TestComponent(0));
// Ensure that the entity location still gets updated even in case of a panic.
panic!("this should get caught by the outer scope")
});
}));
assert!(result.is_err());
// Ensure that the location has been properly updated.
assert_ne!(entity.location(), old_location);
}
// regression test for https://github.com/bevyengine/bevy/pull/7387
#[test]
fn entity_mut_world_scope_panic() {
@ -4983,6 +5067,28 @@ mod tests {
assert_ne!(entity.location(), old_location);
}
#[test]
fn entity_mut_reborrow_scope_panic() {
let mut world = World::new();
let mut entity = world.spawn_empty();
let old_location = entity.location();
let res = std::panic::catch_unwind(AssertUnwindSafe(|| {
entity.reborrow_scope(|mut entity| {
// Change the entity's `EntityLocation`, which invalidates the original `EntityWorldMut`.
// This will get updated at the end of the scope.
entity.insert(TestComponent(0));
// Ensure that the entity location still gets updated even in case of a panic.
panic!("this should get caught by the outer scope")
});
}));
assert!(res.is_err());
// Ensure that the location has been properly updated.
assert_ne!(entity.location(), old_location);
}
// regression test for https://github.com/bevyengine/bevy/pull/7805
#[test]
fn removing_sparse_updates_archetype_row() {
@ -5420,9 +5526,6 @@ mod tests {
#[derive(Component)]
struct A;
#[derive(Resource)]
struct R;
#[test]
fn disjoint_access() {
fn disjoint_readonly(_: Query<EntityMut, With<A>>, _: Query<EntityRef, Without<A>>) {}

View File

@ -2551,6 +2551,11 @@ impl World {
/// This enables safe simultaneous mutable access to both a resource and the rest of the [`World`].
/// For more complex access patterns, consider using [`SystemState`](crate::system::SystemState).
///
/// # Panics
///
/// Panics if the resource does not exist.
/// Use [`try_resource_scope`](Self::try_resource_scope) instead if you want to handle this case.
///
/// # Example
/// ```
/// use bevy_ecs::prelude::*;
@ -2568,8 +2573,6 @@ impl World {
/// });
/// assert_eq!(world.get_resource::<A>().unwrap().0, 2);
/// ```
///
/// See also [`try_resource_scope`](Self::try_resource_scope).
#[track_caller]
pub fn resource_scope<R: Resource, U>(&mut self, f: impl FnOnce(&mut World, Mut<R>) -> U) -> U {
self.try_resource_scope(f)

View File

@ -44,37 +44,38 @@ impl BuildChildrenTransformExt for EntityCommands<'_> {
impl BuildChildrenTransformExt for EntityWorldMut<'_> {
fn set_parent_in_place(&mut self, parent: Entity) -> &mut Self {
let child = self.id();
self.world_scope(|world| {
world.entity_mut(parent).add_child(child);
// FIXME: Replace this closure with a `try` block. See: https://github.com/rust-lang/rust/issues/31436.
let mut update_transform = || {
let parent = *world.get_entity(parent).ok()?.get::<GlobalTransform>()?;
let child_global = *world.get_entity(child).ok()?.get::<GlobalTransform>()?;
let mut child_entity = world.get_entity_mut(child).ok()?;
let mut child = child_entity.get_mut::<Transform>()?;
*child = child_global.reparented_to(&parent);
Some(())
};
update_transform();
});
// FIXME: Replace this closure with a `try` block. See: https://github.com/rust-lang/rust/issues/31436.
let mut update_transform = || {
let child = self.id();
let parent_global = self.world_scope(|world| {
world
.get_entity_mut(parent)
.ok()?
.add_child(child)
.get::<GlobalTransform>()
.copied()
})?;
let child_global = self.get::<GlobalTransform>()?;
let new_child_local = child_global.reparented_to(&parent_global);
let mut child_local = self.get_mut::<Transform>()?;
*child_local = new_child_local;
Some(())
};
update_transform();
self
}
fn remove_parent_in_place(&mut self) -> &mut Self {
let child = self.id();
self.world_scope(|world| {
world.entity_mut(child).remove::<ChildOf>();
// FIXME: Replace this closure with a `try` block. See: https://github.com/rust-lang/rust/issues/31436.
let mut update_transform = || {
let child_global = *world.get_entity(child).ok()?.get::<GlobalTransform>()?;
let mut child_entity = world.get_entity_mut(child).ok()?;
let mut child = child_entity.get_mut::<Transform>()?;
*child = child_global.compute_transform();
Some(())
};
update_transform();
});
self.remove::<ChildOf>();
// FIXME: Replace this closure with a `try` block. See: https://github.com/rust-lang/rust/issues/31436.
let mut update_transform = || {
let global = self.get::<GlobalTransform>()?;
let new_local = global.compute_transform();
let mut local = self.get_mut::<Transform>()?;
*local = new_local;
Some(())
};
update_transform();
self
}
}