Use the capacity from the entities Vec to initialize Table columns

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,6 +132,11 @@ 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 columnts. This
+// means the safety invariant must be enforced even in `TableBuilder`.
 pub(crate) struct TableBuilder {
     columns: SparseSet<ComponentId, ThinColumn>,
     entities: Vec<Entity>,
@@ -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 columnts. 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 columnts. 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,