Don't panic in the ECS where we know a fetch is valid

crates/bevy_ecs/src/archetype.rs

@@ -23,11 +23,12 @@
     bundle::BundleId,
     component::{ComponentId, Components, StorageType},
     entity::{Entity, EntityLocation},
+    query::{DebugCheckedUnwrap, UnsafeVecExtensions},
     storage::{ImmutableSparseSet, SparseArray, SparseSet, SparseSetIndex, TableId, TableRow},
 };
 use std::{
     hash::Hash,
-    ops::{Index, IndexMut, RangeFrom},
+    ops::{Index, RangeFrom},
 };
 
 /// An opaque location within a [`Archetype`].
@@ -516,18 +517,24 @@
     /// Removes the entity at `index` by swapping it out. Returns the table row the entity is stored
     /// in.
     ///
-    /// # Panics
-    /// This function will panic if `index >= self.len()`
+    /// # Safety
+    /// `index < self.len()` must be satisifed.
     #[inline]
-    pub(crate) fn swap_remove(&mut self, row: ArchetypeRow) -> ArchetypeSwapRemoveResult {
+    pub(crate) unsafe fn swap_remove_unchecked(
+        &mut self,
+        row: ArchetypeRow,
+    ) -> ArchetypeSwapRemoveResult {
         let is_last = row.index() == self.entities.len() - 1;
-        let entity = self.entities.swap_remove(row.index());
+        // SAFETY: Caller assures that `row` is in bounds.
+        let entity = unsafe { self.entities.swap_remove_unchecked(row.index()) };
         ArchetypeSwapRemoveResult {
-            swapped_entity: if is_last {
-                None
-            } else {
-                Some(self.entities[row.index()].entity)
-            },
+            swapped_entity: (!is_last).then(|| {
+                let archetype_entity =
+                    // SAFETY: Caller assures that `row` is in bounds, and remains in bounds
+                    // after the prior swap_remove call.
+                    unsafe { self.entities.get(row.index()).debug_checked_unwrap() };
+                archetype_entity.entity
+            }),
             table_row: entity.table_row,
         }
     }
@@ -749,22 +756,32 @@
         self.archetypes.get(id.index())
     }
 
-    /// # Panics
+    /// Fetches an mutable reference to an [`Archetype`] using its
+    /// ID.
     ///
-    /// Panics if `a` and `b` are equal.
+    /// # Safety
+    /// `id` must be valid.
+    #[inline]
+    pub(crate) unsafe fn get_unchecked_mut(&mut self, id: ArchetypeId) -> &mut Archetype {
+        self.archetypes.get_mut(id.index()).debug_checked_unwrap()
+    }
+
+    /// # Safety
+    /// `a` and `b` must both be in bounds and not equal to each other.
     #[inline]
-    pub(crate) fn get_2_mut(
+    pub(crate) unsafe fn get_2_unchecked_mut(
         &mut self,
         a: ArchetypeId,
         b: ArchetypeId,
     ) -> (&mut Archetype, &mut Archetype) {
-        if a.index() > b.index() {
-            let (b_slice, a_slice) = self.archetypes.split_at_mut(a.index());
-            (&mut a_slice[0], &mut b_slice[b.index()])
-        } else {
-            let (a_slice, b_slice) = self.archetypes.split_at_mut(b.index());
-            (&mut a_slice[a.index()], &mut b_slice[0])
-        }
+        debug_assert!(
+            a != b && a.index() < self.archetypes.len() && b.index() < self.archetypes.len()
+        );
+        let ptr = self.archetypes.as_mut_ptr();
+        (
+            ptr.add(a.index()).as_mut().debug_checked_unwrap(),
+            ptr.add(b.index()).as_mut().debug_checked_unwrap(),
+        )
     }
 
     /// Returns a read-only iterator over all archetypes.
@@ -842,18 +859,3 @@
         &self.archetypes[index.start.0.index()..]
     }
 }
-impl Index<ArchetypeId> for Archetypes {
-    type Output = Archetype;
-
-    #[inline]
-    fn index(&self, index: ArchetypeId) -> &Self::Output {
-        &self.archetypes[index.index()]
-    }
-}
-
-impl IndexMut<ArchetypeId> for Archetypes {
-    #[inline]
-    fn index_mut(&mut self, index: ArchetypeId) -> &mut Self::Output {
-        &mut self.archetypes[index.index()]
-    }
-}

