Use entity references everywhere

Author: Suficio

crates/bevy_ecs/src/reflect.rs

@@ -42,11 +42,11 @@ pub struct ReflectComponent(ReflectComponentFns);
 #[derive(Clone)]
 pub struct ReflectComponentFns {
     /// Function pointer implementing [`ReflectComponent::insert()`].
-    pub insert: fn(&mut World, Entity, &dyn Reflect),
+    pub insert: fn(&mut EntityMut, &dyn Reflect),
     /// Function pointer implementing [`ReflectComponent::apply()`].
     pub apply: fn(&mut EntityMut, &dyn Reflect),
     /// Function pointer implementing [`ReflectComponent::apply_or_insert()`].
-    pub apply_or_insert: fn(&mut World, Entity, &dyn Reflect),
+    pub apply_or_insert: fn(&mut EntityMut, &dyn Reflect),
     /// Function pointer implementing [`ReflectComponent::remove()`].
     pub remove: fn(&mut EntityMut),
     /// Function pointer implementing [`ReflectComponent::contains()`].
@@ -56,7 +56,7 @@ pub struct ReflectComponentFns {
     /// Function pointer implementing [`ReflectComponent::reflect_mut()`].
     pub reflect_mut: for<'a> fn(&'a mut EntityMut<'a>) -> Option<Mut<'a, dyn Reflect>>,
     /// Function pointer implementing [`ReflectComponent::reflect_unchecked_mut()`].
-    pub reflect_unchecked_mut: unsafe fn(&World, Entity) -> Option<Mut<dyn Reflect>>,
+    pub reflect_unchecked_mut: unsafe fn(EntityRef) -> Option<Mut<dyn Reflect>>,
     /// Function pointer implementing [`ReflectComponent::copy()`].
     pub copy: fn(&World, &mut World, Entity, Entity),
 }
@@ -74,12 +74,8 @@ impl ReflectComponentFns {
 
 impl ReflectComponent {
     /// Insert a reflected [`Component`] into the entity like [`insert()`](crate::world::EntityMut::insert).
-    ///
-    /// # Panics
-    ///
-    /// Panics if there is no such entity.
-    pub fn insert(&self, world: &mut World, entity: Entity, component: &dyn Reflect) {
-        (self.0.insert)(world, entity, component);
+    pub fn insert(&self, entity: &mut EntityMut, component: &dyn Reflect) {
+        (self.0.insert)(entity, component);
     }
 
     /// Uses reflection to set the value of this [`Component`] type in the entity to the given value.
@@ -92,12 +88,8 @@ impl ReflectComponent {
     }
 
     /// Uses reflection to set the value of this [`Component`] type in the entity to the given value or insert a new one if it does not exist.
-    ///
-    /// # Panics
-    ///
-    /// Panics if the `entity` does not exist.
-    pub fn apply_or_insert(&self, world: &mut World, entity: Entity, component: &dyn Reflect) {
-        (self.0.apply_or_insert)(world, entity, component);
+    pub fn apply_or_insert(&self, entity: &mut EntityMut, component: &dyn Reflect) {
+        (self.0.apply_or_insert)(entity, component);
     }
 
     /// Removes this [`Component`] type from the entity. Does nothing if it doesn't exist.
@@ -132,10 +124,9 @@ impl ReflectComponent {
     /// * Don't call this method more than once in the same scope for a given [`Component`].
     pub unsafe fn reflect_unchecked_mut<'a>(
         &self,
-        world: &'a World,
-        entity: Entity,
+        entity: EntityRef<'a>,
     ) -> Option<Mut<'a, dyn Reflect>> {
-        (self.0.reflect_unchecked_mut)(world, entity)
+        (self.0.reflect_unchecked_mut)(entity)
     }
 
     /// Gets the value of this [`Component`] type from entity from `source_world` and [applies](Self::apply()) it to the value of this [`Component`] type in entity in `destination_world`.
@@ -177,22 +168,22 @@ impl ReflectComponent {
 impl<C: Component + Reflect + FromWorld> FromType<C> for ReflectComponent {
     fn from_type() -> Self {
         ReflectComponent(ReflectComponentFns {
-            insert: |world, entity, reflected_component| {
-                let mut component = C::from_world(world);
+            insert: |entity, reflected_component| {
+                let mut component = entity.world_scope(|world| C::from_world(world));
                 component.apply(reflected_component);
-                world.entity_mut(entity).insert(component);
+                entity.insert(component);
             },
             apply: |entity, reflected_component| {
                 let mut component = entity.get_mut::<C>().unwrap();
                 component.apply(reflected_component);
             },
-            apply_or_insert: |world, entity, reflected_component| {
-                if let Some(mut component) = world.get_mut::<C>(entity) {
+            apply_or_insert: |entity, reflected_component| {
+                if let Some(mut component) = entity.get_mut::<C>() {
                     component.apply(reflected_component);
                 } else {
-                    let mut component = C::from_world(world);
+                    let mut component = entity.world_scope(|world| C::from_world(world));
                     component.apply(reflected_component);
-                    world.entity_mut(entity).insert(component);
+                    entity.insert(component);
                 }
             },
             remove: |entity| {
@@ -214,13 +205,16 @@ impl<C: Component + Reflect + FromWorld> FromType<C> for ReflectComponent {
                     ticks: c.ticks,
                 })
             },
-            reflect_unchecked_mut: |world, entity| {
+            reflect_unchecked_mut: |entity| {
                 // SAFETY: reflect_mut is an unsafe function pointer used by `reflect_unchecked_mut` which promises to never
-                // produce aliasing mutable references, and reflect_mut, which has mutable world access
+                // produce aliasing mutable references, and reflect_mut, which has mutable world access4
+                let world = entity.world();
+                let last_change_tick = world.last_change_tick();
+                let read_change_tick = world.read_change_tick();
+
                 unsafe {
-                    world
-                        .get_entity(entity)?
-                        .get_unchecked_mut::<C>(world.last_change_tick(), world.read_change_tick())
+                    entity
+                        .get_unchecked_mut::<C>(last_change_tick, read_change_tick)
                         .map(|c| Mut {
                             value: c.value as &mut dyn Reflect,
                             ticks: c.ticks,

crates/bevy_ecs/src/world/entity_ref.rs

@@ -639,9 +639,10 @@ impl<'w> EntityMut<'w> {
     }
 
     /// 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);
+    pub fn world_scope<T>(&mut self, f: impl FnOnce(&mut World) -> T) -> T {
+        let result = f(self.world);
         self.update_location();
+        result
     }
 
     /// Updates the internal entity location to match the current location in the internal