bevy_ecs/storage/
blob_array.rs

1use alloc::alloc::handle_alloc_error;
2use bevy_ptr::{OwningPtr, Ptr, PtrMut};
3use bevy_utils::OnDrop;
4use core::{alloc::Layout, cell::UnsafeCell, num::NonZeroUsize, ptr::NonNull};
5
6/// A flat, type-erased data storage type.
7///
8/// Used to densely store homogeneous ECS data. A blob is usually just an arbitrary block of contiguous memory without any identity, and
9/// 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,
10/// and a pointer to the location of the start of the array, similar to a C-style `void*` array.
11///
12/// This type is reliant on its owning type to store the capacity and length information.
13#[derive(Debug)]
14pub(super) struct BlobArray {
15    item_layout: Layout,
16    // the `data` ptr's layout is always `array_layout(item_layout, capacity)`
17    data: NonNull<u8>,
18    // None if the underlying type doesn't need to be dropped
19    pub drop: Option<unsafe fn(OwningPtr<'_>)>,
20    #[cfg(debug_assertions)]
21    capacity: usize,
22}
23
24impl BlobArray {
25    /// Create a new [`BlobArray`] with a specified `capacity`.
26    /// If `capacity` is 0, no allocations will be made.
27    ///
28    /// `drop` is an optional function pointer that is meant to be invoked when any element in the [`BlobArray`]
29    /// should be dropped. For all Rust-based types, this should match 1:1 with the implementation of [`Drop`]
30    /// if present, and should be `None` if `T: !Drop`. For non-Rust based types, this should match any cleanup
31    /// processes typically associated with the stored element.
32    ///
33    /// # Safety
34    /// `drop` should be safe to call with an [`OwningPtr`] pointing to any item that's been placed into this [`BlobArray`].
35    /// If `drop` is `None`, the items will be leaked. This should generally be set as None based on [`needs_drop`].
36    ///
37    /// [`needs_drop`]: std::mem::needs_drop
38    pub unsafe fn with_capacity(
39        item_layout: Layout,
40        drop_fn: Option<unsafe fn(OwningPtr<'_>)>,
41        capacity: usize,
42    ) -> Self {
43        if capacity == 0 {
44            let align = NonZeroUsize::new(item_layout.align()).expect("alignment must be > 0");
45
46            // Create a dangling pointer with the given alignment.
47            let data = NonNull::without_provenance(align);
48
49            Self {
50                item_layout,
51                drop: drop_fn,
52                data,
53                #[cfg(debug_assertions)]
54                capacity,
55            }
56        } else {
57            // SAFETY: Upheld by caller
58            let mut arr = unsafe { Self::with_capacity(item_layout, drop_fn, 0) };
59            // SAFETY: `capacity` > 0
60            unsafe { arr.alloc(NonZeroUsize::new_unchecked(capacity)) }
61            arr
62        }
63    }
64
65    /// Returns the [`Layout`] of the element type stored in the vector.
66    #[inline]
67    pub fn layout(&self) -> Layout {
68        self.item_layout
69    }
70
71    /// Return `true` if this [`BlobArray`] stores `ZSTs`.
72    pub fn is_zst(&self) -> bool {
73        self.item_layout.size() == 0
74    }
75
76    /// Returns the drop function for values stored in the vector,
77    /// or `None` if they don't need to be dropped.
78    #[inline]
79    pub fn get_drop(&self) -> Option<unsafe fn(OwningPtr<'_>)> {
80        self.drop
81    }
82
83    /// Returns a reference to the element at `index`, without doing bounds checking.
84    ///
85    /// *`len` refers to the length of the array, the number of elements that have been initialized, and are safe to read.
86    /// Just like [`Vec::len`].*
87    ///
88    /// # Safety
89    /// - The element at index `index` is safe to access.
90    ///   (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`)
91    ///
92    /// [`Vec::len`]: alloc::vec::Vec::len
93    #[inline]
94    pub unsafe fn get_unchecked(&self, index: usize) -> Ptr<'_> {
95        #[cfg(debug_assertions)]
96        debug_assert!(index < self.capacity);
97        let size = self.item_layout.size();
98        // SAFETY:
99        // - The caller ensures that `index` fits in this array,
100        //   so this operation will not overflow the original allocation.
101        // - `size` is a multiple of the erased type's alignment,
102        //   so adding a multiple of `size` will preserve alignment.
103        unsafe { self.get_ptr().byte_add(index * size) }
104    }
105
106    /// Returns a mutable reference to the element at `index`, without doing bounds checking.
107    ///
108    /// *`len` refers to the length of the array, the number of elements that have been initialized, and are safe to read.
109    /// Just like [`Vec::len`].*
110    ///
111    /// # Safety
112    /// - The element with at index `index` is safe to access.
113    ///   (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`)
114    ///
115    /// [`Vec::len`]: alloc::vec::Vec::len
116    #[inline]
117    pub unsafe fn get_unchecked_mut(&mut self, index: usize) -> PtrMut<'_> {
118        #[cfg(debug_assertions)]
119        debug_assert!(index < self.capacity);
120        let size = self.item_layout.size();
121        // SAFETY:
122        // - The caller ensures that `index` fits in this vector,
123        //   so this operation will not overflow the original allocation.
124        // - `size` is a multiple of the erased type's alignment,
125        //  so adding a multiple of `size` will preserve alignment.
126        unsafe { self.get_ptr_mut().byte_add(index * size) }
127    }
128
129    /// Gets a [`Ptr`] to the start of the array
130    #[inline]
131    pub fn get_ptr(&self) -> Ptr<'_> {
132        // SAFETY: the inner data will remain valid for as long as 'self.
133        unsafe { Ptr::new(self.data) }
134    }
135
136    /// Gets a [`PtrMut`] to the start of the array
137    #[inline]
138    pub fn get_ptr_mut(&mut self) -> PtrMut<'_> {
139        // SAFETY: the inner data will remain valid for as long as 'self.
140        unsafe { PtrMut::new(self.data) }
141    }
142
143    /// Get a slice of the first `slice_len` elements in [`BlobArray`] as if it were an array with elements of type `T`
144    /// To get a slice to the entire array, the caller must plug `len` in `slice_len`.
145    ///
146    /// *`len` refers to the length of the array, the number of elements that have been initialized, and are safe to read.
147    /// Just like [`Vec::len`].*
148    ///
149    /// # Safety
150    /// - The type `T` must be the type of the items in this [`BlobArray`].
151    /// - `slice_len` <= `len`
152    ///
153    /// [`Vec::len`]: alloc::vec::Vec::len
154    pub unsafe fn get_sub_slice<T>(&self, slice_len: usize) -> &[UnsafeCell<T>] {
155        #[cfg(debug_assertions)]
156        debug_assert!(slice_len <= self.capacity);
157        // SAFETY: the inner data will remain valid for as long as 'self.
158        unsafe {
159            core::slice::from_raw_parts(self.data.as_ptr() as *const UnsafeCell<T>, slice_len)
160        }
161    }
162
163    /// Clears the array, i.e. removing (and dropping) all of the elements.
164    /// Note that this method has no effect on the allocated capacity of the vector.
165    ///
166    /// Note that this method will behave exactly the same as [`Vec::clear`].
167    ///
168    /// # Safety
169    /// - For every element with index `i`, if `i` < `len`: It must be safe to call [`Self::get_unchecked_mut`] with `i`.
170    ///   (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.)
171    ///
172    /// [`Vec::clear`]: alloc::vec::Vec::clear
173    pub unsafe fn clear(&mut self, len: usize) {
174        #[cfg(debug_assertions)]
175        debug_assert!(self.capacity >= len);
176        if let Some(drop) = self.drop {
177            // We set `self.drop` to `None` before dropping elements for unwind safety. This ensures we don't
178            // accidentally drop elements twice in the event of a drop impl panicking.
179            self.drop = None;
180            let size = self.item_layout.size();
181            for i in 0..len {
182                // SAFETY:
183                // * 0 <= `i` < `len`, so `i * size` must be in bounds for the allocation.
184                // * `size` is a multiple of the erased type's alignment,
185                //   so adding a multiple of `size` will preserve alignment.
186                // * The item is left unreachable so it can be safely promoted to an `OwningPtr`.
187                let item = unsafe { self.get_ptr_mut().byte_add(i * size).promote() };
188                // SAFETY: `item` was obtained from this `BlobArray`, so its underlying type must match `drop`.
189                unsafe { drop(item) };
190            }
191            self.drop = Some(drop);
192        }
193    }
194
195    /// Because this method needs parameters, it can't be the implementation of the `Drop` trait.
196    /// The owner of this [`BlobArray`] must call this method with the correct information.
197    ///
198    /// # Safety
199    /// - `cap` and `len` are indeed the capacity and length of this [`BlobArray`]
200    /// - This [`BlobArray`] mustn't be used after calling this method.
201    pub unsafe fn drop(&mut self, cap: usize, len: usize) {
202        #[cfg(debug_assertions)]
203        debug_assert_eq!(self.capacity, cap);
204        if cap != 0 {
205            self.clear(len);
206            if !self.is_zst() {
207                let layout =
208                    array_layout(&self.item_layout, cap).expect("array layout should be valid");
209                alloc::alloc::dealloc(self.data.as_ptr().cast(), layout);
210            }
211            #[cfg(debug_assertions)]
212            {
213                self.capacity = 0;
214            }
215        }
216    }
217
218    /// Drops the last element in this [`BlobArray`].
219    ///
220    /// # Safety
221    // - `last_element_index` must correspond to the last element in the array.
222    // - After this method is called, the last element must not be used
223    // unless [`Self::initialize_unchecked`] is called to set the value of the last element.
224    pub unsafe fn drop_last_element(&mut self, last_element_index: usize) {
225        #[cfg(debug_assertions)]
226        debug_assert!(self.capacity > last_element_index);
227        if let Some(drop) = self.drop {
228            // We set `self.drop` to `None` before dropping elements for unwind safety. This ensures we don't
229            // accidentally drop elements twice in the event of a drop impl panicking.
230            self.drop = None;
231            // SAFETY:
232            let item = self.get_unchecked_mut(last_element_index).promote();
233            // SAFETY:
234            unsafe { drop(item) };
235            self.drop = Some(drop);
236        }
237    }
238
239    /// Allocate a block of memory for the array. This should be used to initialize the array, do not use this
240    /// method if there are already elements stored in the array - use [`Self::realloc`] instead.
241    ///
242    /// # Panics
243    /// - Panics if the new capacity overflows `isize::MAX` bytes.
244    /// - Panics if the allocation causes an out-of-memory error.
245    pub(super) fn alloc(&mut self, capacity: NonZeroUsize) {
246        #[cfg(debug_assertions)]
247        debug_assert_eq!(self.capacity, 0);
248        if !self.is_zst() {
249            let new_layout = array_layout(&self.item_layout, capacity.get())
250                .expect("array layout should be valid");
251            // SAFETY: layout has non-zero size because capacity > 0, and the blob isn't ZST (`self.is_zst` == false)
252            let new_data = unsafe { alloc::alloc::alloc(new_layout) };
253            self.data = NonNull::new(new_data).unwrap_or_else(|| handle_alloc_error(new_layout));
254        }
255        #[cfg(debug_assertions)]
256        {
257            self.capacity = capacity.into();
258        }
259    }
260
261    /// Reallocate memory for this array.
262    /// For example, if the length (number of stored elements) reached the capacity (number of elements the current allocation can store),
263    /// you might want to use this method to increase the allocation, so more data can be stored in the array.
264    ///
265    /// # Panics
266    /// - Panics if the new capacity overflows `isize::MAX` bytes.
267    /// - Panics if the allocation causes an out-of-memory error.
268    ///
269    /// # Safety
270    /// - `current_capacity` is indeed the current capacity of this array.
271    /// - After calling this method, the caller must update their saved capacity to reflect the change.
272    pub(super) unsafe fn realloc(
273        &mut self,
274        current_capacity: NonZeroUsize,
275        new_capacity: NonZeroUsize,
276    ) {
277        #[cfg(debug_assertions)]
278        debug_assert_eq!(self.capacity, current_capacity.get());
279        if !self.is_zst() {
280            let new_layout = array_layout(&self.item_layout, new_capacity.get())
281                .expect("array layout should be valid");
282            // SAFETY:
283            // - ptr was be allocated via this allocator
284            // - the layout used to previously allocate this array is equivalent to `array_layout(&self.item_layout, current_capacity.get())`
285            // - `item_layout.size() > 0` (`self.is_zst`==false) and `new_capacity > 0`, so the layout size is non-zero
286            // - "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)",
287            // since the item size is always a multiple of its align, the rounding cannot happen
288            // here and the overflow is handled in `array_layout`
289            let new_data = unsafe {
290                alloc::alloc::realloc(
291                    self.get_ptr_mut().as_ptr(),
292                    // 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 allocation
293                    array_layout_unchecked(&self.item_layout, current_capacity.get()),
294                    new_layout.size(),
295                )
296            };
297            self.data = NonNull::new(new_data).unwrap_or_else(|| handle_alloc_error(new_layout));
298        }
299        #[cfg(debug_assertions)]
300        {
301            self.capacity = new_capacity.into();
302        }
303    }
304
305    /// Initializes the value at `index` to `value`. This function does not do any bounds checking.
306    ///
307    /// # Safety
308    /// - `index` must be in bounds (`index` < capacity)
309    /// - The [`Layout`] of the value must match the layout of the blobs stored in this array,
310    ///   and it must be safe to use the `drop` function of this [`BlobArray`] to drop `value`.
311    /// - `value` must not point to the same value that is being initialized.
312    #[inline]
313    pub unsafe fn initialize_unchecked(&mut self, index: usize, value: OwningPtr<'_>) {
314        #[cfg(debug_assertions)]
315        debug_assert!(self.capacity > index);
316        let size = self.item_layout.size();
317        let dst = self.get_unchecked_mut(index);
318        core::ptr::copy::<u8>(value.as_ptr(), dst.as_ptr(), size);
319    }
320
321    /// Replaces the value at `index` with `value`. This function does not do any bounds checking.
322    ///
323    /// # Safety
324    /// - Index must be in-bounds (`index` < `len`)
325    /// - `value`'s [`Layout`] must match this [`BlobArray`]'s `item_layout`,
326    ///   and it must be safe to use the `drop` function of this [`BlobArray`] to drop `value`.
327    /// - `value` must not point to the same value that is being replaced.
328    pub unsafe fn replace_unchecked(&mut self, index: usize, value: OwningPtr<'_>) {
329        #[cfg(debug_assertions)]
330        debug_assert!(self.capacity > index);
331        // Pointer to the value in the vector that will get replaced.
332        // SAFETY: The caller ensures that `index` fits in this vector.
333        let destination = NonNull::from(unsafe { self.get_unchecked_mut(index) });
334        let source = value.as_ptr();
335
336        if let Some(drop) = self.drop {
337            // We set `self.drop` to `None` before dropping elements for unwind safety. This ensures we don't
338            // accidentally drop elements twice in the event of a drop impl panicking.
339            self.drop = None;
340
341            // Transfer ownership of the old value out of the vector, so it can be dropped.
342            // SAFETY:
343            // - `destination` was obtained from a `PtrMut` in this vector, which ensures it is non-null,
344            //   well-aligned for the underlying type, and has proper provenance.
345            // - The storage location will get overwritten with `value` later, which ensures
346            //   that the element will not get observed or double dropped later.
347            // - If a panic occurs, `self.len` will remain `0`, which ensures a double-drop
348            //   does not occur. Instead, all elements will be forgotten.
349            let old_value = unsafe { OwningPtr::new(destination) };
350
351            // This closure will run in case `drop()` panics,
352            // which ensures that `value` does not get forgotten.
353            let on_unwind = OnDrop::new(|| drop(value));
354
355            drop(old_value);
356
357            // If the above code does not panic, make sure that `value` doesn't get dropped.
358            core::mem::forget(on_unwind);
359
360            self.drop = Some(drop);
361        }
362
363        // Copy the new value into the vector, overwriting the previous value.
364        // SAFETY:
365        // - `source` and `destination` were obtained from `OwningPtr`s, which ensures they are
366        //   valid for both reads and writes.
367        // - The value behind `source` will only be dropped if the above branch panics,
368        //   so it must still be initialized and it is safe to transfer ownership into the vector.
369        // - `source` and `destination` were obtained from different memory locations,
370        //   both of which we have exclusive access to, so they are guaranteed not to overlap.
371        unsafe {
372            core::ptr::copy_nonoverlapping::<u8>(
373                source,
374                destination.as_ptr(),
375                self.item_layout.size(),
376            );
377        }
378    }
379
380    /// This method will swap two elements in the array, and return the one at `index_to_remove`.
381    /// It is the caller's responsibility to drop the returned pointer, if that is desirable.
382    ///
383    /// # Safety
384    /// - `index_to_keep` must be safe to access (within the bounds of the length of the array).
385    /// - `index_to_remove` must be safe to access (within the bounds of the length of the array).
386    /// - `index_to_remove` != `index_to_keep`
387    /// -  The caller should address the inconsistent state of the array that has occurred after the swap, either:
388    ///     1) initialize a different value in `index_to_keep`
389    ///     2) update the saved length of the array if `index_to_keep` was the last element.
390    #[inline]
391    #[must_use = "The returned pointer should be used to drop the removed element"]
392    pub unsafe fn swap_remove_unchecked(
393        &mut self,
394        index_to_remove: usize,
395        index_to_keep: usize,
396    ) -> OwningPtr<'_> {
397        #[cfg(debug_assertions)]
398        {
399            debug_assert!(self.capacity > index_to_keep);
400            debug_assert!(self.capacity > index_to_remove);
401        }
402        if index_to_remove != index_to_keep {
403            return self.swap_remove_unchecked_nonoverlapping(index_to_remove, index_to_keep);
404        }
405        // Now the element that used to be in index `index_to_remove` is now in index `index_to_keep` (after swap)
406        // If we are storing ZSTs than the index doesn't actually matter because the size is 0.
407        self.get_unchecked_mut(index_to_keep).promote()
408    }
409
410    /// The same as [`Self::swap_remove_unchecked`] but the two elements must non-overlapping.
411    ///
412    /// # Safety
413    /// - `index_to_keep` must be safe to access (within the bounds of the length of the array).
414    /// - `index_to_remove` must be safe to access (within the bounds of the length of the array).
415    /// - `index_to_remove` != `index_to_keep`
416    /// -  The caller should address the inconsistent state of the array that has occurred after the swap, either:
417    ///     1) initialize a different value in `index_to_keep`
418    ///     2) update the saved length of the array if `index_to_keep` was the last element.
419    #[inline]
420    pub unsafe fn swap_remove_unchecked_nonoverlapping(
421        &mut self,
422        index_to_remove: usize,
423        index_to_keep: usize,
424    ) -> OwningPtr<'_> {
425        #[cfg(debug_assertions)]
426        {
427            debug_assert!(self.capacity > index_to_keep);
428            debug_assert!(self.capacity > index_to_remove);
429            debug_assert_ne!(index_to_keep, index_to_remove);
430        }
431        debug_assert_ne!(index_to_keep, index_to_remove);
432        core::ptr::swap_nonoverlapping::<u8>(
433            self.get_unchecked_mut(index_to_keep).as_ptr(),
434            self.get_unchecked_mut(index_to_remove).as_ptr(),
435            self.item_layout.size(),
436        );
437        // Now the element that used to be in index `index_to_remove` is now in index `index_to_keep` (after swap)
438        // If we are storing ZSTs than the index doesn't actually matter because the size is 0.
439        self.get_unchecked_mut(index_to_keep).promote()
440    }
441
442    /// This method will call [`Self::swap_remove_unchecked`] and drop the result.
443    ///
444    /// # Safety
445    /// - `index_to_keep` must be safe to access (within the bounds of the length of the array).
446    /// - `index_to_remove` must be safe to access (within the bounds of the length of the array).
447    /// - `index_to_remove` != `index_to_keep`
448    /// -  The caller should address the inconsistent state of the array that has occurred after the swap, either:
449    ///     1) initialize a different value in `index_to_keep`
450    ///     2) update the saved length of the array if `index_to_keep` was the last element.
451    #[inline]
452    pub unsafe fn swap_remove_and_drop_unchecked(
453        &mut self,
454        index_to_remove: usize,
455        index_to_keep: usize,
456    ) {
457        #[cfg(debug_assertions)]
458        {
459            debug_assert!(self.capacity > index_to_keep);
460            debug_assert!(self.capacity > index_to_remove);
461        }
462        let drop = self.drop;
463        let value = self.swap_remove_unchecked(index_to_remove, index_to_keep);
464        if let Some(drop) = drop {
465            drop(value);
466        }
467    }
468
469    /// The same as [`Self::swap_remove_and_drop_unchecked`] but the two elements must non-overlapping.
470    ///
471    /// # Safety
472    /// - `index_to_keep` must be safe to access (within the bounds of the length of the array).
473    /// - `index_to_remove` must be safe to access (within the bounds of the length of the array).
474    /// - `index_to_remove` != `index_to_keep`
475    /// -  The caller should address the inconsistent state of the array that has occurred after the swap, either:
476    ///     1) initialize a different value in `index_to_keep`
477    ///     2) update the saved length of the array if `index_to_keep` was the last element.
478    #[inline]
479    pub unsafe fn swap_remove_and_drop_unchecked_nonoverlapping(
480        &mut self,
481        index_to_remove: usize,
482        index_to_keep: usize,
483    ) {
484        #[cfg(debug_assertions)]
485        {
486            debug_assert!(self.capacity > index_to_keep);
487            debug_assert!(self.capacity > index_to_remove);
488            debug_assert_ne!(index_to_keep, index_to_remove);
489        }
490        let drop = self.drop;
491        let value = self.swap_remove_unchecked_nonoverlapping(index_to_remove, index_to_keep);
492        if let Some(drop) = drop {
493            drop(value);
494        }
495    }
496}
497
498/// From <https://doc.rust-lang.org/beta/src/core/alloc/layout.rs.html>
499pub(super) fn array_layout(layout: &Layout, n: usize) -> Option<Layout> {
500    let (array_layout, offset) = repeat_layout(layout, n)?;
501    debug_assert_eq!(layout.size(), offset);
502    Some(array_layout)
503}
504
505// TODO: replace with `Layout::repeat` if/when it stabilizes
506/// From <https://doc.rust-lang.org/beta/src/core/alloc/layout.rs.html>
507fn repeat_layout(layout: &Layout, n: usize) -> Option<(Layout, usize)> {
508    // This cannot overflow. Quoting from the invariant of Layout:
509    // > `size`, when rounded up to the nearest multiple of `align`,
510    // > must not overflow (i.e., the rounded value must be less than
511    // > `usize::MAX`)
512    let padded_size = layout.size() + padding_needed_for(layout, layout.align());
513    let alloc_size = padded_size.checked_mul(n)?;
514
515    // SAFETY: self.align is already known to be valid and alloc_size has been
516    // padded already.
517    unsafe {
518        Some((
519            Layout::from_size_align_unchecked(alloc_size, layout.align()),
520            padded_size,
521        ))
522    }
523}
524
525/// From <https://doc.rust-lang.org/beta/src/core/alloc/layout.rs.html>
526/// # Safety
527/// The caller must ensure that:
528/// - The resulting [`Layout`] is valid, by ensuring that `(layout.size() + padding_needed_for(layout, layout.align())) * n` doesn't overflow.
529pub(super) unsafe fn array_layout_unchecked(layout: &Layout, n: usize) -> Layout {
530    let (array_layout, offset) = repeat_layout_unchecked(layout, n);
531    debug_assert_eq!(layout.size(), offset);
532    array_layout
533}
534
535// TODO: replace with `Layout::repeat` if/when it stabilizes
536/// From <https://doc.rust-lang.org/beta/src/core/alloc/layout.rs.html>
537/// # Safety
538/// The caller must ensure that:
539/// - The resulting [`Layout`] is valid, by ensuring that `(layout.size() + padding_needed_for(layout, layout.align())) * n` doesn't overflow.
540unsafe fn repeat_layout_unchecked(layout: &Layout, n: usize) -> (Layout, usize) {
541    // This cannot overflow. Quoting from the invariant of Layout:
542    // > `size`, when rounded up to the nearest multiple of `align`,
543    // > must not overflow (i.e., the rounded value must be less than
544    // > `usize::MAX`)
545    let padded_size = layout.size() + padding_needed_for(layout, layout.align());
546    // This may overflow in release builds, that's why this function is unsafe.
547    let alloc_size = padded_size * n;
548
549    // SAFETY: self.align is already known to be valid and alloc_size has been
550    // padded already.
551    unsafe {
552        (
553            Layout::from_size_align_unchecked(alloc_size, layout.align()),
554            padded_size,
555        )
556    }
557}
558
559/// From <https://doc.rust-lang.org/beta/src/core/alloc/layout.rs.html>
560const fn padding_needed_for(layout: &Layout, align: usize) -> usize {
561    let len = layout.size();
562
563    // Rounded up value is:
564    //   len_rounded_up = (len + align - 1) & !(align - 1);
565    // and then we return the padding difference: `len_rounded_up - len`.
566    //
567    // We use modular arithmetic throughout:
568    //
569    // 1. align is guaranteed to be > 0, so align - 1 is always
570    //    valid.
571    //
572    // 2. `len + align - 1` can overflow by at most `align - 1`,
573    //    so the &-mask with `!(align - 1)` will ensure that in the
574    //    case of overflow, `len_rounded_up` will itself be 0.
575    //    Thus the returned padding, when added to `len`, yields 0,
576    //    which trivially satisfies the alignment `align`.
577    //
578    // (Of course, attempts to allocate blocks of memory whose
579    // size and padding overflow in the above manner should cause
580    // the allocator to yield an error anyway.)
581
582    let len_rounded_up = len.wrapping_add(align).wrapping_sub(1) & !align.wrapping_sub(1);
583    len_rounded_up.wrapping_sub(len)
584}
585
586#[cfg(test)]
587mod tests {
588    use bevy_ecs::prelude::*;
589
590    #[derive(Component)]
591    struct PanicOnDrop;
592
593    impl Drop for PanicOnDrop {
594        fn drop(&mut self) {
595            panic!("PanicOnDrop is being Dropped");
596        }
597    }
598
599    #[test]
600    #[should_panic(expected = "PanicOnDrop is being Dropped")]
601    fn make_sure_zst_components_get_dropped() {
602        let mut world = World::new();
603
604        world.spawn(PanicOnDrop);
605    }
606}