From 4cd53cc7e12fb8e4d612fc02b634828cda97af82 Mon Sep 17 00:00:00 2001 From: James Liu Date: Tue, 5 Mar 2024 21:52:18 -0800 Subject: [PATCH] Clean up pointer use in BundleSpawner/BundleInserter (#12269) # Objective Following #10756, we're now using raw pointers in BundleInserter and BundleSpawner. This is primarily to get around the need to split the borrow on the World, but it leaves a lot to be desired in terms of safety guarantees. There's no type level guarantee the code can't dereference a null pointer, and it's restoring them to borrows fairly liberally within the associated functions. ## Solution * Replace the pointers with `NonNull` and a new `bevy_ptr::ConstNonNull` that only allows conversion back to read-only borrows * Remove the closure to avoid potentially aliasing through the closure by restructuring the match expression. * Move all conversions back into borrows as far up as possible to ensure that the borrow checker is at least locally followed. --- crates/bevy_ecs/src/bundle.rs | 406 +++++++++++++++++----------------- crates/bevy_ptr/src/lib.rs | 115 ++++++++++ 2 files changed, 312 insertions(+), 209 deletions(-) diff --git a/crates/bevy_ecs/src/bundle.rs b/crates/bevy_ecs/src/bundle.rs index f5e9a32fd0614..f3d42a66ef29b 100644 --- a/crates/bevy_ecs/src/bundle.rs +++ b/crates/bevy_ecs/src/bundle.rs @@ -15,11 +15,12 @@ use crate::{ prelude::World, query::DebugCheckedUnwrap, storage::{SparseSetIndex, SparseSets, Storages, Table, TableRow}, - world::{unsafe_world_cell::UnsafeWorldCell, DeferredWorld}, + world::unsafe_world_cell::UnsafeWorldCell, }; -use bevy_ptr::OwningPtr; +use bevy_ptr::{ConstNonNull, OwningPtr}; use bevy_utils::all_tuples; use std::any::TypeId; +use std::ptr::NonNull; /// The `Bundle` trait enables insertion and removal of [`Component`]s from an entity. /// @@ -510,10 +511,10 @@ impl BundleInfo { // SAFETY: We have exclusive world access so our pointers can't be invalidated externally pub(crate) struct BundleInserter<'w> { world: UnsafeWorldCell<'w>, - bundle_info: *const BundleInfo, - add_bundle: *const AddBundle, - table: *mut Table, - archetype: *mut Archetype, + bundle_info: ConstNonNull, + add_bundle: ConstNonNull, + table: NonNull, + archetype: NonNull, result: InsertBundleResult, change_tick: Tick, } @@ -521,11 +522,11 @@ pub(crate) struct BundleInserter<'w> { pub(crate) enum InsertBundleResult { SameArchetype, NewArchetypeSameTable { - new_archetype: *mut Archetype, + new_archetype: NonNull, }, NewArchetypeNewTable { - new_archetype: *mut Archetype, - new_table: *mut Table, + new_archetype: NonNull, + new_table: NonNull
, }, } @@ -546,7 +547,7 @@ impl<'w> BundleInserter<'w> { /// Creates a new [`BundleInserter`]. /// /// # Safety - /// Caller must ensure that `bundle_id` exists in `world.bundles` + /// - Caller must ensure that `bundle_id` exists in `world.bundles`. #[inline] pub(crate) unsafe fn new_with_id( world: &'w mut World, @@ -575,10 +576,10 @@ impl<'w> BundleInserter<'w> { let table_id = archetype.table_id(); let table = &mut world.storages.tables[table_id]; Self { - add_bundle, - archetype, - bundle_info, - table, + add_bundle: add_bundle.into(), + archetype: archetype.into(), + bundle_info: bundle_info.into(), + table: table.into(), result: InsertBundleResult::SameArchetype, change_tick, world: world.as_unsafe_world_cell(), @@ -598,24 +599,26 @@ impl<'w> BundleInserter<'w> { if table_id == new_table_id { let table = &mut world.storages.tables[table_id]; Self { - add_bundle, - archetype, - bundle_info, - table, - result: InsertBundleResult::NewArchetypeSameTable { new_archetype }, + add_bundle: add_bundle.into(), + archetype: archetype.into(), + bundle_info: bundle_info.into(), + table: table.into(), + result: InsertBundleResult::NewArchetypeSameTable { + new_archetype: new_archetype.into(), + }, change_tick, world: world.as_unsafe_world_cell(), } } else { let (table, new_table) = world.storages.tables.get_2_mut(table_id, new_table_id); Self { - add_bundle, - archetype, - bundle_info, - table, + add_bundle: add_bundle.into(), + archetype: archetype.into(), + bundle_info: bundle_info.into(), + table: table.into(), result: InsertBundleResult::NewArchetypeNewTable { - new_archetype, - new_table, + new_archetype: new_archetype.into(), + new_table: new_table.into(), }, change_tick, world: world.as_unsafe_world_cell(), @@ -623,6 +626,7 @@ impl<'w> BundleInserter<'w> { } } } + /// # Safety /// `entity` must currently exist in the source archetype for this inserter. `location` /// must be `entity`'s location in the archetype. `T` must match this [`BundleInfo`]'s type @@ -633,191 +637,174 @@ impl<'w> BundleInserter<'w> { location: EntityLocation, bundle: T, ) -> EntityLocation { - // SAFETY: We do not make any structural changes to the archetype graph through self.world so these pointers always remain valid - let trigger_hooks = |archetype: &Archetype, mut world: DeferredWorld| { - let bundle_info = &*self.bundle_info; - let add_bundle = &*self.add_bundle; + let bundle_info = self.bundle_info.as_ref(); + let add_bundle = self.add_bundle.as_ref(); + let table = self.table.as_mut(); + let archetype = self.archetype.as_mut(); + + let (new_archetype, new_location) = match &mut self.result { + InsertBundleResult::SameArchetype => { + // SAFETY: Mutable references do not alias and will be dropped after this block + let sparse_sets = { + let world = self.world.world_mut(); + &mut world.storages.sparse_sets + }; - if archetype.has_on_add() { - world.trigger_on_add( + bundle_info.write_components( + table, + sparse_sets, + add_bundle, entity, - bundle_info - .iter_components() - .zip(add_bundle.bundle_status.iter()) - .filter(|(_, &status)| status == ComponentStatus::Added) - .map(|(id, _)| id), + location.table_row, + self.change_tick, + bundle, ); - } - if archetype.has_on_insert() { - world.trigger_on_insert(entity, bundle_info.iter_components()); - } - }; - match &mut self.result { - InsertBundleResult::SameArchetype => { - { - // SAFETY: Mutable references do not alias and will be dropped after this block - let sparse_sets = { - let world = self.world.world_mut(); - &mut world.storages.sparse_sets - }; - let table = &mut *self.table; - let bundle_info = &*self.bundle_info; - let add_bundle = &*self.add_bundle; - - bundle_info.write_components( - table, - sparse_sets, - add_bundle, - entity, - location.table_row, - self.change_tick, - bundle, - ); - } - // SAFETY: We have no outstanding mutable references to world as they were dropped - unsafe { - let archetype = &*self.archetype; - trigger_hooks(archetype, self.world.into_deferred()); - } - location + + (archetype, location) } InsertBundleResult::NewArchetypeSameTable { new_archetype } => { - let new_location = { - // SAFETY: Mutable references do not alias and will be dropped after this block - let (sparse_sets, entities) = { - let world = self.world.world_mut(); - (&mut world.storages.sparse_sets, &mut world.entities) - }; - let table = &mut *self.table; - let archetype = &mut *self.archetype; - let new_archetype = &mut **new_archetype; - - let result = archetype.swap_remove(location.archetype_row); - if let Some(swapped_entity) = result.swapped_entity { - let swapped_location = - // SAFETY: If the swap was successful, swapped_entity must be valid. - unsafe { entities.get(swapped_entity).debug_checked_unwrap() }; - entities.set( - swapped_entity.index(), - EntityLocation { - archetype_id: swapped_location.archetype_id, - archetype_row: location.archetype_row, - table_id: swapped_location.table_id, - table_row: swapped_location.table_row, - }, - ); - } - let new_location = new_archetype.allocate(entity, result.table_row); - entities.set(entity.index(), new_location); - - let bundle_info = &*self.bundle_info; - let add_bundle = &*self.add_bundle; - bundle_info.write_components( - table, - sparse_sets, - add_bundle, - entity, - result.table_row, - self.change_tick, - bundle, - ); - new_location + let new_archetype = new_archetype.as_mut(); + + // SAFETY: Mutable references do not alias and will be dropped after this block + let (sparse_sets, entities) = { + let world = self.world.world_mut(); + (&mut world.storages.sparse_sets, &mut world.entities) }; - // SAFETY: We have no outstanding mutable references to world as they were dropped - unsafe { trigger_hooks(&**new_archetype, self.world.into_deferred()) }; + let result = archetype.swap_remove(location.archetype_row); + if let Some(swapped_entity) = result.swapped_entity { + let swapped_location = + // SAFETY: If the swap was successful, swapped_entity must be valid. + unsafe { entities.get(swapped_entity).debug_checked_unwrap() }; + entities.set( + swapped_entity.index(), + EntityLocation { + archetype_id: swapped_location.archetype_id, + archetype_row: location.archetype_row, + table_id: swapped_location.table_id, + table_row: swapped_location.table_row, + }, + ); + } + let new_location = new_archetype.allocate(entity, result.table_row); + entities.set(entity.index(), new_location); + bundle_info.write_components( + table, + sparse_sets, + add_bundle, + entity, + result.table_row, + self.change_tick, + bundle, + ); - new_location + (new_archetype, new_location) } InsertBundleResult::NewArchetypeNewTable { new_archetype, new_table, } => { - let new_location = { - // SAFETY: Mutable references do not alias and will be dropped after this block - let (archetypes_ptr, sparse_sets, entities) = { - let world = self.world.world_mut(); - let archetype_ptr: *mut Archetype = - world.archetypes.archetypes.as_mut_ptr(); - ( - archetype_ptr, - &mut world.storages.sparse_sets, - &mut world.entities, - ) - }; - let table = &mut *self.table; - let new_table = &mut **new_table; - let archetype = &mut *self.archetype; - let new_archetype = &mut **new_archetype; - let result = archetype.swap_remove(location.archetype_row); - if let Some(swapped_entity) = result.swapped_entity { - let swapped_location = - // SAFETY: If the swap was successful, swapped_entity must be valid. - unsafe { entities.get(swapped_entity).debug_checked_unwrap() }; - entities.set( - swapped_entity.index(), - EntityLocation { - archetype_id: swapped_location.archetype_id, - archetype_row: location.archetype_row, - table_id: swapped_location.table_id, - table_row: swapped_location.table_row, - }, - ); - } - // PERF: store "non bundle" components in edge, then just move those to avoid - // redundant copies - let move_result = table.move_to_superset_unchecked(result.table_row, new_table); - let new_location = new_archetype.allocate(entity, move_result.new_row); - entities.set(entity.index(), new_location); - - // if an entity was moved into this entity's table spot, update its table row - if let Some(swapped_entity) = move_result.swapped_entity { - let swapped_location = - // SAFETY: If the swap was successful, swapped_entity must be valid. - unsafe { entities.get(swapped_entity).debug_checked_unwrap() }; - let swapped_archetype = if archetype.id() == swapped_location.archetype_id { - archetype - } else if new_archetype.id() == swapped_location.archetype_id { - new_archetype - } else { - // SAFETY: the only two borrowed archetypes are above and we just did collision checks - &mut *archetypes_ptr.add(swapped_location.archetype_id.index()) - }; - - entities.set( - swapped_entity.index(), - EntityLocation { - archetype_id: swapped_location.archetype_id, - archetype_row: swapped_location.archetype_row, - table_id: swapped_location.table_id, - table_row: result.table_row, - }, - ); - swapped_archetype + let new_table = new_table.as_mut(); + let new_archetype = new_archetype.as_mut(); + + // SAFETY: Mutable references do not alias and will be dropped after this block + let (archetypes_ptr, sparse_sets, entities) = { + let world = self.world.world_mut(); + let archetype_ptr: *mut Archetype = world.archetypes.archetypes.as_mut_ptr(); + ( + archetype_ptr, + &mut world.storages.sparse_sets, + &mut world.entities, + ) + }; + let result = archetype.swap_remove(location.archetype_row); + if let Some(swapped_entity) = result.swapped_entity { + let swapped_location = + // SAFETY: If the swap was successful, swapped_entity must be valid. + unsafe { entities.get(swapped_entity).debug_checked_unwrap() }; + entities.set( + swapped_entity.index(), + EntityLocation { + archetype_id: swapped_location.archetype_id, + archetype_row: location.archetype_row, + table_id: swapped_location.table_id, + table_row: swapped_location.table_row, + }, + ); + } + // PERF: store "non bundle" components in edge, then just move those to avoid + // redundant copies + let move_result = table.move_to_superset_unchecked(result.table_row, new_table); + let new_location = new_archetype.allocate(entity, move_result.new_row); + entities.set(entity.index(), new_location); + + // if an entity was moved into this entity's table spot, update its table row + if let Some(swapped_entity) = move_result.swapped_entity { + let swapped_location = + // SAFETY: If the swap was successful, swapped_entity must be valid. + unsafe { entities.get(swapped_entity).debug_checked_unwrap() }; + + entities.set( + swapped_entity.index(), + EntityLocation { + archetype_id: swapped_location.archetype_id, + archetype_row: swapped_location.archetype_row, + table_id: swapped_location.table_id, + table_row: result.table_row, + }, + ); + + if archetype.id() == swapped_location.archetype_id { + archetype + .set_entity_table_row(swapped_location.archetype_row, result.table_row); + } else if new_archetype.id() == swapped_location.archetype_id { + new_archetype + .set_entity_table_row(swapped_location.archetype_row, result.table_row); + } else { + // SAFETY: the only two borrowed archetypes are above and we just did collision checks + (*archetypes_ptr.add(swapped_location.archetype_id.index())) .set_entity_table_row(swapped_location.archetype_row, result.table_row); } + } - let bundle_info = &*self.bundle_info; - let add_bundle = &*self.add_bundle; - bundle_info.write_components( - new_table, - sparse_sets, - add_bundle, - entity, - move_result.new_row, - self.change_tick, - bundle, - ); + bundle_info.write_components( + new_table, + sparse_sets, + add_bundle, + entity, + move_result.new_row, + self.change_tick, + bundle, + ); - new_location - }; + (new_archetype, new_location) + } + }; - // SAFETY: We have no outstanding mutable references to world as they were dropped - unsafe { trigger_hooks(&**new_archetype, self.world.into_deferred()) }; + // SAFETY: We have no outstanding mutable references to world as they were dropped + let mut deferred_world = unsafe { self.world.into_deferred() }; - new_location + if new_archetype.has_on_add() { + // SAFETY: All components in the bundle are guaranteed to exist in the World + // as they must be initialized before creating the BundleInfo. + unsafe { + deferred_world.trigger_on_add( + entity, + bundle_info + .iter_components() + .zip(add_bundle.bundle_status.iter()) + .filter(|(_, &status)| status == ComponentStatus::Added) + .map(|(id, _)| id), + ); } } + if new_archetype.has_on_insert() { + // SAFETY: All components in the bundle are guaranteed to exist in the World + // as they must be initialized before creating the BundleInfo. + unsafe { deferred_world.trigger_on_insert(entity, bundle_info.iter_components()) } + } + + new_location } #[inline] @@ -830,9 +817,9 @@ impl<'w> BundleInserter<'w> { // SAFETY: We have exclusive world access so our pointers can't be invalidated externally pub(crate) struct BundleSpawner<'w> { world: UnsafeWorldCell<'w>, - bundle_info: *const BundleInfo, - table: *mut Table, - archetype: *mut Archetype, + bundle_info: ConstNonNull, + table: NonNull
, + archetype: NonNull, change_tick: Tick, } @@ -866,9 +853,9 @@ impl<'w> BundleSpawner<'w> { let archetype = &mut world.archetypes[new_archetype_id]; let table = &mut world.storages.tables[archetype.table_id()]; Self { - bundle_info, - table, - archetype, + bundle_info: bundle_info.into(), + table: table.into(), + archetype: archetype.into(), change_tick, world: world.as_unsafe_world_cell(), } @@ -877,7 +864,7 @@ impl<'w> BundleSpawner<'w> { #[inline] pub fn reserve_storage(&mut self, additional: usize) { // SAFETY: There are no outstanding world references - let (archetype, table) = unsafe { (&mut *self.archetype, &mut *self.table) }; + let (archetype, table) = unsafe { (self.archetype.as_mut(), self.table.as_mut()) }; archetype.reserve(additional); table.reserve(additional); } @@ -890,6 +877,10 @@ impl<'w> BundleSpawner<'w> { entity: Entity, bundle: T, ) -> EntityLocation { + let table = self.table.as_mut(); + let archetype = self.archetype.as_mut(); + let bundle_info = self.bundle_info.as_ref(); + // SAFETY: We do not make any structural changes to the archetype graph through self.world so this pointer always remain valid let location = { // SAFETY: Mutable references do not alias and will be dropped after this block @@ -897,11 +888,8 @@ impl<'w> BundleSpawner<'w> { let world = self.world.world_mut(); (&mut world.storages.sparse_sets, &mut world.entities) }; - let table = &mut *self.table; - let archetype = &mut *self.archetype; let table_row = table.allocate(entity); let location = archetype.allocate(entity, table_row); - let bundle_info = &*self.bundle_info; bundle_info.write_components( table, sparse_sets, @@ -916,16 +904,16 @@ impl<'w> BundleSpawner<'w> { }; // SAFETY: We have no outstanding mutable references to world as they were dropped - unsafe { - let archetype = &*self.archetype; - let bundle_info = &*self.bundle_info; - let mut world = self.world.into_deferred(); - if archetype.has_on_add() { - world.trigger_on_add(entity, bundle_info.iter_components()); - } - if archetype.has_on_insert() { - world.trigger_on_insert(entity, bundle_info.iter_components()); - } + let mut deferred_world = unsafe { self.world.into_deferred() }; + if archetype.has_on_add() { + // SAFETY: All components in the bundle are guaranteed to exist in the World + // as they must be initialized before creating the BundleInfo. + unsafe { deferred_world.trigger_on_add(entity, bundle_info.iter_components()) }; + } + if archetype.has_on_insert() { + // SAFETY: All components in the bundle are guaranteed to exist in the World + // as they must be initialized before creating the BundleInfo. + unsafe { deferred_world.trigger_on_insert(entity, bundle_info.iter_components()) }; } location diff --git a/crates/bevy_ptr/src/lib.rs b/crates/bevy_ptr/src/lib.rs index d26cc1fcb2a5e..d6c0322461528 100644 --- a/crates/bevy_ptr/src/lib.rs +++ b/crates/bevy_ptr/src/lib.rs @@ -26,6 +26,119 @@ mod sealed { impl Sealed for super::Unaligned {} } +/// A newtype around [`NonNull`] that only allows conversion to read-only borrows or pointers. +/// +/// This type can be thought of as the `*const T` to [`NonNull`]'s `*mut T`. +#[repr(transparent)] +pub struct ConstNonNull(NonNull); + +impl ConstNonNull { + /// Creates a new `ConstNonNull` if `ptr` is non-null. + /// + /// # Examples + /// + /// ``` + /// use bevy_ptr::ConstNonNull; + /// + /// let x = 0u32; + /// let ptr = ConstNonNull::::new(&x as *const _).expect("ptr is null!"); + /// + /// if let Some(ptr) = ConstNonNull::::new(std::ptr::null()) { + /// unreachable!(); + /// } + /// ``` + pub fn new(ptr: *const T) -> Option { + NonNull::new(ptr.cast_mut()).map(Self) + } + + /// Creates a new `ConstNonNull`. + /// + /// # Safety + /// + /// `ptr` must be non-null. + /// + /// # Examples + /// + /// ``` + /// use bevy_ptr::ConstNonNull; + /// + /// let x = 0u32; + /// let ptr = unsafe { ConstNonNull::new_unchecked(&x as *const _) }; + /// ``` + /// + /// *Incorrect* usage of this function: + /// + /// ```rust,no_run + /// use bevy_ptr::ConstNonNull; + /// + /// // NEVER DO THAT!!! This is undefined behavior. ⚠️ + /// let ptr = unsafe { ConstNonNull::::new_unchecked(std::ptr::null()) }; + /// ``` + pub const unsafe fn new_unchecked(ptr: *const T) -> Self { + // SAFETY: This function's safety invariants are identical to `NonNull::new_unchecked` + // The caller must satisfy all of them. + unsafe { Self(NonNull::new_unchecked(ptr.cast_mut())) } + } + + /// Returns a shared reference to the value. + /// + /// # Safety + /// + /// When calling this method, you have to ensure that all of the following is true: + /// + /// * The pointer must be properly aligned. + /// + /// * It must be "dereferenceable" in the sense defined in [the module documentation]. + /// + /// * The pointer must point to an initialized instance of `T`. + /// + /// * You must enforce Rust's aliasing rules, since the returned lifetime `'a` is + /// arbitrarily chosen and does not necessarily reflect the actual lifetime of the data. + /// In particular, while this reference exists, the memory the pointer points to must + /// not get mutated (except inside `UnsafeCell`). + /// + /// This applies even if the result of this method is unused! + /// (The part about being initialized is not yet fully decided, but until + /// it is, the only safe approach is to ensure that they are indeed initialized.) + /// + /// # Examples + /// + /// ``` + /// use bevy_ptr::ConstNonNull; + /// + /// let mut x = 0u32; + /// let ptr = ConstNonNull::new(&mut x as *mut _).expect("ptr is null!"); + /// + /// let ref_x = unsafe { ptr.as_ref() }; + /// println!("{ref_x}"); + /// ``` + /// + /// [the module documentation]: core::ptr#safety + #[inline] + pub unsafe fn as_ref<'a>(&self) -> &'a T { + // SAFETY: This function's safety invariants are identical to `NonNull::as_ref` + // The caller must satisfy all of them. + unsafe { self.0.as_ref() } + } +} + +impl From> for ConstNonNull { + fn from(value: NonNull) -> ConstNonNull { + ConstNonNull(value) + } +} + +impl<'a, T: ?Sized> From<&'a T> for ConstNonNull { + fn from(value: &'a T) -> ConstNonNull { + ConstNonNull(NonNull::from(value)) + } +} + +impl<'a, T: ?Sized> From<&'a mut T> for ConstNonNull { + fn from(value: &'a mut T) -> ConstNonNull { + ConstNonNull(NonNull::from(value)) + } +} /// Type-erased borrow of some unknown type chosen when constructing this type. /// /// This type tries to act "borrow-like" which means that: @@ -280,6 +393,7 @@ impl<'a> OwningPtr<'a> { f(unsafe { PtrMut::from(&mut *temp).promote() }) } } + impl<'a, A: IsAligned> OwningPtr<'a, A> { /// Creates a new instance from a raw pointer. /// @@ -346,6 +460,7 @@ impl<'a, A: IsAligned> OwningPtr<'a, A> { unsafe { PtrMut::new(self.0) } } } + impl<'a> OwningPtr<'a, Unaligned> { /// Consumes the [`OwningPtr`] to obtain ownership of the underlying data of type `T`. ///