diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index e3050339916ed..e60bfbfd7dc41 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -2,9 +2,12 @@ use crate::{ archetype::{Archetype, ArchetypeId, Archetypes}, bundle::{Bundle, BundleInfo}, change_detection::{MutUntyped, Ticks}, - component::{Component, ComponentId, ComponentTicks, Components, StorageType, TickCells}, + component::{ + Component, ComponentId, ComponentStorage, ComponentTicks, Components, StorageType, + TickCells, + }, entity::{Entities, Entity, EntityLocation}, - storage::{SparseSet, Storages}, + storage::{Column, ComponentSparseSet, SparseSet, Storages}, world::{Mut, World}, }; use bevy_ptr::{OwningPtr, Ptr}; @@ -67,10 +70,19 @@ impl<'w> EntityRef<'w> { #[inline] pub fn get(&self) -> Option<&'w T> { - // SAFETY: entity location is valid and returned component is of type T + // SAFETY: + // - entity location and entity is valid + // - returned component is of type T + // - the storage type provided is correct for T unsafe { - get_component_with_type(self.world, TypeId::of::(), self.entity, self.location) - .map(|value| value.deref::()) + get_component_with_type( + self.world, + TypeId::of::(), + T::Storage::STORAGE_TYPE, + self.entity, + self.location, + ) + .map(|value| value.deref::()) } } @@ -78,8 +90,18 @@ impl<'w> EntityRef<'w> { /// detection in custom runtimes. #[inline] pub fn get_change_ticks(&self) -> Option { - // SAFETY: entity location is valid - unsafe { get_ticks_with_type(self.world, TypeId::of::(), self.entity, self.location) } + // SAFETY: + // - entity location and entity is valid + // - the storage type provided is correct for T + unsafe { + get_ticks_with_type( + self.world, + TypeId::of::(), + T::Storage::STORAGE_TYPE, + self.entity, + self.location, + ) + } } /// Retrieves the change ticks for the given [`ComponentId`]. This can be useful for implementing change @@ -94,8 +116,17 @@ impl<'w> EntityRef<'w> { return None; } + let info = self.world.components().get_info(component_id)?; // SAFETY: Entity location is valid and component_id exists. - unsafe { get_ticks(self.world, component_id, self.entity, self.location) } + unsafe { + get_ticks( + self.world, + component_id, + info.storage_type(), + self.entity, + self.location, + ) + } } /// Gets a mutable reference to the component of type `T` associated with @@ -114,11 +145,21 @@ impl<'w> EntityRef<'w> { last_change_tick: u32, change_tick: u32, ) -> Option> { - get_component_and_ticks_with_type(self.world, TypeId::of::(), self.entity, self.location) - .map(|(value, ticks)| Mut { - value: value.assert_unique().deref_mut::(), - ticks: Ticks::from_tick_cells(ticks, last_change_tick, change_tick), - }) + // SAFETY: + // - entity location and entity is valid + // - returned component is of type T + // - the storage type provided is correct for T + get_component_and_ticks_with_type( + self.world, + TypeId::of::(), + T::Storage::STORAGE_TYPE, + self.entity, + self.location, + ) + .map(|(value, ticks)| Mut { + value: value.assert_unique().deref_mut::(), + ticks: Ticks::from_tick_cells(ticks, last_change_tick, change_tick), + }) } } @@ -133,9 +174,20 @@ impl<'w> EntityRef<'w> { /// which is only valid while the `'w` borrow of the lifetime is active. #[inline] pub fn get_by_id(&self, component_id: ComponentId) -> Option> { - self.world.components().get_info(component_id)?; - // SAFETY: entity_location is valid, component_id is valid as checked by the line above - unsafe { get_component(self.world, component_id, self.entity, self.location) } + let info = self.world.components().get_info(component_id)?; + // SAFETY: + // - entity_location is valid, + // - component_id is valid as checked by the line above + // - the storage type is accurate as checked by the fetched ComponentInfo + unsafe { + get_component( + self.world, + component_id, + info.storage_type(), + self.entity, + self.location, + ) + } } } @@ -201,10 +253,20 @@ impl<'w> EntityMut<'w> { #[inline] pub fn get(&self) -> Option<&'_ T> { - // SAFETY: lifetimes enforce correct usage of returned borrow + // SAFETY: + // - lifetimes enforce correct usage of returned borrow + // - entity location and entity is valid + // - returned component is of type T + // - the storage type provided is correct for T unsafe { - get_component_with_type(self.world, TypeId::of::(), self.entity, self.location) - .map(|value| value.deref::()) + get_component_with_type( + self.world, + TypeId::of::(), + T::Storage::STORAGE_TYPE, + self.entity, + self.location, + ) + .map(|value| value.deref::()) } } @@ -218,8 +280,18 @@ impl<'w> EntityMut<'w> { /// detection in custom runtimes. #[inline] pub fn get_change_ticks(&self) -> Option { - // SAFETY: entity location is valid - unsafe { get_ticks_with_type(self.world, TypeId::of::(), self.entity, self.location) } + // SAFETY: + // - entity location and entity is valid + // - the storage type provided is correct for T + unsafe { + get_ticks_with_type( + self.world, + TypeId::of::(), + T::Storage::STORAGE_TYPE, + self.entity, + self.location, + ) + } } /// Retrieves the change ticks for the given [`ComponentId`]. This can be useful for implementing change @@ -234,8 +306,17 @@ impl<'w> EntityMut<'w> { return None; } + let info = self.world.components().get_info(component_id)?; // SAFETY: Entity location is valid and component_id exists. - unsafe { get_ticks(self.world, component_id, self.entity, self.location) } + unsafe { + get_ticks( + self.world, + component_id, + info.storage_type(), + self.entity, + self.location, + ) + } } /// Gets a mutable reference to the component of type `T` associated with @@ -250,15 +331,25 @@ impl<'w> EntityMut<'w> { /// operation on this world (non-exhaustive list). #[inline] pub unsafe fn get_unchecked_mut(&self) -> Option> { - get_component_and_ticks_with_type(self.world, TypeId::of::(), self.entity, self.location) - .map(|(value, ticks)| Mut { - value: value.assert_unique().deref_mut::(), - ticks: Ticks::from_tick_cells( - ticks, - self.world.last_change_tick(), - self.world.read_change_tick(), - ), - }) + // SAFETY: + // - entity location and entity is valid + // - returned component is of type T + // - the storage type provided is correct for T + get_component_and_ticks_with_type( + self.world, + TypeId::of::(), + T::Storage::STORAGE_TYPE, + self.entity, + self.location, + ) + .map(|(value, ticks)| Mut { + value: value.assert_unique().deref_mut::(), + ticks: Ticks::from_tick_cells( + ticks, + self.world.last_change_tick(), + self.world.read_change_tick(), + ), + }) } /// Adds a [`Bundle`] of components to the entity. @@ -573,9 +664,20 @@ impl<'w> EntityMut<'w> { /// which is only valid while the [`EntityMut`] is alive. #[inline] pub fn get_by_id(&self, component_id: ComponentId) -> Option> { - self.world.components().get_info(component_id)?; - // SAFETY: entity_location is valid, component_id is valid as checked by the line above - unsafe { get_component(self.world, component_id, self.entity, self.location) } + let info = self.world.components().get_info(component_id)?; + // SAFETY: + // - entity_location is valid + // - component_id is valid as checked by the line above + // - the storage type is accurate as checked by the fetched ComponentInfo + unsafe { + get_component( + self.world, + component_id, + info.storage_type(), + self.entity, + self.location, + ) + } } /// Gets a [`MutUntyped`] of the component of the given [`ComponentId`] from the entity. @@ -594,6 +696,24 @@ impl<'w> EntityMut<'w> { } } +#[inline] +fn fetch_table( + world: &World, + location: EntityLocation, + component_id: ComponentId, +) -> Option<(&Column, usize)> { + let archetype = &world.archetypes[location.archetype_id]; + let table = &world.storages.tables[archetype.table_id()]; + let components = table.get_column(component_id)?; + let table_row = archetype.entity_table_row(location.index); + Some((components, table_row)) +} + +#[inline] +fn fetch_sparse_set(world: &World, component_id: ComponentId) -> Option<&ComponentSparseSet> { + world.storages.sparse_sets.get(component_id) +} + // TODO: move to Storages? /// Get a raw pointer to a particular [`Component`] on a particular [`Entity`] in the provided [`World`]. /// @@ -601,29 +721,22 @@ impl<'w> EntityMut<'w> { /// - `entity_location` must be within bounds of the given archetype and `entity` must exist inside /// the archetype /// - `component_id` must be valid +/// - `storage_type` must accurately reflect where the components for `component_id` are stored. #[inline] pub(crate) unsafe fn get_component( world: &World, component_id: ComponentId, + storage_type: StorageType, entity: Entity, location: EntityLocation, ) -> Option> { - let archetype = &world.archetypes[location.archetype_id]; - // SAFETY: component_id exists and is therefore valid - let component_info = world.components.get_info_unchecked(component_id); - match component_info.storage_type() { + match storage_type { StorageType::Table => { - let table = &world.storages.tables[archetype.table_id()]; - let components = table.get_column(component_id)?; - let table_row = archetype.entity_table_row(location.index); + let (components, table_row) = fetch_table(world, location, component_id)?; // SAFETY: archetypes only store valid table_rows and the stored component type is T Some(components.get_data_unchecked(table_row)) } - StorageType::SparseSet => world - .storages - .sparse_sets - .get(component_id) - .and_then(|sparse_set| sparse_set.get(entity)), + StorageType::SparseSet => fetch_sparse_set(world, component_id)?.get(entity), } } @@ -631,21 +744,21 @@ pub(crate) unsafe fn get_component( /// Get a raw pointer to the [`ComponentTicks`] of a particular [`Component`] on a particular [`Entity`] in the provided [World]. /// /// # Safety -/// Caller must ensure that `component_id` is valid +/// - `entity_location` must be within bounds of the given archetype and `entity` must exist inside +/// the archetype +/// - `component_id` must be valid +/// - `storage_type` must accurately reflect where the components for `component_id` are stored. #[inline] unsafe fn get_component_and_ticks( world: &World, component_id: ComponentId, + storage_type: StorageType, entity: Entity, location: EntityLocation, ) -> Option<(Ptr<'_>, TickCells<'_>)> { - let archetype = &world.archetypes[location.archetype_id]; - let component_info = world.components.get_info_unchecked(component_id); - match component_info.storage_type() { + match storage_type { StorageType::Table => { - let table = &world.storages.tables[archetype.table_id()]; - let components = table.get_column(component_id)?; - let table_row = archetype.entity_table_row(location.index); + let (components, table_row) = fetch_table(world, location, component_id)?; // SAFETY: archetypes only store valid table_rows and the stored component type is T Some(( components.get_data_unchecked(table_row), @@ -655,36 +768,29 @@ unsafe fn get_component_and_ticks( }, )) } - StorageType::SparseSet => world - .storages - .sparse_sets - .get(component_id) - .and_then(|sparse_set| sparse_set.get_with_ticks(entity)), + StorageType::SparseSet => fetch_sparse_set(world, component_id)?.get_with_ticks(entity), } } -#[inline] +/// # Safety +/// - `entity_location` must be within bounds of the given archetype and `entity` must exist inside +/// the archetype +/// - `component_id` must be valid +/// - `storage_type` must accurately reflect where the components for `component_id` are stored. unsafe fn get_ticks( world: &World, component_id: ComponentId, + storage_type: StorageType, entity: Entity, location: EntityLocation, ) -> Option { - let archetype = &world.archetypes[location.archetype_id]; - let component_info = world.components.get_info_unchecked(component_id); - match component_info.storage_type() { + match storage_type { StorageType::Table => { - let table = &world.storages.tables[archetype.table_id()]; - let components = table.get_column(component_id)?; - let table_row = archetype.entity_table_row(location.index); + let (components, table_row) = fetch_table(world, location, component_id)?; // SAFETY: archetypes only store valid table_rows and the stored component type is T Some(components.get_ticks_unchecked(table_row)) } - StorageType::SparseSet => world - .storages - .sparse_sets - .get(component_id) - .and_then(|sparse_set| sparse_set.get_ticks(entity)), + StorageType::SparseSet => fetch_sparse_set(world, component_id)?.get_ticks(entity), } } @@ -733,41 +839,71 @@ unsafe fn take_component<'a>( /// Get a raw pointer to a particular [`Component`] by [`TypeId`] on a particular [`Entity`] in the provided [`World`]. /// /// # Safety -/// `entity_location` must be within bounds of an archetype that exists. +/// - `entity_location` must be within bounds of the given archetype and `entity` must exist inside +/// the archetype +/// - `type_id` must be correspond to a type that implements [`Component`] +/// - `storage_type` must accurately reflect where the components for `component_id` are stored. +#[inline] unsafe fn get_component_with_type( world: &World, type_id: TypeId, + storage_type: StorageType, entity: Entity, location: EntityLocation, ) -> Option> { - let component_id = world.components.get_id(type_id)?; - get_component(world, component_id, entity, location) + get_component( + world, + world.components.get_id(type_id)?, + storage_type, + entity, + location, + ) } /// Get a raw pointer to the [`ComponentTicks`] of a particular [`Component`] by [`TypeId`] on a particular [`Entity`] in the provided [`World`]. /// /// # Safety -/// `entity_location` must be within bounds of an archetype that exists. -pub(crate) unsafe fn get_component_and_ticks_with_type( +/// - `entity_location` must be within bounds of the given archetype and `entity` must exist inside +/// the archetype +/// - `type_id` must be correspond to a type that implements [`Component`] +/// - `storage_type` must accurately reflect where the components for `component_id` are stored. +#[inline] +unsafe fn get_component_and_ticks_with_type( world: &World, type_id: TypeId, + storage_type: StorageType, entity: Entity, location: EntityLocation, ) -> Option<(Ptr<'_>, TickCells<'_>)> { - let component_id = world.components.get_id(type_id)?; - get_component_and_ticks(world, component_id, entity, location) + get_component_and_ticks( + world, + world.components.get_id(type_id)?, + storage_type, + entity, + location, + ) } /// # Safety -/// `entity_location` must be within bounds of an archetype that exists. -pub(crate) unsafe fn get_ticks_with_type( +/// - `entity_location` must be within bounds of the given archetype and `entity` must exist inside +/// the archetype +/// - `type_id` must be correspond to a type that implements [`Component`] +/// - `storage_type` must accurately reflect where the components for `component_id` are stored. +#[inline] +unsafe fn get_ticks_with_type( world: &World, type_id: TypeId, + storage_type: StorageType, entity: Entity, location: EntityLocation, ) -> Option { - let component_id = world.components.get_id(type_id)?; - get_ticks(world, component_id, entity, location) + get_ticks( + world, + world.components.get_id(type_id)?, + storage_type, + entity, + location, + ) } fn contains_component_with_type(world: &World, type_id: TypeId, location: EntityLocation) -> bool { @@ -914,12 +1050,17 @@ pub(crate) unsafe fn get_mut( // T let change_tick = world.change_tick(); let last_change_tick = world.last_change_tick(); - get_component_and_ticks_with_type(world, TypeId::of::(), entity, location).map( - |(value, ticks)| Mut { - value: value.assert_unique().deref_mut::(), - ticks: Ticks::from_tick_cells(ticks, last_change_tick, change_tick), - }, + get_component_and_ticks_with_type( + world, + TypeId::of::(), + T::Storage::STORAGE_TYPE, + entity, + location, ) + .map(|(value, ticks)| Mut { + value: value.assert_unique().deref_mut::(), + ticks: Ticks::from_tick_cells(ticks, last_change_tick, change_tick), + }) } // SAFETY: EntityLocation must be valid, component_id must be valid @@ -931,13 +1072,14 @@ pub(crate) unsafe fn get_mut_by_id( component_id: ComponentId, ) -> Option { let change_tick = world.change_tick(); + let info = world.components.get_info_unchecked(component_id); // SAFETY: world access is unique, entity location and component_id required to be valid - get_component_and_ticks(world, component_id, entity, location).map(|(value, ticks)| { - MutUntyped { + get_component_and_ticks(world, component_id, info.storage_type(), entity, location).map( + |(value, ticks)| MutUntyped { value: value.assert_unique(), ticks: Ticks::from_tick_cells(ticks, world.last_change_tick(), change_tick), - } - }) + }, + ) } #[cfg(test)] diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index c0ca9d4b85301..435694f23c2c2 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -1507,12 +1507,16 @@ impl World { /// use this in cases where the actual types are not known at compile time.** #[inline] pub fn get_by_id(&self, entity: Entity, component_id: ComponentId) -> Option> { - self.components().get_info(component_id)?; - // SAFETY: entity_location is valid, component_id is valid as checked by the line above + let info = self.components().get_info(component_id)?; + // SAFETY: + // - entity_location is valid + // - component_id is valid as checked by the line above + // - the storage type is accurate as checked by the fetched ComponentInfo unsafe { get_component( self, component_id, + info.storage_type(), entity, self.get_entity(entity)?.location(), )