diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index 1308198186c40..5e21f6b9987a7 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, removal_detection::RemovedComponentEvents, storage::Storages, world::{Mut, World}, @@ -1824,8 +1824,9 @@ impl<'w> FilteredEntityRef<'w> { let id = self.entity.world().components().get_id(TypeId::of::())?; self.access .has_read(id) - // SAFETY: We have read access so we must have the component - .then(|| unsafe { self.entity.get().debug_checked_unwrap() }) + // SAFETY: We have read access + .then(|| unsafe { self.entity.get() }) + .flatten() } /// Gets access to the component of type `T` for the current entity, @@ -1837,8 +1838,9 @@ impl<'w> FilteredEntityRef<'w> { let id = self.entity.world().components().get_id(TypeId::of::())?; self.access .has_read(id) - // SAFETY: We have read access so we must have the component - .then(|| unsafe { self.entity.get_ref().debug_checked_unwrap() }) + // SAFETY: We have read access + .then(|| unsafe { self.entity.get_ref() }) + .flatten() } /// Retrieves the change ticks for the given component. This can be useful for implementing change @@ -1848,8 +1850,9 @@ impl<'w> FilteredEntityRef<'w> { let id = self.entity.world().components().get_id(TypeId::of::())?; self.access .has_read(id) - // SAFETY: We have read access so we must have the component - .then(|| unsafe { self.entity.get_change_ticks::().debug_checked_unwrap() }) + // SAFETY: We have read access + .then(|| unsafe { self.entity.get_change_ticks::() }) + .flatten() } /// Retrieves the change ticks for the given [`ComponentId`]. This can be useful for implementing change @@ -1860,12 +1863,11 @@ impl<'w> FilteredEntityRef<'w> { /// compile time.** #[inline] pub fn get_change_ticks_by_id(&self, component_id: ComponentId) -> Option { - // SAFETY: We have read access so we must have the component - self.access.has_read(component_id).then(|| unsafe { - self.entity - .get_change_ticks_by_id(component_id) - .debug_checked_unwrap() - }) + self.access + .has_read(component_id) + // SAFETY: We have read access + .then(|| unsafe { self.entity.get_change_ticks_by_id(component_id) }) + .flatten() } /// Gets the component of the given [`ComponentId`] from the entity. @@ -1880,8 +1882,9 @@ impl<'w> FilteredEntityRef<'w> { pub fn get_by_id(&self, component_id: ComponentId) -> Option> { self.access .has_read(component_id) - // SAFETY: We have read access so we must have the component - .then(|| unsafe { self.entity.get_by_id(component_id).debug_checked_unwrap() }) + // SAFETY: We have read access + .then(|| unsafe { self.entity.get_by_id(component_id) }) + .flatten() } } @@ -2094,8 +2097,9 @@ impl<'w> FilteredEntityMut<'w> { let id = self.entity.world().components().get_id(TypeId::of::())?; self.access .has_write(id) - // SAFETY: We have write access so we must have the component - .then(|| unsafe { self.entity.get_mut().debug_checked_unwrap() }) + // SAFETY: We have write access + .then(|| unsafe { self.entity.get_mut() }) + .flatten() } /// Retrieves the change ticks for the given component. This can be useful for implementing change @@ -2139,12 +2143,11 @@ impl<'w> FilteredEntityMut<'w> { /// which is only valid while the [`FilteredEntityMut`] is alive. #[inline] pub fn get_mut_by_id(&mut self, component_id: ComponentId) -> Option> { - // SAFETY: We have write access so we must have the component - self.access.has_write(component_id).then(|| unsafe { - self.entity - .get_mut_by_id(component_id) - .debug_checked_unwrap() - }) + self.access + .has_write(component_id) + // SAFETY: We have write access + .then(|| unsafe { self.entity.get_mut_by_id(component_id) }) + .flatten() } } @@ -2416,6 +2419,7 @@ mod tests { use bevy_ptr::OwningPtr; use std::panic::AssertUnwindSafe; + use crate::world::{FilteredEntityMut, FilteredEntityRef}; use crate::{self as bevy_ecs, component::ComponentId, prelude::*, system::assert_is_system}; #[test] @@ -2918,4 +2922,64 @@ mod tests { assert_is_system(incompatible_system); } + + #[test] + fn filtered_entity_ref_normal() { + let mut world = World::new(); + let a_id = world.init_component::(); + + let e: FilteredEntityRef = world.spawn(A).into(); + + assert!(e.get::().is_some()); + assert!(e.get_ref::().is_some()); + assert!(e.get_change_ticks::().is_some()); + assert!(e.get_by_id(a_id).is_some()); + assert!(e.get_change_ticks_by_id(a_id).is_some()); + } + + #[test] + fn filtered_entity_ref_missing() { + let mut world = World::new(); + let a_id = world.init_component::(); + + let e: FilteredEntityRef = world.spawn(()).into(); + + assert!(e.get::().is_none()); + assert!(e.get_ref::().is_none()); + assert!(e.get_change_ticks::().is_none()); + assert!(e.get_by_id(a_id).is_none()); + assert!(e.get_change_ticks_by_id(a_id).is_none()); + } + + #[test] + fn filtered_entity_mut_normal() { + let mut world = World::new(); + let a_id = world.init_component::(); + + let mut e: FilteredEntityMut = world.spawn(A).into(); + + assert!(e.get::().is_some()); + assert!(e.get_ref::().is_some()); + assert!(e.get_mut::().is_some()); + assert!(e.get_change_ticks::().is_some()); + assert!(e.get_by_id(a_id).is_some()); + assert!(e.get_mut_by_id(a_id).is_some()); + assert!(e.get_change_ticks_by_id(a_id).is_some()); + } + + #[test] + fn filtered_entity_mut_missing() { + let mut world = World::new(); + let a_id = world.init_component::(); + + let mut e: FilteredEntityMut = world.spawn(()).into(); + + assert!(e.get::().is_none()); + assert!(e.get_ref::().is_none()); + assert!(e.get_mut::().is_none()); + assert!(e.get_change_ticks::().is_none()); + assert!(e.get_by_id(a_id).is_none()); + assert!(e.get_mut_by_id(a_id).is_none()); + assert!(e.get_change_ticks_by_id(a_id).is_none()); + } }