diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index 1308198186c40..5a6c2f8df9339 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/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(&self) -> Q::Item<'w> { + self.get_components::().expect(QUERY_MISMATCH_ERROR) + } + + /// Returns read-only components for the current entity that match the query `Q`. + pub fn get_components(&self) -> Option> { + // SAFETY: We have read-only access to all components of this entity. + unsafe { self.0.get_components::() } + } } impl<'w> From> 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(&self) -> Q::Item<'_> { + self.get_components::().expect(QUERY_MISMATCH_ERROR) + } + + /// Returns components for the current entity that match the query `Q`. + pub fn components_mut(&mut self) -> Q::Item<'_> { + self.get_components_mut::().expect(QUERY_MISMATCH_ERROR) + } + + /// Returns read-only components for the current entity that match the query `Q`. + pub fn get_components(&self) -> Option> { + // SAFETY: We have read-only access to all components of this entity. + unsafe { self.0.get_components::() } + } + + /// Returns components for the current entity that match the query `Q`. + pub fn get_components_mut(&mut self) -> Option> { + // SAFETY: &mut self implies exclusive access for duration of returned value + unsafe { self.0.get_components::() } + } + /// 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(&self) -> Q::Item<'_> { + EntityRef::from(self).components::() + } + + /// Returns components for the current entity that match the query `Q`. + #[inline] + pub fn components_mut(&mut self) -> Q::Item<'_> { + self.get_components_mut::().expect(QUERY_MISMATCH_ERROR) + } + + /// Returns read-only components for the current entity that match the query `Q`. + #[inline] + pub fn get_components(&self) -> Option> { + EntityRef::from(self).get_components::() + } + + /// Returns components for the current entity that match the query `Q`. + #[inline] + pub fn get_components_mut(&mut self) -> Option> { + // SAFETY: &mut self implies exclusive access for duration of returned value + unsafe { self.as_unsafe_entity_cell().get_components::() } + } + /// 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 + '_ { + pub fn accessed_components(&self) -> impl Iterator + '_ { 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 + '_ { + pub fn accessed_components(&self) -> impl Iterator + '_ { 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)>()); + } } diff --git a/crates/bevy_ecs/src/world/unsafe_world_cell.rs b/crates/bevy_ecs/src/world/unsafe_world_cell.rs index 429a9658fdf78..16f33a13affcb 100644 --- a/crates/bevy_ecs/src/world/unsafe_world_cell.rs +++ b/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(&self) -> Option> { + 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> { diff --git a/examples/ecs/dynamic.rs b/examples/ecs/dynamic.rs index c804062aefe91..b3cc6c2b2770d 100644 --- a/examples/ecs/dynamic.rs +++ b/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();