EntityRef/Mut get_components

crates/bevy_ecs/src/world/entity_ref.rs

@@ -4,7 +4,7 @@ use crate::{
     change_detection::MutUntyped,
     component::{Component, ComponentId, ComponentTicks, Components, StorageType},
     entity::{Entities, Entity, EntityLocation},
-    query::{Access, DebugCheckedUnwrap},
+    query::{Access, DebugCheckedUnwrap, QueryData, ReadOnlyQueryData},
     removal_detection::RemovedComponentEvents,
     storage::Storages,
     world::{Mut, World},
@@ -153,6 +153,17 @@ impl<'w> EntityRef<'w> {
         // SAFETY: We have read-only access to all components of this entity.
         unsafe { self.0.get_by_id(component_id) }
     }
+
+    /// Returns read-only components for the current entity that match the query `Q`.
+    pub fn components<Q: ReadOnlyQueryData>(&self) -> Q::Item<'w> {
+        self.get_components::<Q>().expect(QUERY_MISMATCH_ERROR)
+    }
+
+    /// Returns read-only components for the current entity that match the query `Q`.
+    pub fn get_components<Q: ReadOnlyQueryData>(&self) -> Option<Q::Item<'w>> {
+        // SAFETY: We have read-only access to all components of this entity.
+        unsafe { self.0.get_components::<Q>() }
+    }
 }
 
 impl<'w> From<EntityWorldMut<'w>> for EntityRef<'w> {
@@ -348,6 +359,28 @@ impl<'w> EntityMut<'w> {
         self.as_readonly().get()
     }
 
+    /// Returns read-only components for the current entity that match the query `Q`.
+    pub fn components<Q: ReadOnlyQueryData>(&self) -> Q::Item<'_> {
+        self.get_components::<Q>().expect(QUERY_MISMATCH_ERROR)
+    }
+
+    /// Returns components for the current entity that match the query `Q`.
+    pub fn components_mut<Q: QueryData>(&mut self) -> Q::Item<'_> {
+        self.get_components_mut::<Q>().expect(QUERY_MISMATCH_ERROR)
+    }
+
+    /// Returns read-only components for the current entity that match the query `Q`.
+    pub fn get_components<Q: ReadOnlyQueryData>(&self) -> Option<Q::Item<'_>> {
+        // SAFETY: We have read-only access to all components of this entity.
+        unsafe { self.0.get_components::<Q>() }
+    }
+
+    /// Returns components for the current entity that match the query `Q`.
+    pub fn get_components_mut<Q: QueryData>(&mut self) -> Option<Q::Item<'_>> {
+        // SAFETY: &mut self implies exclusive access for duration of returned value
+        unsafe { self.0.get_components::<Q>() }
+    }
+
     /// Consumes `self` and gets access to the component of type `T` with the
     /// world `'w` lifetime for the current entity.
     ///
@@ -645,6 +678,31 @@ impl<'w> EntityWorldMut<'w> {
         EntityRef::from(self).get()
     }
 
+    /// Returns read-only components for the current entity that match the query `Q`.
+    #[inline]
+    pub fn components<Q: ReadOnlyQueryData>(&self) -> Q::Item<'_> {
+        EntityRef::from(self).components::<Q>()
+    }
+
+    /// Returns components for the current entity that match the query `Q`.
+    #[inline]
+    pub fn components_mut<Q: QueryData>(&mut self) -> Q::Item<'_> {
+        self.get_components_mut::<Q>().expect(QUERY_MISMATCH_ERROR)
+    }
+
+    /// Returns read-only components for the current entity that match the query `Q`.
+    #[inline]
+    pub fn get_components<Q: ReadOnlyQueryData>(&self) -> Option<Q::Item<'_>> {
+        EntityRef::from(self).get_components::<Q>()
+    }
+
+    /// Returns components for the current entity that match the query `Q`.
+    #[inline]
+    pub fn get_components_mut<Q: QueryData>(&mut self) -> Option<Q::Item<'_>> {
+        // SAFETY: &mut self implies exclusive access for duration of returned value
+        unsafe { self.as_unsafe_entity_cell().get_components::<Q>() }
+    }
+
     /// Consumes `self` and gets access to the component of type `T` with
     /// the world `'w` lifetime for the current entity.
     /// Returns `None` if the entity does not have a component of type `T`.
@@ -1383,6 +1441,8 @@ impl<'w> EntityWorldMut<'w> {
     }
 }
 