crates/bevy_ecs/src/bundle.rs

@@ -420,6 +420,7 @@ impl BundleInfo {
     /// Adds a bundle to the given archetype and returns the resulting archetype. This could be the
     /// same [`ArchetypeId`], in the event that adding the given bundle does not result in an
     /// [`Archetype`] change. Results are cached in the [`Archetype`] graph to avoid redundant work.
+    ///
     /// # Safety
     /// `components` must be the same components as passed in [`Self::new`]
     pub(crate) unsafe fn add_bundle_to_archetype(
@@ -429,14 +430,19 @@ impl BundleInfo {
         components: &Components,
         archetype_id: ArchetypeId,
     ) -> ArchetypeId {
-        if let Some(add_bundle_id) = archetypes[archetype_id].edges().get_add_bundle(self.id) {
-            return add_bundle_id;
+        {
+            // SAFETY:  Caller guarantees that archetype_id is valid.
+            let current_archetype = unsafe { archetypes.get(archetype_id).debug_checked_unwrap() };
+            if let Some(add_bundle_id) = current_archetype.edges().get_add_bundle(self.id) {
+                return add_bundle_id;
+            }
         }
         let mut new_table_components = Vec::new();
         let mut new_sparse_set_components = Vec::new();
         let mut bundle_status = Vec::with_capacity(self.component_ids.len());
 
-        let current_archetype = &mut archetypes[archetype_id];
+        // SAFETY:  Caller guarantees that archetype_id is valid.
+        let current_archetype = unsafe { archetypes.get_unchecked_mut(archetype_id) };
         for component_id in self.component_ids.iter().cloned() {
             if current_archetype.contains(component_id) {
                 bundle_status.push(ComponentStatus::Mutated);
@@ -462,7 +468,9 @@ impl BundleInfo {
             let sparse_set_components;
             // the archetype changes when we add this bundle. prepare the new archetype and storages
             {
-                let current_archetype = &archetypes[archetype_id];
+                // SAFETY:  Caller guarantees that archetype_id is valid.
+                let current_archetype =
+                    unsafe { archetypes.get(archetype_id).debug_checked_unwrap() };
                 table_components = if new_table_components.is_empty() {
                     // if there are no new table components, we can keep using this table
                     table_id = current_archetype.table_id();
@@ -497,12 +505,16 @@ impl BundleInfo {
                 table_components,
                 sparse_set_components,
             );
+
             // add an edge from the old archetype to the new archetype
-            archetypes[archetype_id].edges_mut().insert_add_bundle(
-                self.id,
-                new_archetype_id,
-                bundle_status,
-            );
+            // SAFETY:  Caller guarantees that archetype_id is valid.
+            unsafe {
+                archetypes
+                    .get_unchecked_mut(archetype_id)
+                    .edges_mut()
+                    .insert_add_bundle(self.id, new_archetype_id, bundle_status);
+            }
+
             new_archetype_id
         }
     }
@@ -565,7 +577,7 @@ impl<'w> BundleInserter<'w> {
             archetype_id,
         );
         if new_archetype_id == archetype_id {
-            let archetype = &mut world.archetypes[archetype_id];
+            let archetype = world.archetypes.get_unchecked_mut(archetype_id);
             // SAFETY: The edge is assured to be initialized when we called add_bundle_to_archetype
             let add_bundle = unsafe {
                 archetype
@@ -574,7 +586,7 @@ impl<'w> BundleInserter<'w> {
                     .debug_checked_unwrap()
             };
             let table_id = archetype.table_id();
-            let table = &mut world.storages.tables[table_id];
+            let table = world.storages.tables.get_unchecked_mut(table_id);
             Self {
                 add_bundle: add_bundle.into(),
                 archetype: archetype.into(),
@@ -585,8 +597,9 @@ impl<'w> BundleInserter<'w> {
                 world: world.as_unsafe_world_cell(),
             }
         } else {
-            let (archetype, new_archetype) =
-                world.archetypes.get_2_mut(archetype_id, new_archetype_id);
+            let (archetype, new_archetype) = world
+                .archetypes
+                .get_2_unchecked_mut(archetype_id, new_archetype_id);
             // SAFETY: The edge is assured to be initialized when we called add_bundle_to_archetype
             let add_bundle = unsafe {
                 archetype
@@ -597,7 +610,7 @@ impl<'w> BundleInserter<'w> {
             let table_id = archetype.table_id();
             let new_table_id = new_archetype.table_id();
             if table_id == new_table_id {
-                let table = &mut world.storages.tables[table_id];
+                let table = world.storages.tables.get_unchecked_mut(table_id);
                 Self {
                     add_bundle: add_bundle.into(),
                     archetype: archetype.into(),
@@ -610,7 +623,10 @@ impl<'w> BundleInserter<'w> {
                     world: world.as_unsafe_world_cell(),
                 }
             } else {
-                let (table, new_table) = world.storages.tables.get_2_mut(table_id, new_table_id);
+                let (table, new_table) = world
+                    .storages
+                    .tables
+                    .get_2_unchecked_mut(table_id, new_table_id);
                 Self {
                     add_bundle: add_bundle.into(),
                     archetype: archetype.into(),
@@ -671,7 +687,7 @@ impl<'w> BundleInserter<'w> {
                     (&mut world.storages.sparse_sets, &mut world.entities)
                 };
 
-                let result = archetype.swap_remove(location.archetype_row);
+                let result = archetype.swap_remove_unchecked(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.
@@ -717,7 +733,7 @@ impl<'w> BundleInserter<'w> {
                         &mut world.entities,
                     )
                 };
-                let result = archetype.swap_remove(location.archetype_row);
+                let result = archetype.swap_remove_unchecked(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.
@@ -850,8 +866,11 @@ impl<'w> BundleSpawner<'w> {
             &world.components,
             ArchetypeId::EMPTY,
         );
-        let archetype = &mut world.archetypes[new_archetype_id];
-        let table = &mut world.storages.tables[archetype.table_id()];
+        let archetype = world.archetypes.get_unchecked_mut(new_archetype_id);
+        let table = world
+            .storages
+            .tables
+            .get_unchecked_mut(archetype.table_id());
         Self {
             bundle_info: bundle_info.into(),
             table: table.into(),

crates/bevy_ecs/src/entity/mod.rs

@@ -53,6 +53,7 @@ use crate::{
         masks::{IdentifierMask, HIGH_MASK},
         Identifier,
     },
+    query::UnsafeVecExtensions,
     storage::{SparseSetIndex, TableId, TableRow},
 };
 use serde::{Deserialize, Serialize};
@@ -617,7 +618,9 @@ impl Entities {
             self.len += 1;
             None
         } else if let Some(index) = self.pending.iter().position(|item| *item == entity.index()) {
-            self.pending.swap_remove(index);
+            // SAFETY: The returned index means there is at least one item in the pending Vec and that
+            // index is within bounds.
+            unsafe { self.pending.swap_remove_unchecked(index) };
             let new_free_cursor = self.pending.len() as IdCursor;
             *self.free_cursor.get_mut() = new_free_cursor;
             self.len += 1;
@@ -653,7 +656,9 @@ impl Entities {
             self.len += 1;
             AllocAtWithoutReplacement::DidNotExist
         } else if let Some(index) = self.pending.iter().position(|item| *item == entity.index()) {
-            self.pending.swap_remove(index);
+            // SAFETY: The returned index means there is at least one item in the pending Vec and that
+            // index is within bounds.
+            unsafe { self.pending.swap_remove_unchecked(index) };
             let new_free_cursor = self.pending.len() as IdCursor;
             *self.free_cursor.get_mut() = new_free_cursor;
             self.len += 1;

crates/bevy_ecs/src/query/iter.rs

@@ -743,10 +743,21 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> {
         let ids = self.storage_id_iter.clone();
         let remaining_matched: usize = if Self::IS_DENSE {
             // SAFETY: The if check ensures that storage_id_iter stores TableIds
-            unsafe { ids.map(|id| tables[id.table_id].entity_count()).sum() }
+            unsafe {
+                ids.map(|id| {
+                    tables
+                        .get(id.table_id)
+                        .debug_checked_unwrap()
+                        .entity_count()
+                })
+                .sum()
+            }
         } else {
             // SAFETY: The if check ensures that storage_id_iter stores ArchetypeIds
-            unsafe { ids.map(|id| archetypes[id.archetype_id].len()).sum() }
+            unsafe {
+                ids.map(|id| archetypes.get(id.archetype_id).debug_checked_unwrap().len())
+                    .sum()
+            }
         };
         remaining_matched + self.current_len - self.current_row
     }

crates/bevy_ecs/src/query/mod.rs

@@ -67,6 +67,27 @@ impl<T> DebugCheckedUnwrap for Option<T> {
     }
 }
 
+pub(crate) trait UnsafeVecExtensions<T> {
+    unsafe fn swap_remove_unchecked(&mut self, index: usize) -> T;
+}
+
+impl<T> UnsafeVecExtensions<T> for Vec<T> {
+    /// # Safety
+    /// `index` must be in the bounds `0 <= index < self.len()`.
+    unsafe fn swap_remove_unchecked(&mut self, index: usize) -> T {
+        let len = self.len();
+        // SAFETY: Caller guarantees that `index` is in bounds. We replace self[index] with the last element.
+        // If index is in bounds, there must be a last element (which can be self[index] itself).
+        unsafe {
+            let value = std::ptr::read(self.as_ptr().add(index));
+            let base_ptr = self.as_mut_ptr();
+            std::ptr::copy(base_ptr.add(len - 1), base_ptr.add(index), 1);
+            self.set_len(len - 1);
+            value
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use bevy_ecs_macros::{QueryData, QueryFilter};

crates/bevy_ecs/src/query/par_iter.rs

@@ -152,6 +152,8 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryParIter<'w, 's, D, F> {
 
     #[cfg(all(not(target_arch = "wasm32"), feature = "multi-threaded"))]
     fn get_batch_size(&self, thread_count: usize) -> usize {
+        use crate::query::DebugCheckedUnwrap;
+
         if self.batching_strategy.batch_size_limits.is_empty() {
             return self.batching_strategy.batch_size_limits.start;
         }
@@ -166,13 +168,18 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryParIter<'w, 's, D, F> {
             let tables = unsafe { &self.world.world_metadata().storages().tables };
             id_iter
                 // SAFETY: The if check ensures that matched_storage_ids stores TableIds
-                .map(|id| unsafe { tables[id.table_id].entity_count() })
+                .map(|id| unsafe {
+                    tables
+                        .get(id.table_id)
+                        .debug_checked_unwrap()
+                        .entity_count()
+                })
                 .max()
         } else {
             let archetypes = &self.world.archetypes();
             id_iter
                 // SAFETY: The if check ensures that matched_storage_ids stores ArchetypeIds
-                .map(|id| unsafe { archetypes[id.archetype_id].len() })
+                .map(|id| unsafe { archetypes.get(id.archetype_id).debug_checked_unwrap().len() })
                 .max()
         };
         let max_size = max_size.unwrap_or(0);

crates/bevy_ecs/src/query/state.rs

@@ -1395,7 +1395,9 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
             for storage_id in &self.matched_storage_ids {
                 if D::IS_DENSE && F::IS_DENSE {
                     let table_id = storage_id.table_id;
-                    let table = &tables[table_id];
+                    // SAFETY: The table has been matched with the query and tables cannot be deleted,
+                    // and so table_id must be valid.
+                    let table = unsafe { tables.get(table_id).debug_checked_unwrap() };
                     if table.is_empty() {
                         continue;
                     }
@@ -1417,7 +1419,9 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
                     }
                 } else {
                     let archetype_id = storage_id.archetype_id;
-                    let archetype = &archetypes[archetype_id];
+                    // SAFETY: The table has been matched with the query and tables cannot be deleted,
+                    // and so archetype_id must be valid.
+                    let archetype = unsafe { archetypes.get(archetype_id).debug_checked_unwrap() };
                     if archetype.is_empty() {
                         continue;
                     }