From f1b2de4898382d8ec63650e66054b2795cc3651f Mon Sep 17 00:00:00 2001
From: devil-ira <justthecooldude@gmail.com>
Date: Fri, 4 Nov 2022 13:53:59 +0100
Subject: [PATCH] Add `EntityMut::world_scope`

---
 crates/bevy_ecs/src/world/entity_ref.rs    | 19 ++++++++++++-
 crates/bevy_hierarchy/src/child_builder.rs | 33 ++++++++--------------
 2 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs
index 29ff36fb903e24..f037fcbca458bd 100644
--- a/crates/bevy_ecs/src/world/entity_ref.rs
+++ b/crates/bevy_ecs/src/world/entity_ref.rs
@@ -526,7 +526,7 @@ impl<'w> EntityMut<'w> {
 
     /// Returns this `EntityMut`'s world.
     ///
-    /// See [`EntityMut::into_world_mut`] for a safe alternative.
+    /// See [`EntityMut::world_scope`] or [`EntityMut::into_world_mut`] for a safe alternative.
     ///
     /// # Safety
     /// Caller must not modify the world in a way that changes the current entity's location
@@ -543,6 +543,12 @@ impl<'w> EntityMut<'w> {
         self.world
     }
 
+    /// Gives mutable access to this `EntityMut`'s [`World`] in a temporary scope.
+    pub fn world_scope(&mut self, f: impl FnOnce(&mut World)) {
+        f(self.world);
+        self.update_location();
+    }
+
     /// Updates the internal entity location to match the current location in the internal
     /// [`World`]. This is only needed if the user called [`EntityMut::world`], which enables the
     /// location to change.
@@ -1026,4 +1032,15 @@ mod tests {
         assert!(entity.get_by_id(invalid_component_id).is_none());
         assert!(entity.get_mut_by_id(invalid_component_id).is_none());
     }
+
+    #[test]
+    fn world_scope() {
+        let mut world = World::new();
+        let mut entity = world.spawn_empty();
+        // let new_world = None;
+        entity.world_scope(|world| {
+            // new_world = Some(world);
+            entity.insert(());
+        });
+    }
 }
diff --git a/crates/bevy_hierarchy/src/child_builder.rs b/crates/bevy_hierarchy/src/child_builder.rs
index b23c2ef268e1d6..a5c5d6708a6674 100644
--- a/crates/bevy_hierarchy/src/child_builder.rs
+++ b/crates/bevy_hierarchy/src/child_builder.rs
@@ -456,29 +456,23 @@ pub trait BuildWorldChildren {
 
 impl<'w> BuildWorldChildren for EntityMut<'w> {
     fn with_children(&mut self, spawn_children: impl FnOnce(&mut WorldChildBuilder)) -> &mut Self {
-        {
-            let entity = self.id();
+        let entity = self.id();
+        self.world_scope(|world| {
             let mut builder = WorldChildBuilder {
                 current_entity: None,
                 parent_entities: vec![entity],
-                // SAFETY: The EntityLocation is updated before any other methods are called on self.
-                world: unsafe { self.world_mut() },
+                world,
             };
-
             spawn_children(&mut builder);
-        }
-        self.update_location();
+        });
         self
     }
 
     fn push_children(&mut self, children: &[Entity]) -> &mut Self {
         let parent = self.id();
-        {
-            // SAFETY: The EntityLocation is updated before any other methods are called on self.
-            let world = unsafe { self.world_mut() };
+        self.world_scope(|world| {
             update_old_parents(world, parent, children);
-            self.update_location();
-        }
+        });
         if let Some(mut children_component) = self.get_mut::<Children>() {
             children_component
                 .0
@@ -492,13 +486,9 @@ impl<'w> BuildWorldChildren for EntityMut<'w> {
 
     fn insert_children(&mut self, index: usize, children: &[Entity]) -> &mut Self {
         let parent = self.id();
-        {
-            // SAFETY: The EntityLocation is updated before any other methods are called on self.
-            let world = unsafe { self.world_mut() };
+        self.world_scope(|world| {
             update_old_parents(world, parent, children);
-            self.update_location();
-        }
-
+        });
         if let Some(mut children_component) = self.get_mut::<Children>() {
             children_component
                 .0
@@ -512,10 +502,9 @@ impl<'w> BuildWorldChildren for EntityMut<'w> {
 
     fn remove_children(&mut self, children: &[Entity]) -> &mut Self {
         let parent = self.id();
-        // SAFETY: The EntityLocation is updated before any other methods are called on self.
-        let world = unsafe { self.world_mut() };
-        remove_children(parent, children, world);
-        self.update_location();
+        self.world_scope(|world| {
+            remove_children(parent, children, world);
+        });
         self
     }
 }