add InteriorMutableEntityRef

Author: jakobhellermann

crates/bevy_ecs/src/world/entity_ref.rs

@@ -673,15 +673,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,

crates/bevy_ecs/src/world/interior_mutable_world.rs

@@ -1,9 +1,12 @@
 #![warn(unsafe_op_in_unsafe_fn)]
 
-use super::{Mut, World};
+use super::{entity_ref, Mut, World};
 use crate::{
+    archetype::Archetype,
     change_detection::{MutUntyped, Ticks},
-    component::ComponentId,
+    component::{ComponentId, ComponentTicks},
+    entity::{Entity, EntityLocation},
+    prelude::Component,
     system::Resource,
 };
 use bevy_ptr::Ptr;
@@ -81,6 +84,17 @@ impl<'w> InteriorMutableWorld<'w> {
         self.0
     }
 
+    /// Retrieves an [`InteriorMutableEntityRef`] that exposes read and write operations for the given `entity`.
+    /// Similar to the [`InteriorMutableWorld`], you are in charge of making sure that no aliasing rules are violated.
+    pub fn get_entity(&self, entity: Entity) -> Option<InteriorMutableEntityRef<'w>> {
+        let location = self.0.entities.get(entity)?;
+        Some(InteriorMutableEntityRef {
+            world: InteriorMutableWorld(self.0),
+            entity,
+            location,
+        })
+    }
+
     /// Gets a reference to the resource of the given type if it exists
     #[inline]
     pub unsafe fn get_resource<R: Resource>(&self) -> Option<&'w R> {
@@ -192,3 +206,192 @@ impl<'w> InteriorMutableWorld<'w> {
         unsafe { self.get_resource_mut_with_id(component_id) }
     }
 }
