Path: blob/main/crates/bevy_ecs/src/storage/blob_array.rs
9353 views
use alloc::alloc::handle_alloc_error;1use bevy_ptr::{OwningPtr, Ptr, PtrMut};2use bevy_utils::OnDrop;3use core::{alloc::Layout, cell::UnsafeCell, num::NonZeroUsize, ptr::NonNull};45/// A flat, type-erased data storage type.6///7/// Used to densely store homogeneous ECS data. A blob is usually just an arbitrary block of contiguous memory without any identity, and8/// could be used to represent any arbitrary data (i.e. string, arrays, etc). This type only stores meta-data about the blob that it stores,9/// and a pointer to the location of the start of the array, similar to a C-style `void*` array.10///11/// This type is reliant on its owning type to store the capacity and length information.12#[derive(Debug)]13pub(super) struct BlobArray {14item_layout: Layout,15// the `data` ptr's layout is always `array_layout(item_layout, capacity)`16data: NonNull<u8>,17// None if the underlying type doesn't need to be dropped18pub drop: Option<unsafe fn(OwningPtr<'_>)>,19#[cfg(debug_assertions)]20capacity: usize,21}2223impl BlobArray {24/// Create a new [`BlobArray`] with a specified `capacity`.25/// If `capacity` is 0, no allocations will be made.26///27/// `drop` is an optional function pointer that is meant to be invoked when any element in the [`BlobArray`]28/// should be dropped. For all Rust-based types, this should match 1:1 with the implementation of [`Drop`]29/// if present, and should be `None` if `T: !Drop`. For non-Rust based types, this should match any cleanup30/// processes typically associated with the stored element.31///32/// # Safety33/// `drop` should be safe to call with an [`OwningPtr`] pointing to any item that's been placed into this [`BlobArray`].34/// If `drop` is `None`, the items will be leaked. This should generally be set as None based on [`needs_drop`].35///36/// [`needs_drop`]: std::mem::needs_drop37pub unsafe fn with_capacity(38item_layout: Layout,39drop_fn: Option<unsafe fn(OwningPtr<'_>)>,40capacity: usize,41) -> Self {42if capacity == 0 {43let align = NonZeroUsize::new(item_layout.align()).expect("alignment must be > 0");4445// Create a dangling pointer with the given alignment.46let data = NonNull::without_provenance(align);4748Self {49item_layout,50drop: drop_fn,51data,52#[cfg(debug_assertions)]53capacity,54}55} else {56// SAFETY: Upheld by caller57let mut arr = unsafe { Self::with_capacity(item_layout, drop_fn, 0) };58// SAFETY: `capacity` > 059unsafe { arr.alloc(NonZeroUsize::new_unchecked(capacity)) }60arr61}62}6364/// Returns the [`Layout`] of the element type stored in the vector.65#[inline]66pub fn layout(&self) -> Layout {67self.item_layout68}6970/// Return `true` if this [`BlobArray`] stores `ZSTs`.71pub fn is_zst(&self) -> bool {72self.item_layout.size() == 073}7475/// Returns the drop function for values stored in the vector,76/// or `None` if they don't need to be dropped.77#[inline]78pub fn get_drop(&self) -> Option<unsafe fn(OwningPtr<'_>)> {79self.drop80}8182/// Returns a reference to the element at `index`, without doing bounds checking.83///84/// *`len` refers to the length of the array, the number of elements that have been initialized, and are safe to read.85/// Just like [`Vec::len`].*86///87/// # Safety88/// - The element at index `index` is safe to access.89/// (If the safety requirements of every method that has been used on `Self` have been fulfilled, the caller just needs to ensure that `index` < `len`)90///91/// [`Vec::len`]: alloc::vec::Vec::len92#[inline]93pub unsafe fn get_unchecked(&self, index: usize) -> Ptr<'_> {94#[cfg(debug_assertions)]95debug_assert!(index < self.capacity);96let size = self.item_layout.size();97// SAFETY:98// - The caller ensures that `index` fits in this array,99// so this operation will not overflow the original allocation.100// - `size` is a multiple of the erased type's alignment,101// so adding a multiple of `size` will preserve alignment.102unsafe { self.get_ptr().byte_add(index * size) }103}104105/// Returns a mutable reference to the element at `index`, without doing bounds checking.106///107/// *`len` refers to the length of the array, the number of elements that have been initialized, and are safe to read.108/// Just like [`Vec::len`].*109///110/// # Safety111/// - The element with at index `index` is safe to access.112/// (If the safety requirements of every method that has been used on `Self` have been fulfilled, the caller just needs to ensure that `index` < `len`)113///114/// [`Vec::len`]: alloc::vec::Vec::len115#[inline]116pub unsafe fn get_unchecked_mut(&mut self, index: usize) -> PtrMut<'_> {117#[cfg(debug_assertions)]118debug_assert!(index < self.capacity);119let size = self.item_layout.size();120// SAFETY:121// - The caller ensures that `index` fits in this vector,122// so this operation will not overflow the original allocation.123// - `size` is a multiple of the erased type's alignment,124// so adding a multiple of `size` will preserve alignment.125unsafe { self.get_ptr_mut().byte_add(index * size) }126}127128/// Gets a [`Ptr`] to the start of the array129#[inline]130pub fn get_ptr(&self) -> Ptr<'_> {131// SAFETY: the inner data will remain valid for as long as 'self.132unsafe { Ptr::new(self.data) }133}134135/// Gets a [`PtrMut`] to the start of the array136#[inline]137pub fn get_ptr_mut(&mut self) -> PtrMut<'_> {138// SAFETY: the inner data will remain valid for as long as 'self.139unsafe { PtrMut::new(self.data) }140}141142/// Get a slice of the first `slice_len` elements in [`BlobArray`] as if it were an array with elements of type `T`143/// To get a slice to the entire array, the caller must plug `len` in `slice_len`.144///145/// *`len` refers to the length of the array, the number of elements that have been initialized, and are safe to read.146/// Just like [`Vec::len`].*147///148/// # Safety149/// - The type `T` must be the type of the items in this [`BlobArray`].150/// - `slice_len` <= `len`151///152/// [`Vec::len`]: alloc::vec::Vec::len153pub unsafe fn get_sub_slice<T>(&self, slice_len: usize) -> &[UnsafeCell<T>] {154#[cfg(debug_assertions)]155debug_assert!(slice_len <= self.capacity);156// SAFETY: the inner data will remain valid for as long as 'self.157unsafe {158core::slice::from_raw_parts(self.data.as_ptr() as *const UnsafeCell<T>, slice_len)159}160}161162/// Clears the array, i.e. removing (and dropping) all of the elements.163/// Note that this method has no effect on the allocated capacity of the vector.164///165/// Note that this method will behave exactly the same as [`Vec::clear`].166///167/// # Safety168/// - For every element with index `i`, if `i` < `len`: It must be safe to call [`Self::get_unchecked_mut`] with `i`.169/// (If the safety requirements of every method that has been used on `Self` have been fulfilled, the caller just needs to ensure that `len` is correct.)170///171/// [`Vec::clear`]: alloc::vec::Vec::clear172pub unsafe fn clear(&mut self, len: usize) {173#[cfg(debug_assertions)]174debug_assert!(self.capacity >= len);175if let Some(drop) = self.drop {176// We set `self.drop` to `None` before dropping elements for unwind safety. This ensures we don't177// accidentally drop elements twice in the event of a drop impl panicking.178self.drop = None;179let size = self.item_layout.size();180for i in 0..len {181// SAFETY:182// * 0 <= `i` < `len`, so `i * size` must be in bounds for the allocation.183// * `size` is a multiple of the erased type's alignment,184// so adding a multiple of `size` will preserve alignment.185// * The item is left unreachable so it can be safely promoted to an `OwningPtr`.186let item = unsafe { self.get_ptr_mut().byte_add(i * size).promote() };187// SAFETY: `item` was obtained from this `BlobArray`, so its underlying type must match `drop`.188unsafe { drop(item) };189}190self.drop = Some(drop);191}192}193194/// Because this method needs parameters, it can't be the implementation of the `Drop` trait.195/// The owner of this [`BlobArray`] must call this method with the correct information.196///197/// # Safety198/// - `cap` and `len` are indeed the capacity and length of this [`BlobArray`]199/// - This [`BlobArray`] mustn't be used after calling this method.200pub unsafe fn drop(&mut self, cap: usize, len: usize) {201#[cfg(debug_assertions)]202debug_assert_eq!(self.capacity, cap);203if cap != 0 {204self.clear(len);205if !self.is_zst() {206let layout =207array_layout(&self.item_layout, cap).expect("array layout should be valid");208alloc::alloc::dealloc(self.data.as_ptr().cast(), layout);209}210#[cfg(debug_assertions)]211{212self.capacity = 0;213}214}215}216217/// Drops the last element in this [`BlobArray`].218///219/// # Safety220// - `last_element_index` must correspond to the last element in the array.221// - After this method is called, the last element must not be used222// unless [`Self::initialize_unchecked`] is called to set the value of the last element.223pub unsafe fn drop_last_element(&mut self, last_element_index: usize) {224#[cfg(debug_assertions)]225debug_assert!(self.capacity > last_element_index);226if let Some(drop) = self.drop {227// We set `self.drop` to `None` before dropping elements for unwind safety. This ensures we don't228// accidentally drop elements twice in the event of a drop impl panicking.229self.drop = None;230// SAFETY:231let item = self.get_unchecked_mut(last_element_index).promote();232// SAFETY:233unsafe { drop(item) };234self.drop = Some(drop);235}236}237238/// Allocate a block of memory for the array. This should be used to initialize the array, do not use this239/// method if there are already elements stored in the array - use [`Self::realloc`] instead.240///241/// # Panics242/// - Panics if the new capacity overflows `isize::MAX` bytes.243/// - Panics if the allocation causes an out-of-memory error.244pub(super) fn alloc(&mut self, capacity: NonZeroUsize) {245#[cfg(debug_assertions)]246debug_assert_eq!(self.capacity, 0);247if !self.is_zst() {248let new_layout = array_layout(&self.item_layout, capacity.get())249.expect("array layout should be valid");250// SAFETY: layout has non-zero size because capacity > 0, and the blob isn't ZST (`self.is_zst` == false)251let new_data = unsafe { alloc::alloc::alloc(new_layout) };252self.data = NonNull::new(new_data).unwrap_or_else(|| handle_alloc_error(new_layout));253}254#[cfg(debug_assertions)]255{256self.capacity = capacity.into();257}258}259260/// Reallocate memory for this array.261/// For example, if the length (number of stored elements) reached the capacity (number of elements the current allocation can store),262/// you might want to use this method to increase the allocation, so more data can be stored in the array.263///264/// # Panics265/// - Panics if the new capacity overflows `isize::MAX` bytes.266/// - Panics if the allocation causes an out-of-memory error.267///268/// # Safety269/// - `current_capacity` is indeed the current capacity of this array.270/// - After calling this method, the caller must update their saved capacity to reflect the change.271pub(super) unsafe fn realloc(272&mut self,273current_capacity: NonZeroUsize,274new_capacity: NonZeroUsize,275) {276#[cfg(debug_assertions)]277debug_assert_eq!(self.capacity, current_capacity.get());278if !self.is_zst() {279let new_layout = array_layout(&self.item_layout, new_capacity.get())280.expect("array layout should be valid");281// SAFETY:282// - ptr was be allocated via this allocator283// - the layout used to previously allocate this array is equivalent to `array_layout(&self.item_layout, current_capacity.get())`284// - `item_layout.size() > 0` (`self.is_zst`==false) and `new_capacity > 0`, so the layout size is non-zero285// - "new_size, when rounded up to the nearest multiple of layout.align(), must not overflow (i.e., the rounded value must be less than usize::MAX)",286// since the item size is always a multiple of its align, the rounding cannot happen287// here and the overflow is handled in `array_layout`288let new_data = unsafe {289alloc::alloc::realloc(290self.get_ptr_mut().as_ptr(),291// SAFETY: This is the Layout of the current array, it must be valid, if it hadn't have been, there would have been a panic on a previous allocation292array_layout_unchecked(&self.item_layout, current_capacity.get()),293new_layout.size(),294)295};296self.data = NonNull::new(new_data).unwrap_or_else(|| handle_alloc_error(new_layout));297}298#[cfg(debug_assertions)]299{300self.capacity = new_capacity.into();301}302}303304/// Initializes the value at `index` to `value`. This function does not do any bounds checking.305///306/// # Safety307/// - `index` must be in bounds (`index` < capacity)308/// - The [`Layout`] of the value must match the layout of the blobs stored in this array,309/// and it must be safe to use the `drop` function of this [`BlobArray`] to drop `value`.310/// - `value` must not point to the same value that is being initialized.311#[inline]312pub unsafe fn initialize_unchecked(&mut self, index: usize, value: OwningPtr<'_>) {313#[cfg(debug_assertions)]314debug_assert!(self.capacity > index);315let size = self.item_layout.size();316let dst = self.get_unchecked_mut(index);317core::ptr::copy::<u8>(value.as_ptr(), dst.as_ptr(), size);318}319320/// Replaces the value at `index` with `value`. This function does not do any bounds checking.321///322/// # Safety323/// - Index must be in-bounds (`index` < `len`)324/// - `value`'s [`Layout`] must match this [`BlobArray`]'s `item_layout`,325/// and it must be safe to use the `drop` function of this [`BlobArray`] to drop `value`.326/// - `value` must not point to the same value that is being replaced.327pub unsafe fn replace_unchecked(&mut self, index: usize, value: OwningPtr<'_>) {328#[cfg(debug_assertions)]329debug_assert!(self.capacity > index);330// Pointer to the value in the vector that will get replaced.331// SAFETY: The caller ensures that `index` fits in this vector.332let destination = NonNull::from(unsafe { self.get_unchecked_mut(index) });333let source = value.as_ptr();334335if let Some(drop) = self.drop {336// We set `self.drop` to `None` before dropping elements for unwind safety. This ensures we don't337// accidentally drop elements twice in the event of a drop impl panicking.338self.drop = None;339340// Transfer ownership of the old value out of the vector, so it can be dropped.341// SAFETY:342// - `destination` was obtained from a `PtrMut` in this vector, which ensures it is non-null,343// well-aligned for the underlying type, and has proper provenance.344// - The storage location will get overwritten with `value` later, which ensures345// that the element will not get observed or double dropped later.346// - If a panic occurs, `self.len` will remain `0`, which ensures a double-drop347// does not occur. Instead, all elements will be forgotten.348let old_value = unsafe { OwningPtr::new(destination) };349350// This closure will run in case `drop()` panics,351// which ensures that `value` does not get forgotten.352let on_unwind = OnDrop::new(|| drop(value));353354drop(old_value);355356// If the above code does not panic, make sure that `value` doesn't get dropped.357core::mem::forget(on_unwind);358359self.drop = Some(drop);360}361362// Copy the new value into the vector, overwriting the previous value.363// SAFETY:364// - `source` and `destination` were obtained from `OwningPtr`s, which ensures they are365// valid for both reads and writes.366// - The value behind `source` will only be dropped if the above branch panics,367// so it must still be initialized and it is safe to transfer ownership into the vector.368// - `source` and `destination` were obtained from different memory locations,369// both of which we have exclusive access to, so they are guaranteed not to overlap.370unsafe {371core::ptr::copy_nonoverlapping::<u8>(372source,373destination.as_ptr(),374self.item_layout.size(),375);376}377}378379/// This method will swap two elements in the array, and return the one at `index_to_remove`.380/// It is the caller's responsibility to drop the returned pointer, if that is desirable.381///382/// # Safety383/// - `index_to_keep` must be safe to access (within the bounds of the length of the array).384/// - `index_to_remove` must be safe to access (within the bounds of the length of the array).385/// - `index_to_remove` != `index_to_keep`386/// - The caller should address the inconsistent state of the array that has occurred after the swap, either:387/// 1) initialize a different value in `index_to_keep`388/// 2) update the saved length of the array if `index_to_keep` was the last element.389#[inline]390#[must_use = "The returned pointer should be used to drop the removed element"]391pub unsafe fn swap_remove_unchecked(392&mut self,393index_to_remove: usize,394index_to_keep: usize,395) -> OwningPtr<'_> {396#[cfg(debug_assertions)]397{398debug_assert!(self.capacity > index_to_keep);399debug_assert!(self.capacity > index_to_remove);400}401if index_to_remove != index_to_keep {402return self.swap_remove_unchecked_nonoverlapping(index_to_remove, index_to_keep);403}404// Now the element that used to be in index `index_to_remove` is now in index `index_to_keep` (after swap)405// If we are storing ZSTs than the index doesn't actually matter because the size is 0.406self.get_unchecked_mut(index_to_keep).promote()407}408409/// The same as [`Self::swap_remove_unchecked`] but the two elements must non-overlapping.410///411/// # Safety412/// - `index_to_keep` must be safe to access (within the bounds of the length of the array).413/// - `index_to_remove` must be safe to access (within the bounds of the length of the array).414/// - `index_to_remove` != `index_to_keep`415/// - The caller should address the inconsistent state of the array that has occurred after the swap, either:416/// 1) initialize a different value in `index_to_keep`417/// 2) update the saved length of the array if `index_to_keep` was the last element.418#[inline]419pub unsafe fn swap_remove_unchecked_nonoverlapping(420&mut self,421index_to_remove: usize,422index_to_keep: usize,423) -> OwningPtr<'_> {424#[cfg(debug_assertions)]425{426debug_assert!(self.capacity > index_to_keep);427debug_assert!(self.capacity > index_to_remove);428debug_assert_ne!(index_to_keep, index_to_remove);429}430debug_assert_ne!(index_to_keep, index_to_remove);431core::ptr::swap_nonoverlapping::<u8>(432self.get_unchecked_mut(index_to_keep).as_ptr(),433self.get_unchecked_mut(index_to_remove).as_ptr(),434self.item_layout.size(),435);436// Now the element that used to be in index `index_to_remove` is now in index `index_to_keep` (after swap)437// If we are storing ZSTs than the index doesn't actually matter because the size is 0.438self.get_unchecked_mut(index_to_keep).promote()439}440441/// This method will call [`Self::swap_remove_unchecked`] and drop the result.442///443/// # Safety444/// - `index_to_keep` must be safe to access (within the bounds of the length of the array).445/// - `index_to_remove` must be safe to access (within the bounds of the length of the array).446/// - `index_to_remove` != `index_to_keep`447/// - The caller should address the inconsistent state of the array that has occurred after the swap, either:448/// 1) initialize a different value in `index_to_keep`449/// 2) update the saved length of the array if `index_to_keep` was the last element.450#[inline]451pub unsafe fn swap_remove_and_drop_unchecked(452&mut self,453index_to_remove: usize,454index_to_keep: usize,455) {456#[cfg(debug_assertions)]457{458debug_assert!(self.capacity > index_to_keep);459debug_assert!(self.capacity > index_to_remove);460}461let drop = self.drop;462let value = self.swap_remove_unchecked(index_to_remove, index_to_keep);463if let Some(drop) = drop {464drop(value);465}466}467468/// The same as [`Self::swap_remove_and_drop_unchecked`] but the two elements must non-overlapping.469///470/// # Safety471/// - `index_to_keep` must be safe to access (within the bounds of the length of the array).472/// - `index_to_remove` must be safe to access (within the bounds of the length of the array).473/// - `index_to_remove` != `index_to_keep`474/// - The caller should address the inconsistent state of the array that has occurred after the swap, either:475/// 1) initialize a different value in `index_to_keep`476/// 2) update the saved length of the array if `index_to_keep` was the last element.477#[inline]478pub unsafe fn swap_remove_and_drop_unchecked_nonoverlapping(479&mut self,480index_to_remove: usize,481index_to_keep: usize,482) {483#[cfg(debug_assertions)]484{485debug_assert!(self.capacity > index_to_keep);486debug_assert!(self.capacity > index_to_remove);487debug_assert_ne!(index_to_keep, index_to_remove);488}489let drop = self.drop;490let value = self.swap_remove_unchecked_nonoverlapping(index_to_remove, index_to_keep);491if let Some(drop) = drop {492drop(value);493}494}495}496497/// From <https://doc.rust-lang.org/beta/src/core/alloc/layout.rs.html>498pub(super) fn array_layout(layout: &Layout, n: usize) -> Option<Layout> {499let (array_layout, offset) = repeat_layout(layout, n)?;500debug_assert_eq!(layout.size(), offset);501Some(array_layout)502}503504// TODO: replace with `Layout::repeat` if/when it stabilizes505/// From <https://doc.rust-lang.org/beta/src/core/alloc/layout.rs.html>506fn repeat_layout(layout: &Layout, n: usize) -> Option<(Layout, usize)> {507// This cannot overflow. Quoting from the invariant of Layout:508// > `size`, when rounded up to the nearest multiple of `align`,509// > must not overflow (i.e., the rounded value must be less than510// > `usize::MAX`)511let padded_size = layout.size() + padding_needed_for(layout, layout.align());512let alloc_size = padded_size.checked_mul(n)?;513514// SAFETY: self.align is already known to be valid and alloc_size has been515// padded already.516unsafe {517Some((518Layout::from_size_align_unchecked(alloc_size, layout.align()),519padded_size,520))521}522}523524/// From <https://doc.rust-lang.org/beta/src/core/alloc/layout.rs.html>525/// # Safety526/// The caller must ensure that:527/// - The resulting [`Layout`] is valid, by ensuring that `(layout.size() + padding_needed_for(layout, layout.align())) * n` doesn't overflow.528pub(super) unsafe fn array_layout_unchecked(layout: &Layout, n: usize) -> Layout {529let (array_layout, offset) = repeat_layout_unchecked(layout, n);530debug_assert_eq!(layout.size(), offset);531array_layout532}533534// TODO: replace with `Layout::repeat` if/when it stabilizes535/// From <https://doc.rust-lang.org/beta/src/core/alloc/layout.rs.html>536/// # Safety537/// The caller must ensure that:538/// - The resulting [`Layout`] is valid, by ensuring that `(layout.size() + padding_needed_for(layout, layout.align())) * n` doesn't overflow.539unsafe fn repeat_layout_unchecked(layout: &Layout, n: usize) -> (Layout, usize) {540// This cannot overflow. Quoting from the invariant of Layout:541// > `size`, when rounded up to the nearest multiple of `align`,542// > must not overflow (i.e., the rounded value must be less than543// > `usize::MAX`)544let padded_size = layout.size() + padding_needed_for(layout, layout.align());545// This may overflow in release builds, that's why this function is unsafe.546let alloc_size = padded_size * n;547548// SAFETY: self.align is already known to be valid and alloc_size has been549// padded already.550unsafe {551(552Layout::from_size_align_unchecked(alloc_size, layout.align()),553padded_size,554)555}556}557558/// From <https://doc.rust-lang.org/beta/src/core/alloc/layout.rs.html>559const fn padding_needed_for(layout: &Layout, align: usize) -> usize {560let len = layout.size();561562// Rounded up value is:563// len_rounded_up = (len + align - 1) & !(align - 1);564// and then we return the padding difference: `len_rounded_up - len`.565//566// We use modular arithmetic throughout:567//568// 1. align is guaranteed to be > 0, so align - 1 is always569// valid.570//571// 2. `len + align - 1` can overflow by at most `align - 1`,572// so the &-mask with `!(align - 1)` will ensure that in the573// case of overflow, `len_rounded_up` will itself be 0.574// Thus the returned padding, when added to `len`, yields 0,575// which trivially satisfies the alignment `align`.576//577// (Of course, attempts to allocate blocks of memory whose578// size and padding overflow in the above manner should cause579// the allocator to yield an error anyway.)580581let len_rounded_up = len.wrapping_add(align).wrapping_sub(1) & !align.wrapping_sub(1);582len_rounded_up.wrapping_sub(len)583}584585#[cfg(test)]586mod tests {587use bevy_ecs::prelude::*;588589#[derive(Component)]590struct PanicOnDrop;591592impl Drop for PanicOnDrop {593fn drop(&mut self) {594panic!("PanicOnDrop is being Dropped");595}596}597598#[test]599#[should_panic(expected = "PanicOnDrop is being Dropped")]600fn make_sure_zst_components_get_dropped() {601let mut world = World::new();602603world.spawn(PanicOnDrop);604}605}606607608