replace UnsafeCell with Cell in ComponentTicks and fix missing drop on component overwrite

Author: Frizi

crates/bevy_ecs/src/bundle.rs

@@ -140,20 +140,20 @@ impl BundleInfo {
         // bundle_info.component_ids are also in "bundle order"
         let mut bundle_component = 0;
         bundle.get_components(|component_ptr| {
-            // SAFE: component_id was initialized by get_dynamic_bundle_info
             let component_id = *self.component_ids.get_unchecked(bundle_component);
-            let component_status = bundle_status.get_unchecked(bundle_component);
             match self.storage_types[bundle_component] {
                 StorageType::Table => {
                     let column = table.get_column_mut(component_id).unwrap();
-                    column.set_data_unchecked(table_row, component_ptr);
-                    let column_status = column.get_ticks_unchecked_mut(table_row);
-                    match component_status {
+                    match bundle_status.get_unchecked(bundle_component) {
                         ComponentStatus::Added => {
-                            *column_status = ComponentTicks::new(change_tick);
+                            column.initialize(
+                                table_row,
+                                component_ptr,
+                                ComponentTicks::new(change_tick),
+                            );
                         }
                         ComponentStatus::Mutated => {
-                            column_status.set_changed(change_tick);
+                            column.replace(table_row, component_ptr, change_tick);
                         }
                     }
                 }

crates/bevy_ecs/src/component/mod.rs

@@ -6,6 +6,7 @@ use crate::storage::SparseSetIndex;
 use std::{
     alloc::Layout,
     any::{Any, TypeId},
+    cell::Cell,
     collections::hash_map::Entry,
 };
 use thiserror::Error;
@@ -306,10 +307,10 @@ impl Components {
     }
 }
 
-#[derive(Copy, Clone, Debug)]
+#[derive(Clone, Debug)]
 pub struct ComponentTicks {
     pub(crate) added: u32,
-    pub(crate) changed: u32,
+    pub(crate) changed: Cell<u32>,
 }
 
 impl ComponentTicks {
@@ -326,7 +327,7 @@ impl ComponentTicks {
 
     #[inline]
     pub fn is_changed(&self, last_change_tick: u32, change_tick: u32) -> bool {
-        let component_delta = change_tick.wrapping_sub(self.changed);
+        let component_delta = change_tick.wrapping_sub(self.changed.get());
         let system_delta = change_tick.wrapping_sub(last_change_tick);
 
         component_delta < system_delta
@@ -335,13 +336,13 @@ impl ComponentTicks {
     pub(crate) fn new(change_tick: u32) -> Self {
         Self {
             added: change_tick,
-            changed: change_tick,
+            changed: Cell::new(change_tick),
         }
     }
 
     pub(crate) fn check_ticks(&mut self, change_tick: u32) {
         check_tick(&mut self.added, change_tick);
-        check_tick(&mut self.changed, change_tick);
+        check_tick(self.changed.get_mut(), change_tick);
     }
 
     /// Manually sets the change tick.
@@ -357,8 +358,8 @@ impl ComponentTicks {
     /// component_ticks.set_changed(world.read_change_tick());
     /// ```
     #[inline]
-    pub fn set_changed(&mut self, change_tick: u32) {
-        self.changed = change_tick;
+    pub fn set_changed(&self, change_tick: u32) {
+        self.changed.set(change_tick);
     }
 }
 

crates/bevy_ecs/src/lib.rs

@@ -47,13 +47,34 @@ mod tests {
     };
     use bevy_tasks::TaskPool;
     use parking_lot::Mutex;
-    use std::{any::TypeId, sync::Arc};
+    use std::{
+        any::TypeId,
+        sync::{
+            atomic::{AtomicUsize, Ordering},
+            Arc,
+        },
+    };
 
     #[derive(Debug, PartialEq, Eq)]
     struct A(usize);
     struct B(usize);
     struct C;
 
+    #[derive(Clone, Debug)]
+    struct DropCk(Arc<AtomicUsize>);
+    impl DropCk {
+        fn new_pair() -> (Self, Arc<AtomicUsize>) {
+            let atomic = Arc::new(AtomicUsize::new(0));
+            (DropCk(atomic.clone()), atomic)
+        }
+    }
+
+    impl Drop for DropCk {
+        fn drop(&mut self) {
+            self.0.as_ref().fetch_add(1, Ordering::Relaxed);
+        }
+    }
+
     #[test]
     fn random_access() {
         let mut world = World::new();
@@ -1174,4 +1195,33 @@ mod tests {
         });
         assert_eq!(*world.get_resource::<i32>().unwrap(), 1);
     }
+
+    #[test]
+    fn insert_overwrite_drop() {
+        let (dropck1, dropped1) = DropCk::new_pair();
+        let (dropck2, dropped2) = DropCk::new_pair();
+        let mut world = World::default();
+        world.spawn().insert(dropck1).insert(dropck2);
+        assert_eq!(dropped1.load(Ordering::Relaxed), 1);
+        assert_eq!(dropped2.load(Ordering::Relaxed), 0);
+        drop(world);
+        assert_eq!(dropped1.load(Ordering::Relaxed), 1);
+        assert_eq!(dropped2.load(Ordering::Relaxed), 1);
+    }
+
+    #[test]
+    fn insert_overwrite_drop_sparse() {
+        let (dropck1, dropped1) = DropCk::new_pair();
+        let (dropck2, dropped2) = DropCk::new_pair();
+        let mut world = World::default();
+        world
+            .register_component(ComponentDescriptor::new::<DropCk>(StorageType::SparseSet))
+            .unwrap();
+        world.spawn().insert(dropck1).insert(dropck2);
+        assert_eq!(dropped1.load(Ordering::Relaxed), 1);
+        assert_eq!(dropped2.load(Ordering::Relaxed), 0);
+        drop(world);
+        assert_eq!(dropped1.load(Ordering::Relaxed), 1);
+        assert_eq!(dropped2.load(Ordering::Relaxed), 1);
+    }
 }

crates/bevy_ecs/src/query/fetch.rs

@@ -342,7 +342,7 @@ impl<'w, T: Component> Fetch<'w> for ReadFetch<T> {
                 let column = tables[archetype.table_id()]
                     .get_column(state.component_id)
                     .unwrap();
-                self.table_components = column.get_ptr().cast::<T>();
+                self.table_components = column.get_data_ptr().cast::<T>();
             }
             StorageType::SparseSet => self.entities = archetype.entities().as_ptr(),
         }
@@ -353,7 +353,7 @@ impl<'w, T: Component> Fetch<'w> for ReadFetch<T> {
         self.table_components = table
             .get_column(state.component_id)
             .unwrap()
-            .get_ptr()
+            .get_data_ptr()
             .cast::<T>();
     }
 
@@ -386,7 +386,7 @@ impl<T: Component> WorldQuery for &mut T {
 pub struct WriteFetch<T> {
     storage_type: StorageType,
     table_components: NonNull<T>,
-    table_ticks: *mut ComponentTicks,
+    table_ticks: *const ComponentTicks,
     entities: *const Entity,
     entity_table_rows: *const usize,
     sparse_set: *const ComponentSparseSet,
@@ -508,8 +508,8 @@ impl<'w, T: Component> Fetch<'w> for WriteFetch<T> {
                 let column = tables[archetype.table_id()]
                     .get_column(state.component_id)
                     .unwrap();
-                self.table_components = column.get_ptr().cast::<T>();
-                self.table_ticks = column.get_ticks_mut_ptr();
+                self.table_components = column.get_data_ptr().cast::<T>();
+                self.table_ticks = column.get_ticks_ptr();
             }
             StorageType::SparseSet => self.entities = archetype.entities().as_ptr(),
         }
@@ -518,8 +518,8 @@ impl<'w, T: Component> Fetch<'w> for WriteFetch<T> {
     #[inline]
     unsafe fn set_table(&mut self, state: &Self::State, table: &Table) {
         let column = table.get_column(state.component_id).unwrap();
-        self.table_components = column.get_ptr().cast::<T>();
-        self.table_ticks = column.get_ticks_mut_ptr();
+        self.table_components = column.get_data_ptr().cast::<T>();
+        self.table_ticks = column.get_ticks_ptr();
     }
 
     #[inline]
@@ -529,7 +529,7 @@ impl<'w, T: Component> Fetch<'w> for WriteFetch<T> {
                 let table_row = *self.entity_table_rows.add(archetype_index);
                 Mut {
                     value: &mut *self.table_components.as_ptr().add(table_row),
-                    component_ticks: &mut *self.table_ticks.add(table_row),
+                    component_ticks: &*self.table_ticks.add(table_row),
                     change_tick: self.change_tick,
                     last_change_tick: self.last_change_tick,
                 }
@@ -540,7 +540,7 @@ impl<'w, T: Component> Fetch<'w> for WriteFetch<T> {
                     (*self.sparse_set).get_with_ticks(entity).unwrap();
                 Mut {
                     value: &mut *component.cast::<T>(),
-                    component_ticks: &mut *component_ticks,
+                    component_ticks: &*component_ticks,
                     change_tick: self.change_tick,
                     last_change_tick: self.last_change_tick,
                 }
@@ -552,7 +552,7 @@ impl<'w, T: Component> Fetch<'w> for WriteFetch<T> {
     unsafe fn table_fetch(&mut self, table_row: usize) -> Self::Item {
         Mut {
             value: &mut *self.table_components.as_ptr().add(table_row),
-            component_ticks: &mut *self.table_ticks.add(table_row),
+            component_ticks: &*self.table_ticks.add(table_row),
             change_tick: self.change_tick,
             last_change_tick: self.last_change_tick,
         }
@@ -853,7 +853,7 @@ impl<'w, T: Component> Fetch<'w> for ChangeTrackersFetch<T> {
                 let column = tables[archetype.table_id()]
                     .get_column(state.component_id)
                     .unwrap();
-                self.table_ticks = column.get_ticks_mut_ptr().cast::<ComponentTicks>();
+                self.table_ticks = column.get_ticks_ptr().cast::<ComponentTicks>();
             }
             StorageType::SparseSet => self.entities = archetype.entities().as_ptr(),
         }
@@ -864,7 +864,7 @@ impl<'w, T: Component> Fetch<'w> for ChangeTrackersFetch<T> {
         self.table_ticks = table
             .get_column(state.component_id)
             .unwrap()
-            .get_ticks_mut_ptr()
+            .get_ticks_ptr()
             .cast::<ComponentTicks>();
     }
 
@@ -874,7 +874,7 @@ impl<'w, T: Component> Fetch<'w> for ChangeTrackersFetch<T> {
             StorageType::Table => {
                 let table_row = *self.entity_table_rows.add(archetype_index);
                 ChangeTrackers {
-                    component_ticks: *self.table_ticks.add(table_row),
+                    component_ticks: (&*self.table_ticks.add(table_row)).clone(),
                     marker: PhantomData,
                     last_change_tick: self.last_change_tick,
                     change_tick: self.change_tick,
@@ -883,7 +883,7 @@ impl<'w, T: Component> Fetch<'w> for ChangeTrackersFetch<T> {
             StorageType::SparseSet => {
                 let entity = *self.entities.add(archetype_index);
                 ChangeTrackers {
-                    component_ticks: *(*self.sparse_set).get_ticks(entity).unwrap(),
+                    component_ticks: (&*self.sparse_set).get_ticks(entity).cloned().unwrap(),
                     marker: PhantomData,
                     last_change_tick: self.last_change_tick,
                     change_tick: self.change_tick,
@@ -895,7 +895,7 @@ impl<'w, T: Component> Fetch<'w> for ChangeTrackersFetch<T> {
     #[inline]
     unsafe fn table_fetch(&mut self, table_row: usize) -> Self::Item {
         ChangeTrackers {
-            component_ticks: *self.table_ticks.add(table_row),
+            component_ticks: (&*self.table_ticks.add(table_row)).clone(),
             marker: PhantomData,
             last_change_tick: self.last_change_tick,
             change_tick: self.change_tick,

crates/bevy_ecs/src/query/filter.rs

@@ -561,7 +561,7 @@ macro_rules! impl_tick_filter {
         $(#[$fetch_meta])*
         pub struct $fetch_name<T> {
             storage_type: StorageType,
-            table_ticks: *mut ComponentTicks,
+            table_ticks: *const ComponentTicks,
             entity_table_rows: *const usize,
             marker: PhantomData<T>,
             entities: *const Entity,
@@ -630,7 +630,7 @@ macro_rules! impl_tick_filter {
             unsafe fn init(world: &World, state: &Self::State, last_change_tick: u32, change_tick: u32) -> Self {
                 let mut value = Self {
                     storage_type: state.storage_type,
-                    table_ticks: ptr::null_mut::<ComponentTicks>(),
+                    table_ticks: ptr::null::<ComponentTicks>(),
                     entities: ptr::null::<Entity>(),
                     entity_table_rows: ptr::null::<usize>(),
                     sparse_set: ptr::null::<ComponentSparseSet>(),
@@ -655,7 +655,7 @@ macro_rules! impl_tick_filter {
             unsafe fn set_table(&mut self, state: &Self::State, table: &Table) {
                 self.table_ticks = table
                     .get_column(state.component_id).unwrap()
-                    .get_ticks_mut_ptr();
+                    .get_ticks_ptr();
             }
 
             unsafe fn set_archetype(&mut self, state: &Self::State, archetype: &Archetype, tables: &Tables) {
@@ -665,7 +665,7 @@ macro_rules! impl_tick_filter {
                         let table = &tables[archetype.table_id()];
                         self.table_ticks = table
                             .get_column(state.component_id).unwrap()
-                            .get_ticks_mut_ptr();
+                            .get_ticks_ptr();
                     }
                     StorageType::SparseSet => self.entities = archetype.entities().as_ptr(),
                 }
@@ -683,7 +683,7 @@ macro_rules! impl_tick_filter {
                     }
                     StorageType::SparseSet => {
                         let entity = *self.entities.add(archetype_index);
-                        let ticks = (*(*self.sparse_set).get_ticks(entity).unwrap());
+                        let ticks = (&*self.sparse_set).get_ticks(entity).cloned().unwrap();
                         $is_detected(&ticks, self.last_change_tick, self.change_tick)
                     }
                 }

crates/bevy_ecs/src/reflect.rs

@@ -116,7 +116,7 @@ impl<C: Component + Reflect + FromWorld> FromType<C> for ReflectComponent {
 /// Unique borrow of a Reflected component
 pub struct ReflectMut<'a> {
     pub(crate) value: &'a mut dyn Reflect,
-    pub(crate) component_ticks: &'a mut ComponentTicks,
+    pub(crate) component_ticks: &'a ComponentTicks,
     pub(crate) last_change_tick: u32,
     pub(crate) change_tick: u32,
 }

crates/bevy_ecs/src/schedule/stage.rs

@@ -2050,7 +2050,7 @@ mod tests {
                 assert!(time_since_last_check <= MAX_DELTA);
                 let time_since_last_check = tracker
                     .change_tick
-                    .wrapping_sub(tracker.component_ticks.changed);
+                    .wrapping_sub(tracker.component_ticks.changed.get());
                 assert!(time_since_last_check <= MAX_DELTA);
             }
             let change_tick = world.change_tick.get_mut();

crates/bevy_ecs/src/storage/blob_vec.rs

@@ -89,12 +89,22 @@ impl BlobVec {
     /// `index` must be in bounds
     /// Allows aliased mutable access to `index`'s data. Caller must ensure this does not happen
     #[inline]
-    pub unsafe fn set_unchecked(&self, index: usize, value: *mut u8) {
+    pub unsafe fn initialize_unchecked(&mut self, index: usize, value: *mut u8) {
         debug_assert!(index < self.len());
         let ptr = self.get_unchecked(index);
         std::ptr::copy_nonoverlapping(value, ptr, self.item_layout.size());
     }
 
+    /// # Safety
+    /// index must be in-bounds
+    ///
+    pub unsafe fn replace_unchecked(&mut self, index: usize, value: *mut u8) {
+        debug_assert!(index < self.len());
+        let ptr = self.get_unchecked(index);
+        (self.drop)(ptr);
+        std::ptr::copy_nonoverlapping(value, ptr, self.item_layout.size());
+    }
+
     /// increases the length by one (and grows the vec if needed) with uninitialized memory and
     /// returns the index
     ///
@@ -267,7 +277,7 @@ mod tests {
     /// `blob_vec` must have a layout that matches Layout::new::<T>()
     unsafe fn push<T>(blob_vec: &mut BlobVec, mut value: T) {
         let index = blob_vec.push_uninit();
-        blob_vec.set_unchecked(index, (&mut value as *mut T).cast::<u8>());
+        blob_vec.initialize_unchecked(index, (&mut value as *mut T).cast::<u8>());
         std::mem::forget(value);
     }