diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs index 3f942299f3142f..1dfef36c96c86a 100644 --- a/crates/bevy_ecs/src/storage/sparse_set.rs +++ b/crates/bevy_ecs/src/storage/sparse_set.rs @@ -6,6 +6,8 @@ use crate::{ use bevy_ptr::{OwningPtr, Ptr}; use std::{cell::UnsafeCell, hash::Hash, marker::PhantomData}; +type EntityId = u32; + #[derive(Debug)] pub struct SparseArray { values: Vec>, @@ -96,8 +98,14 @@ impl SparseArray { pub struct ComponentSparseSet { dense: BlobVec, ticks: Vec>, + // Internally this only relies on the Entity ID to keep track of where the component data is + // stored for entities that are alive. The generation is not required, but is stored + // in debug builds to validate that access is correct. + #[cfg(not(debug_assertions))] + entities: Vec, + #[cfg(debug_assertions)] entities: Vec, - sparse: SparseArray, + sparse: SparseArray, } impl ComponentSparseSet { @@ -137,38 +145,55 @@ impl ComponentSparseSet { /// The `value` pointer must point to a valid address that matches the [`Layout`](std::alloc::Layout) /// inside the [`ComponentInfo`] given when constructing this sparse set. pub unsafe fn insert(&mut self, entity: Entity, value: OwningPtr<'_>, change_tick: u32) { - if let Some(&dense_index) = self.sparse.get(entity) { - self.dense.replace_unchecked(dense_index, value); - *self.ticks.get_unchecked_mut(dense_index) = + if let Some(&dense_index) = self.sparse.get(entity.id()) { + debug_assert_eq!(entity, self.entities[dense_index as usize]); + let _entity = self.dense.replace_unchecked(dense_index as usize, value); + *self.ticks.get_unchecked_mut(dense_index as usize) = UnsafeCell::new(ComponentTicks::new(change_tick)); } else { let dense_index = self.dense.len(); self.dense.push(value); - self.sparse.insert(entity, dense_index); + self.sparse.insert(entity.id(), dense_index as u32); debug_assert_eq!(self.ticks.len(), dense_index); debug_assert_eq!(self.entities.len(), dense_index); self.ticks .push(UnsafeCell::new(ComponentTicks::new(change_tick))); + #[cfg(not(debug_assertions))] + self.entities.push(entity.id()); + #[cfg(debug_assertions)] self.entities.push(entity); } } #[inline] pub fn contains(&self, entity: Entity) -> bool { - self.sparse.contains(entity) + #[cfg(debug_assertions)] + { + if let Some(&dense_index) = self.sparse.get(entity.id()) { + debug_assert_eq!(entity, self.entities[dense_index as usize]); + true + } else { + false + } + } + #[cfg(not(debug_assertions))] + self.sparse.contains(entity.id()) } #[inline] pub fn get(&self, entity: Entity) -> Option> { - self.sparse.get(entity).map(|dense_index| { + self.sparse.get(entity.id()).map(|dense_index| { + let dense_index = *dense_index as usize; + debug_assert_eq!(entity, self.entities[dense_index]); // SAFE: if the sparse index points to something in the dense vec, it exists - unsafe { self.dense.get_unchecked(*dense_index) } + unsafe { self.dense.get_unchecked(dense_index) } }) } #[inline] pub fn get_with_ticks(&self, entity: Entity) -> Option<(Ptr<'_>, &UnsafeCell)> { - let dense_index = *self.sparse.get(entity)?; + let dense_index = *self.sparse.get(entity.id())? as usize; + debug_assert_eq!(entity, self.entities[dense_index]); // SAFE: if the sparse index points to something in the dense vec, it exists unsafe { Some(( @@ -180,7 +205,8 @@ impl ComponentSparseSet { #[inline] pub fn get_ticks(&self, entity: Entity) -> Option<&UnsafeCell> { - let dense_index = *self.sparse.get(entity)?; + let dense_index = *self.sparse.get(entity.id())? as usize; + debug_assert_eq!(entity, self.entities[dense_index]); // SAFE: if the sparse index points to something in the dense vec, it exists unsafe { Some(self.ticks.get_unchecked(dense_index)) } } @@ -189,7 +215,9 @@ impl ComponentSparseSet { /// it exists). #[must_use = "The returned pointer must be used to drop the removed component."] pub fn remove_and_forget(&mut self, entity: Entity) -> Option> { - self.sparse.remove(entity).map(|dense_index| { + self.sparse.remove(entity.id()).map(|dense_index| { + let dense_index = dense_index as usize; + debug_assert_eq!(entity, self.entities[dense_index]); self.ticks.swap_remove(dense_index); self.entities.swap_remove(dense_index); let is_last = dense_index == self.dense.len() - 1; @@ -197,14 +225,20 @@ impl ComponentSparseSet { let value = unsafe { self.dense.swap_remove_and_forget_unchecked(dense_index) }; if !is_last { let swapped_entity = self.entities[dense_index]; - *self.sparse.get_mut(swapped_entity).unwrap() = dense_index; + #[cfg(not(debug_assertions))] + let idx = swapped_entity; + #[cfg(debug_assertions)] + let idx = swapped_entity.id(); + *self.sparse.get_mut(idx).unwrap() = dense_index as u32; } value }) } pub fn remove(&mut self, entity: Entity) -> bool { - if let Some(dense_index) = self.sparse.remove(entity) { + if let Some(dense_index) = self.sparse.remove(entity.id()) { + let dense_index = dense_index as usize; + debug_assert_eq!(entity, self.entities[dense_index]); self.ticks.swap_remove(dense_index); self.entities.swap_remove(dense_index); let is_last = dense_index == self.dense.len() - 1; @@ -212,7 +246,11 @@ impl ComponentSparseSet { unsafe { self.dense.swap_remove_and_drop_unchecked(dense_index) } if !is_last { let swapped_entity = self.entities[dense_index]; - *self.sparse.get_mut(swapped_entity).unwrap() = dense_index; + #[cfg(not(debug_assertions))] + let idx = swapped_entity; + #[cfg(debug_assertions)] + let idx = swapped_entity.id(); + *self.sparse.get_mut(idx).unwrap() = dense_index as u32; } true } else {