+const QUERY_MISMATCH_ERROR: &str = "Query does not match the current entity";
+
 /// A view into a single entity and component in a world, which may either be vacant or occupied.
 ///
 /// This `enum` can only be constructed from the [`entry`] method on [`EntityWorldMut`].
@@ -1770,7 +1830,7 @@ impl<'w> FilteredEntityRef<'w> {
 
     /// Returns an iterator over the component ids that are accessed by self.
     #[inline]
-    pub fn components(&self) -> impl Iterator<Item = ComponentId> + '_ {
+    pub fn accessed_components(&self) -> impl Iterator<Item = ComponentId> + '_ {
         self.access.reads_and_writes()
     }
 
@@ -2024,7 +2084,7 @@ impl<'w> FilteredEntityMut<'w> {
 
     /// Returns an iterator over the component ids that are accessed by self.
     #[inline]
-    pub fn components(&self) -> impl Iterator<Item = ComponentId> + '_ {
+    pub fn accessed_components(&self) -> impl Iterator<Item = ComponentId> + '_ {
         self.access.reads_and_writes()
     }
 
@@ -2918,4 +2978,24 @@ mod tests {
 
         assert_is_system(incompatible_system);
     }
+
+    #[test]
+    fn get_components() {
+        #[derive(Component, PartialEq, Eq, Debug)]
+        struct X(usize);
+
+        #[derive(Component, PartialEq, Eq, Debug)]
+        struct Y(usize);
+        let mut world = World::default();
+        let e1 = world.spawn((X(7), Y(10))).id();
+        let e2 = world.spawn(X(8)).id();
+        let e3 = world.spawn_empty().id();
+
+        assert_eq!(
+            Some((&X(7), &Y(10))),
+            world.entity(e1).get_components::<(&X, &Y)>()
+        );
+        assert_eq!(None, world.entity(e2).get_components::<(&X, &Y)>());
+        assert_eq!(None, world.entity(e3).get_components::<(&X, &Y)>());
+    }
 }

crates/bevy_ecs/src/world/unsafe_world_cell.rs

@@ -10,6 +10,7 @@ use crate::{
     component::{ComponentId, ComponentTicks, Components, StorageType, Tick, TickCells},
     entity::{Entities, Entity, EntityLocation},
     prelude::Component,
+    query::{DebugCheckedUnwrap, QueryData},
     removal_detection::RemovedComponentEvents,
     storage::{Column, ComponentSparseSet, Storages},
     system::{Res, Resource},
@@ -818,6 +819,48 @@ impl<'w> UnsafeEntityCell<'w> {
             })
         }
     }
+
+    /// # Safety
+    /// It is the callers responsibility to ensure that
+    /// - the [`UnsafeEntityCell`] has permission to access the queried data mutably
+    /// - no other references to the queried data exist at the same time
+    pub(crate) unsafe fn get_components<Q: QueryData>(&self) -> Option<Q::Item<'w>> {
+        let state = Q::get_state(self.world().components())?;
+        let location = self.location();
+        // SAFETY: Location is guaranteed to exist
+        let archetype = unsafe {
+            self.world
+                .archetypes()
+                .get(location.archetype_id)
+                .debug_checked_unwrap()
+        };
+        if Q::matches_component_set(&state, &|id| archetype.contains(id)) {
+            // SAFETY: state was initialized above using the world passed into this function
+            let mut fetch = unsafe {
+                Q::init_fetch(
+                    self.world,
+                    &state,
+                    self.world.last_change_tick(),
+                    self.world.change_tick(),
+                )
+            };
+            // SAFETY: Table is guaranteed to exist
+            let table = unsafe {
+                self.world
+                    .storages()
+                    .tables
+                    .get(location.table_id)
+                    .debug_checked_unwrap()
+            };
+            // SAFETY: Archetype and table are from the same world used to initialize state and fetch.
+            // Table corresponds to archetype. State is the same state used to init fetch above.
+            unsafe { Q::set_archetype(&mut fetch, &state, archetype, table) }
+            // SAFETY: Called after set_archetype above. Entity and location are guaranteed to exist.
+            unsafe { Some(Q::fetch(&mut fetch, self.id(), location.table_row)) }
+        } else {
+            None
+        }
+    }
 }
 
 impl<'w> UnsafeEntityCell<'w> {

examples/ecs/dynamic.rs

@@ -151,7 +151,7 @@ fn main() {
 
                 query.iter_mut(&mut world).for_each(|filtered_entity| {
                     let terms = filtered_entity
-                        .components()
+                        .accessed_components()
                         .map(|id| {
                             let ptr = filtered_entity.get_by_id(id).unwrap();
                             let info = component_info.get(&id).unwrap();