Use the capacity from the entities Vec to initialize Table columns (#20528)

# Objective
When working with `realloc`, it's a safety invariant to pass in the
existing layout of the allocation that is being reallocated. This may
not be the case with newly created `Table`s. `Vec::with_capacity`'s
documentation states that it will return an allocation with enough space
for *at least* `capacity` elements, not exactly `capacity`. This means
that `entities.capacity()` may be greater than the provided capacity. As
the `ThinColumn`s use this as their capacity, the new Layout fed to
`realloc` will not match the allocation originally provided to `alloc`.
This is unsound.

While investigating this, I also found that we were not validating that
the total capacity of `BlobArray`'s layout upon reallocation were less
than `isize::MAX` via `array_layout_unchecked`.

## Solution
Begin `Table` construction by allocating the `entities` Vec, and use
it's capacity to allocate the columns instead of directly feeding the
provided capacity into `ThinColumn::with_capacity`.

Replace the `array_layout_unchecked` call with a safe call to
`array_layout`, and panic if it fails.

## Testing
Tested this locally against existing unit tests and miri.

---------

Co-authored-by: Giacomo Stevanato <[email protected]>

Author: james7132

crates/bevy_ecs/src/storage/blob_array.rs

@@ -258,9 +258,8 @@ impl BlobArray {
         #[cfg(debug_assertions)]
         debug_assert_eq!(self.capacity, current_capacity.get());
         if !self.is_zst() {
-            // SAFETY: `new_capacity` can't overflow usize
-            let new_layout =
-                unsafe { array_layout_unchecked(&self.item_layout, new_capacity.get()) };
+            let new_layout = array_layout(&self.item_layout, new_capacity.get())
+                .expect("array layout should be valid");
             // SAFETY:
             // - ptr was be allocated via this allocator
             // - the layout used to previously allocate this array is equivalent to `array_layout(&self.item_layout, current_capacity.get())`

crates/bevy_ecs/src/storage/table/mod.rs

@@ -132,17 +132,22 @@ impl TableRow {
 /// [`with_capacity`]: Self::with_capacity
 /// [`add_column`]: Self::add_column
 /// [`build`]: Self::build
+//
+// # Safety
+// The capacity of all columns is determined by that of the `entities` Vec. This means that
+// it must be the correct capacity to allocate, reallocate, and deallocate all columns. This
+// means the safety invariant must be enforced even in `TableBuilder`.
 pub(crate) struct TableBuilder {
     columns: SparseSet<ComponentId, ThinColumn>,
-    capacity: usize,
+    entities: Vec<Entity>,
 }
 
 impl TableBuilder {
     /// Start building a new [`Table`] with a specified `column_capacity` (How many components per column?) and a `capacity` (How many columns?)
     pub fn with_capacity(capacity: usize, column_capacity: usize) -> Self {
         Self {
             columns: SparseSet::with_capacity(column_capacity),
-            capacity,
+            entities: Vec::with_capacity(capacity),
         }
     }
 
@@ -151,7 +156,7 @@ impl TableBuilder {
     pub fn add_column(mut self, component_info: &ComponentInfo) -> Self {
         self.columns.insert(
             component_info.id(),
-            ThinColumn::with_capacity(component_info, self.capacity),
+            ThinColumn::with_capacity(component_info, self.entities.capacity()),
         );
         self
     }
@@ -161,7 +166,7 @@ impl TableBuilder {
     pub fn build(self) -> Table {
         Table {
             columns: self.columns.into_immutable(),
-            entities: Vec::with_capacity(self.capacity),
+            entities: self.entities,
         }
     }
 }
@@ -178,6 +183,11 @@ impl TableBuilder {
 /// [structure-of-arrays]: https://en.wikipedia.org/wiki/AoS_and_SoA#Structure_of_arrays
 /// [`Component`]: crate::component::Component
 /// [`World`]: crate::world::World
+//
+// # Safety
+// The capacity of all columns is determined by that of the `entities` Vec. This means that
+// it must be the correct capacity to allocate, reallocate, and deallocate all columns. This
+// means the safety invariant must be enforced even in `TableBuilder`.
 pub struct Table {
     columns: ImmutableSparseSet<ComponentId, ThinColumn>,
     entities: Vec<Entity>,
@@ -529,6 +539,11 @@ impl Table {
     /// Allocate memory for the columns in the [`Table`]
     ///
     /// The current capacity of the columns should be 0, if it's not 0, then the previous data will be overwritten and leaked.
+    ///
+    /// # Safety
+    /// The capacity of all columns is determined by that of the `entities` Vec. This means that
+    /// it must be the correct capacity to allocate, reallocate, and deallocate all columns. This
+    /// means the safety invariant must be enforced even in `TableBuilder`.
     fn alloc_columns(&mut self, new_capacity: NonZeroUsize) {
         // If any of these allocations trigger an unwind, the wrong capacity will be used while dropping this table - UB.
         // To avoid this, we use `AbortOnPanic`. If the allocation triggered a panic, the `AbortOnPanic`'s Drop impl will be
@@ -544,6 +559,10 @@ impl Table {
     ///
     /// # Safety
     /// - `current_column_capacity` is indeed the capacity of the columns
+    ///
+    /// The capacity of all columns is determined by that of the `entities` Vec. This means that
+    /// it must be the correct capacity to allocate, reallocate, and deallocate all columnts. This
+    /// means the safety invariant must be enforced even in `TableBuilder`.
     unsafe fn realloc_columns(
         &mut self,
         current_column_capacity: NonZeroUsize,