From 1c0b8dae26029d22354a48d0185d78377bab5a6b Mon Sep 17 00:00:00 2001
From: james7132 <contact@jamessliu.com>
Date: Sun, 17 Mar 2024 01:21:46 -0700
Subject: [PATCH] Remove panics from SparseSet construction

---
 crates/bevy_ecs/src/archetype.rs          |  1 -
 crates/bevy_ecs/src/storage/sparse_set.rs | 24 ++++++++++++++---------
 crates/bevy_ecs/src/world/entity_ref.rs   |  4 ++--
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/crates/bevy_ecs/src/archetype.rs b/crates/bevy_ecs/src/archetype.rs
index b4098276a59cd..1beb719944162 100644
--- a/crates/bevy_ecs/src/archetype.rs
+++ b/crates/bevy_ecs/src/archetype.rs
@@ -26,7 +26,6 @@ use crate::{
     query::{DebugCheckedUnwrap, UnsafeVecExtensions},
     storage::{ImmutableSparseSet, SparseArray, SparseSet, SparseSetIndex, TableId, TableRow},
 };
-use core::arch;
 use std::{
     hash::Hash,
     ops::{Index, RangeFrom},
diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs
index 67918343fe85b..8f0a19238cc96 100644
--- a/crates/bevy_ecs/src/storage/sparse_set.rs
+++ b/crates/bevy_ecs/src/storage/sparse_set.rs
@@ -1,7 +1,7 @@
 use crate::{
     component::{ComponentId, ComponentInfo, ComponentTicks, Tick, TickCells},
     entity::Entity,
-    query::UnsafeVecExtensions,
+    query::{DebugCheckedUnwrap, UnsafeVecExtensions},
     storage::{Column, TableRow},
 };
 use bevy_ptr::{OwningPtr, Ptr};
@@ -289,12 +289,14 @@ impl ComponentSparseSet {
             // SAFETY: dense_index was just removed from `sparse`, which ensures that it is valid
             let (value, _) = unsafe { self.dense.swap_remove_and_forget_unchecked(dense_index) };
             if !is_last {
-                let swapped_entity = self.entities[dense_index.as_usize()];
+                // SAFETY: This index was just swapped to above, it must be valid.
+                let swapped_entity = unsafe { self.entities.get_unchecked(dense_index.as_usize()) };
                 #[cfg(not(debug_assertions))]
-                let index = swapped_entity;
+                let index = *swapped_entity;
                 #[cfg(debug_assertions)]
                 let index = swapped_entity.index();
-                *self.sparse.get_mut(index).unwrap() = dense_index;
+                // SAFETY: if the sparse index points to something in the dense vec, it exists
+                unsafe { *self.sparse.get_mut(index).debug_checked_unwrap() = dense_index }
             }
             value
         })
@@ -316,12 +318,14 @@ impl ComponentSparseSet {
                 self.dense.swap_remove_unchecked(dense_index);
             }
             if !is_last {
-                let swapped_entity = self.entities[dense_index.as_usize()];
+                // SAFETY: This index was just swapped to above, it must be valid.
+                let swapped_entity = unsafe { self.entities.get_unchecked(dense_index.as_usize()) };
                 #[cfg(not(debug_assertions))]
-                let index = swapped_entity;
+                let index = *swapped_entity;
                 #[cfg(debug_assertions)]
                 let index = swapped_entity.index();
-                *self.sparse.get_mut(index).unwrap() = dense_index;
+                // SAFETY: if the sparse index points to something in the dense vec, it exists
+                unsafe { *self.sparse.get_mut(index).debug_checked_unwrap() = dense_index }
             }
             true
         } else {
@@ -508,8 +512,10 @@ impl<I: SparseSetIndex, V> SparseSet<I, V> {
             // in the indicies Vec.
             unsafe { self.indices.swap_remove_unchecked(index) };
             if !is_last {
-                let swapped_index = self.indices[index].clone();
-                *self.sparse.get_mut(swapped_index).unwrap() = dense_index;
+                // SAFETY: This index was just swapped to above, it must be valid.
+                let swapped_index = unsafe { self.indices.get_unchecked(index).clone() };
+                // SAFETY: The swapped index must be valid in the sparse array.
+                unsafe { *self.sparse.get_mut(swapped_index).debug_checked_unwrap() = dense_index };
             }
             value
         })
diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs
index 2557ef2c8f728..66fad17d1c119 100644
--- a/crates/bevy_ecs/src/world/entity_ref.rs
+++ b/crates/bevy_ecs/src/world/entity_ref.rs
@@ -1562,8 +1562,8 @@ impl<'w, 'a, T: Component> OccupiedEntry<'w, 'a, T> {
     /// ```
     #[inline]
     pub fn into_mut(self) -> Mut<'a, T> {
-        // This shouldn't panic because if we have an OccupiedEntry the component must exist.
-        self.entity_world.get_mut().unwrap()
+        // SAFETY: If we have an OccupiedEntry the component must exist.
+        unsafe { self.entity_world.get_mut().debug_checked_unwrap() }
     }
 
     /// Replaces the component of the entry.