add InteriorMutableEntityRef

Author: jakobhellermann

crates/bevy_ecs/src/reflect.rs

@@ -210,19 +210,14 @@ impl<C: Component + Reflect + FromWorld> FromType<C> for ReflectComponent {
                     .map(|c| c as &dyn Reflect)
             },
             reflect_mut: |world, 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
+                // SAFETY: reflect_mut is an unsafe function pointer used by
+                // 1. `reflect_unchecked_mut` which must be called with an InteriorMutableWorld with access the the component `C` on the `entity`, and
+                // 2. reflect_mut, which has mutable world access
                 unsafe {
-                    // SAFETY: entity access through the InteriorMutableWorld is not implemented yet.
-                    // The following code only accesses the component `C` through the entity `entity`
-                    let world = world.world();
-                    world
-                        .get_entity(entity)?
-                        .get_unchecked_mut::<C>(world.last_change_tick(), world.read_change_tick())
-                        .map(|c| Mut {
-                            value: c.value as &mut dyn Reflect,
-                            ticks: c.ticks,
-                        })
+                    world.get_entity(entity)?.get_mut::<C>().map(|c| Mut {
+                        value: c.value as &mut dyn Reflect,
+                        ticks: c.ticks,
+                    })
                 }
             },
         })

crates/bevy_ecs/src/system/query.rs

@@ -1173,6 +1173,7 @@ impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> Query<'w, 's, Q, F> {
         }
         let world = self.world;
         let entity_ref = world
+            .as_interior_mutable()
             .get_entity(entity)
             .ok_or(QueryComponentError::NoSuchEntity)?;
         let component_id = world
@@ -1189,7 +1190,7 @@ impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> Query<'w, 's, Q, F> {
             .has_write(archetype_component)
         {
             entity_ref
-                .get_unchecked_mut::<T>(self.last_change_tick, self.change_tick)
+                .get_mut_using_ticks::<T>(self.last_change_tick, self.change_tick)
                 .ok_or(QueryComponentError::MissingComponent)
         } else {
             Err(QueryComponentError::MissingWriteAccess)

crates/bevy_ecs/src/world/entity_ref.rs

@@ -126,39 +126,6 @@ impl<'w> EntityRef<'w> {
             )
         }
     }
-
-    /// Gets a mutable reference to the component of type `T` associated with
-    /// this entity without ensuring there are no other borrows active and without
-    /// ensuring that the returned reference will stay valid.
-    ///
-    /// # Safety
-    ///
-    /// - The returned reference must never alias a mutable borrow of this component.
-    /// - The returned reference must not be used after this component is moved which
-    ///   may happen from **any** `insert_component`, `remove_component` or `despawn`
-    ///   operation on this world (non-exhaustive list).
-    #[inline]
-    pub unsafe fn get_unchecked_mut<T: Component>(
-        &self,
-        last_change_tick: u32,
-        change_tick: u32,
-    ) -> Option<Mut<'w, T>> {
-        // SAFETY:
-        // - entity location and entity is valid
-        // - returned component is of type T
-        // - the storage type provided is correct for T
-        self.world
-            .get_component_and_ticks_with_type(
-                TypeId::of::<T>(),
-                T::Storage::STORAGE_TYPE,
-                self.entity,
-                self.location,
-            )
-            .map(|(value, ticks)| Mut {
-                value: value.assert_unique().deref_mut::<T>(),
-                ticks: Ticks::from_tick_cells(ticks, last_change_tick, change_tick),
-            })
-    }
 }
 
 impl<'w> EntityRef<'w> {
@@ -270,7 +237,23 @@ impl<'w> EntityMut<'w> {
     #[inline]
     pub fn get_mut<T: Component>(&mut self) -> Option<Mut<'_, T>> {
         // SAFETY: world access is unique, and lifetimes enforce correct usage of returned borrow
-        unsafe { self.get_unchecked_mut::<T>() }
+        unsafe {
+            self.world
+                .get_component_and_ticks_with_type(
+                    TypeId::of::<T>(),
+                    T::Storage::STORAGE_TYPE,
+                    self.entity,
+                    self.location,
+                )
+                .map(|(value, cells)| Mut {
+                    value: value.assert_unique().deref_mut::<T>(),
+                    ticks: Ticks::from_tick_cells(
+                        cells,
+                        self.world.last_change_tick,
+                        self.world.read_change_tick(),
+                    ),
+                })
+        }
     }
 
     /// Retrieves the change ticks for the given component. This can be useful for implementing change
@@ -314,39 +297,6 @@ impl<'w> EntityMut<'w> {
         }
     }
 
-    /// Gets a mutable reference to the component of type `T` associated with
-    /// this entity without ensuring there are no other borrows active and without
-    /// ensuring that the returned reference will stay valid.
-    ///
-    /// # Safety
-    ///
-    /// - The returned reference must never alias a mutable borrow of this component.
-    /// - The returned reference must not be used after this component is moved which
-    ///   may happen from **any** `insert_component`, `remove_component` or `despawn`
-    ///   operation on this world (non-exhaustive list).
-    #[inline]
-    pub unsafe fn get_unchecked_mut<T: Component>(&self) -> Option<Mut<'_, T>> {
-        // SAFETY:
-        // - entity location and entity is valid
-        // - returned component is of type T
-        // - the storage type provided is correct for T
-        self.world
-            .get_component_and_ticks_with_type(
-                TypeId::of::<T>(),
-                T::Storage::STORAGE_TYPE,
-                self.entity,
-                self.location,
-            )
-            .map(|(value, ticks)| Mut {
-                value: value.assert_unique().deref_mut::<T>(),
-                ticks: Ticks::from_tick_cells(
-                    ticks,
-                    self.world.last_change_tick(),
-                    self.world.read_change_tick(),
-                ),
-            })
-    }
-
     /// Adds a [`Bundle`] of components to the entity.
     ///
     /// This will overwrite any previous value(s) of the same component type.
@@ -694,15 +644,19 @@ impl<'w> EntityMut<'w> {
     }
 }
 
-fn contains_component_with_type(world: &World, type_id: TypeId, location: EntityLocation) -> bool {
+pub(crate) fn contains_component_with_type(
+    world: &World,
+    type_id: TypeId,
+    location: EntityLocation,
+) -> bool {
     if let Some(component_id) = world.components.get_id(type_id) {
         contains_component_with_id(world, component_id, location)
     } else {
         false
     }
 }
 
-fn contains_component_with_id(
+pub(crate) fn contains_component_with_id(
     world: &World,
     component_id: ComponentId,
     location: EntityLocation,