Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

EntityRef/Mut get_components (immutable variants only) #15089

Merged
merged 9 commits into from
Sep 9, 2024
77 changes: 74 additions & 3 deletions crates/bevy_ecs/src/world/entity_ref.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use crate::{
entity::{Entities, Entity, EntityLocation},
event::Event,
observer::{Observer, Observers},
query::Access,
query::{Access, ReadOnlyQueryData},
removal_detection::RemovedComponentEvents,
storage::Storages,
system::IntoObserverSystem,
Expand Down Expand Up @@ -156,6 +156,22 @@ 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`.
///
/// # Panics
///
/// If the entity does not have the components required by 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`,
/// or `None` if the entity does not have the components required by 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> {
Expand Down Expand Up @@ -351,6 +367,22 @@ impl<'w> EntityMut<'w> {
self.as_readonly().get()
}

/// Returns read-only components for the current entity that match the query `Q`.
///
/// # Panics
///
/// If the entity does not have the components required by the query `Q`.
pub fn components<Q: ReadOnlyQueryData>(&self) -> Q::Item<'_> {
self.get_components::<Q>().expect(QUERY_MISMATCH_ERROR)
}

/// Returns read-only components for the current entity that match the query `Q`,
/// or `None` if the entity does not have the components required by 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>() }
}

/// Consumes `self` and gets access to the component of type `T` with the
/// world `'w` lifetime for the current entity.
///
Expand Down Expand Up @@ -648,6 +680,23 @@ impl<'w> EntityWorldMut<'w> {
EntityRef::from(self).get()
}

/// Returns read-only components for the current entity that match the query `Q`.
///
/// # Panics
///
/// If the entity does not have the components required by the query `Q`.
#[inline]
pub fn components<Q: ReadOnlyQueryData>(&self) -> Q::Item<'_> {
EntityRef::from(self).components::<Q>()
}

/// Returns read-only components for the current entity that match the query `Q`,
/// or `None` if the entity does not have the components required by the query `Q`.
#[inline]
pub fn get_components<Q: ReadOnlyQueryData>(&self) -> Option<Q::Item<'_>> {
EntityRef::from(self).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`.
Expand Down Expand Up @@ -1491,6 +1540,8 @@ unsafe fn trigger_on_replace_and_on_remove_hooks_and_observers(
}
}

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`].
Expand Down Expand Up @@ -1878,7 +1929,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.component_reads_and_writes()
}

Expand Down Expand Up @@ -2135,7 +2186,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.component_reads_and_writes()
}

Expand Down Expand Up @@ -3115,4 +3166,24 @@ mod tests {
assert!(e.get_mut_by_id(a_id).is_none());
assert!(e.get_change_ticks_by_id(a_id).is_none());
}

#[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)>());
}
}
50 changes: 50 additions & 0 deletions crates/bevy_ecs/src/world/unsafe_world_cell.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ use crate::{
entity::{Entities, Entity, EntityLocation},
observer::Observers,
prelude::Component,
query::{DebugCheckedUnwrap, ReadOnlyQueryData},
removal_detection::RemovedComponentEvents,
storage::{Column, ComponentSparseSet, Storages},
system::{Res, Resource},
Expand Down Expand Up @@ -882,6 +883,55 @@ impl<'w> UnsafeEntityCell<'w> {
})
}
}

/// Returns read-only components for the current entity that match the query `Q`,
/// or `None` if the entity does not have the components required by the query `Q`.
///
/// # Safety
/// It is the callers responsibility to ensure that
/// - the [`UnsafeEntityCell`] has permission to access the queried data immutably
/// - no mutable references to the queried data exist at the same time
pub(crate) unsafe fn get_components<Q: ReadOnlyQueryData>(&self) -> Option<Q::Item<'w>> {
// SAFETY: World is only used to access query data and initialize query state
let state = unsafe {
let world = self.world().world();
Q::get_state(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> {
Expand Down
2 changes: 1 addition & 1 deletion examples/ecs/dynamic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down