From cafc95013857a2fbbbdbb3df20b5e6e4ba77e040 Mon Sep 17 00:00:00 2001
From: Urben1680 <urben1680@gmail.com>
Date: Mon, 14 Jul 2025 21:45:34 +0200
Subject: [PATCH] assert no errors before doing mutations

---
 crates/bevy_ecs/src/bundle/info.rs | 48 ++++++++++++++++++++----------
 crates/bevy_ecs/src/world/mod.rs   | 19 +++++++++---
 2 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/crates/bevy_ecs/src/bundle/info.rs b/crates/bevy_ecs/src/bundle/info.rs
index 6805cba2b8..44b8688536 100644
--- a/crates/bevy_ecs/src/bundle/info.rs
+++ b/crates/bevy_ecs/src/bundle/info.rs
@@ -625,25 +625,24 @@ impl Bundles {
         *bundle_id
     }
 
-    /// Updates the required components of bundles that contain `requiree`.
+    /// Checks if the bundles containing `requiree` can have their required components be updated,
+    /// in which case this returns `Ok(true)`.
     ///
-    /// # Safety
+    /// Returns `Ok(false)` if there are no known bundles with this component.
     ///
-    /// The caller must pass the same `storages` and `components` to this method that were used
-    /// to create the stored bundles.
-    pub(crate) unsafe fn refresh_required_components(
-        &mut self,
-        storages: &mut Storages,
-        components: &Components,
+    /// Returns the [`RequiredComponentsError::RemovedFromArchetype`] error if bundles cannot be updated
+    /// in which case the registration of any new required component must be refused.
+    pub(crate) fn verify_to_refresh_required_components(
+        &self,
         requiree: ComponentId,
-    ) -> Result<(), RequiredComponentsError> {
+    ) -> Result<bool, RequiredComponentsError> {
         let Some(bundles_with_requiree) = self
             .components_in_bundles
-            .get_mut(requiree.index())
+            .get(requiree.index())
             .filter(|bundles| !bundles.is_empty())
         else {
             // no bundle with `requiree` exists
-            return Ok(());
+            return Ok(false);
         };
 
         // `EntityWorldMut::remove_with_requires` generate edges between archetypes where the required
@@ -660,8 +659,29 @@ impl Bundles {
             return Err(RequiredComponentsError::RemovedFromArchetype(requiree));
         }
 
-        // take it to update `Self::bundles_with_requiree` while iterating this vector
-        let taken_bundles_with_requiree = core::mem::take(bundles_with_requiree);
+        Ok(true)
+    }
+
+    /// Updates the required components of bundles that contain `requiree`.
+    ///
+    /// # Safety
+    ///
+    /// The caller must have confirmed [`Self::verify_to_refresh_required_components`] returned `Ok(true)`
+    /// for this `requiree` and must pass the same `storages` and `components` to this method that were used
+    /// to create the stored bundles.
+    pub(crate) unsafe fn refresh_required_components(
+        &mut self,
+        storages: &mut Storages,
+        components: &Components,
+        requiree: ComponentId,
+    ) {
+        // take list of bundles to update `Self::bundles_with_requiree` while iterating this
+        let taken_bundles_with_requiree = self
+            .components_in_bundles
+            .get_mut(requiree.index())
+            .filter(|bundles| !bundles.is_empty())
+            .map(core::mem::take)
+            .expect("verify_to_refresh_required_components confirmed existing bundles to refresh");
 
         for bundle_id in taken_bundles_with_requiree.iter() {
             let bundle_info = self.bundle_infos.get_mut(bundle_id.index());
@@ -686,8 +706,6 @@ impl Bundles {
         let replaced = unsafe { replaced.debug_checked_unwrap() };
 
         *replaced = taken_bundles_with_requiree;
-
-        Ok(())
     }
 }
 
diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs
index cc3e21485f..94acafe29b 100644
--- a/crates/bevy_ecs/src/world/mod.rs
+++ b/crates/bevy_ecs/src/world/mod.rs
@@ -542,17 +542,28 @@ impl World {
 
         let required = self.register_component::<R>();
 
+        let update_bundles = self
+            .bundles
+            .verify_to_refresh_required_components(requiree)?;
+
         // SAFETY: We just created the `required` and `requiree` components.
         unsafe {
             self.components
                 .register_required_components::<R>(requiree, required, constructor)?;
         }
 
-        // SAFETY: all bundles are created with Self::storages and Self::components
-        unsafe {
-            self.bundles
-                .refresh_required_components(&mut self.storages, &self.components, requiree)
+        if update_bundles {
+            // SAFETY: all bundles are created with Self::storages and Self::components
+            unsafe {
+                self.bundles.refresh_required_components(
+                    &mut self.storages,
+                    &self.components,
+                    requiree,
+                );
+            }
         }
+
+        Ok(())
     }
 
     /// Retrieves the [required components](RequiredComponents) for the given component type, if it exists.