+
+/// A interior-mutable reference to a particular [`Entity`] and all of its components
+#[derive(Copy, Clone)]
+pub struct InteriorMutableEntityRef<'w> {
+    world: InteriorMutableWorld<'w>,
+    entity: Entity,
+    location: EntityLocation,
+}
+
+impl<'w> InteriorMutableEntityRef<'w> {
+    #[inline]
+    #[must_use = "Omit the .id() call if you do not need to store the `Entity` identifier."]
+    pub fn id(&self) -> Entity {
+        self.entity
+    }
+
+    #[inline]
+    pub fn location(&self) -> EntityLocation {
+        self.location
+    }
+
+    #[inline]
+    pub fn archetype(&self) -> &Archetype {
+        &self.world.0.archetypes[self.location.archetype_id]
+    }
+
+    #[inline]
+    pub fn world(&self) -> InteriorMutableWorld<'w> {
+        self.world
+    }
+
+    #[inline]
+    pub fn contains<T: Component>(&self) -> bool {
+        self.contains_type_id(TypeId::of::<T>())
+    }
+
+    #[inline]
+    pub fn contains_id(&self, component_id: ComponentId) -> bool {
+        entity_ref::contains_component_with_id(self.world.0, component_id, self.location)
+    }
+
+    #[inline]
+    pub fn contains_type_id(&self, type_id: TypeId) -> bool {
+        entity_ref::contains_component_with_type(self.world.0, type_id, self.location)
+    }
+
+    #[inline]
+    pub unsafe fn get<T: Component>(&self) -> Option<&'w T> {
+        // SAFETY:
+        // - entity location is valid
+        // - archetypes and components come from the same world
+        // - world access is immutable, lifetime tied to `&self`
+        unsafe {
+            self.world
+                .0
+                .storages
+                .get_component_with_type(
+                    &self.world.0.archetypes,
+                    &self.world.0.components,
+                    TypeId::of::<T>(),
+                    self.entity,
+                    self.location,
+                )
+                // SAFETY: returned component is of type T
+                .map(|value| value.deref::<T>())
+        }
+    }
+
+    /// Retrieves the change ticks for the given component. This can be useful for implementing change
+    /// detection in custom runtimes.
+    #[inline]
+    pub unsafe fn get_change_ticks<T: Component>(&self) -> Option<&'w ComponentTicks> {
+        // SAFETY:
+        // - entity location is valid
+        // - archetypes and components come from the same world
+        // - world access is immutable, lifetime tied to `&self`
+        unsafe {
+            self.world
+                .0
+                .storages
+                .get_ticks_with_type(
+                    &self.world.0.archetypes,
+                    &self.world.0.components,
+                    TypeId::of::<T>(),
+                    self.entity,
+                    self.location,
+                )
+                // SAFETY: we have read access, lifetime tied to `&self`
+                .map(|ticks| ticks.deref())
+        }
+    }
+
+    #[inline]
+    pub unsafe fn get_mut<T: Component>(&self) -> Option<Mut<'w, T>> {
+        // SAFETY: same safety requirements
+        unsafe {
+            self.get_mut_using_ticks(
+                self.world.0.last_change_tick(),
+                self.world.0.read_change_tick(),
+            )
+        }
+    }
+
+    #[inline]
+    pub unsafe fn get_mut_using_ticks<T: Component>(
+        &self,
+        last_change_tick: u32,
+        change_tick: u32,
+    ) -> Option<Mut<'w, T>> {
+        unsafe {
+            self.world
+                .0
+                .storages
+                .get_component_and_ticks_with_type(
+                    &self.world.0.archetypes,
+                    &self.world.0.components,
+                    TypeId::of::<T>(),
+                    self.entity,
+                    self.location,
+                )
+                .map(|(value, ticks)| Mut {
+                    value: value.assert_unique().deref_mut::<T>(),
+                    ticks: Ticks {
+                        component_ticks: ticks.deref_mut(),
+                        last_change_tick,
+                        change_tick,
+                    },
+                })
+        }
+    }
+}
+
+impl<'w> InteriorMutableEntityRef<'w> {
+    /// Gets the component of the given [`ComponentId`] from the entity.
+    ///
+    /// **You should prefer to use the typed API where possible and only
+    /// use this in cases where the actual component types are not known at
+    /// compile time.**
+    ///
+    /// Unlike [`InteriorMutableEntityRef::get`], this returns a raw pointer to the component,
+    /// which is only valid while the `'w` borrow of the lifetime is active.
+    #[inline]
+    pub unsafe fn get_by_id(&self, component_id: ComponentId) -> Option<Ptr<'w>> {
+        self.world.0.components.get_info(component_id)?;
+        // SAFETY: entity_location is valid, component_id is valid as checked by the line above
+        unsafe {
+            self.world.0.storages.get_component(
+                &self.world.0.archetypes,
+                &self.world.0.components,
+                component_id,
+                self.entity,
+                self.location,
+            )
+        }
+    }
+
+    /// Retrieves a mutable untyped reference to the given `entity`'s [Component] of the given [`ComponentId`].
+    /// Returns [None] if the `entity` does not have a [Component] of the given type.
+    ///
+    /// **You should prefer to use the typed API [`InteriorMutableEntityRef::get_mut`] where possible and only
+    /// use this in cases where the actual types are not known at compile time.**
+    #[inline]
+    pub unsafe fn get_mut_by_id(&self, component_id: ComponentId) -> Option<MutUntyped<'w>> {
+        self.world.0.components.get_info(component_id)?;
+        // SAFETY: entity_location is valid, component_id is valid as checked by the line above
+        unsafe {
+            self.world
+                .0
+                .storages
+                .get_component_and_ticks(
+                    &self.world.0.archetypes,
+                    &self.world.0.components,
+                    component_id,
+                    self.entity,
+                    self.location,
+                )
+                .map(|(value, ticks)| MutUntyped {
+                    // SAFETY: world access validated by caller and ties world lifetime to `MutUntyped` lifetime
+                    value: value.assert_unique(),
+                    ticks: Ticks {
+                        // SAFETY: world access is validated by caller and ties world lifetime to `Mut` lifetime
+                        component_ticks: ticks.deref_mut(),
+                        last_change_tick: self.world.0.last_change_tick(),
+                        change_tick: self.world.0.read_change_tick(),
+                    },
+                })
+        }
+    }
+}