From b30ba78e16ce37b6c15f4e2b75083087570407c4 Mon Sep 17 00:00:00 2001
From: CatThingy <andyc.andychen@gmail.com>
Date: Sun, 5 Feb 2023 15:37:07 +0000
Subject: [PATCH] Add a wrapper around `Entity` for `RemovedComponents` (#7503)

# Objective

- Make the internals of `RemovedComponents` clearer


## Solution

- Add a wrapper around `Entity`, used in `RemovedComponents` as `Events<RemovedComponentsEntity>`

---

## Changelog

- `RemovedComponents` now internally uses an `Events<RemovedComponentsEntity>` instead of an `Events<Entity>`
---
 crates/bevy_ecs/src/removal_detection.rs | 35 ++++++++++++++++++------
 crates/bevy_ecs/src/world/mod.rs         |  1 +
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/crates/bevy_ecs/src/removal_detection.rs b/crates/bevy_ecs/src/removal_detection.rs
index e0e0bd7430cde..f9953ff77cf9e 100644
--- a/crates/bevy_ecs/src/removal_detection.rs
+++ b/crates/bevy_ecs/src/removal_detection.rs
@@ -19,14 +19,25 @@ use std::{
     option,
 };
 
-/// Wrapper around a [`ManualEventReader<Entity>`] so that we
+/// Wrapper around [`Entity`] for [`RemovedComponents`].
+/// Internally, `RemovedComponents` uses these as an `Events<RemovedComponentEntity>`.
+#[derive(Debug, Clone)]
+pub struct RemovedComponentEntity(Entity);
+
+impl From<RemovedComponentEntity> for Entity {
+    fn from(value: RemovedComponentEntity) -> Self {
+        value.0
+    }
+}
+
+/// Wrapper around a [`ManualEventReader<RemovedComponentEntity>`] so that we
 /// can differentiate events between components.
 #[derive(Debug)]
 pub struct RemovedComponentReader<T>
 where
     T: Component,
 {
-    reader: ManualEventReader<Entity>,
+    reader: ManualEventReader<RemovedComponentEntity>,
     marker: PhantomData<T>,
 }
 
@@ -40,7 +51,7 @@ impl<T: Component> Default for RemovedComponentReader<T> {
 }
 
 impl<T: Component> Deref for RemovedComponentReader<T> {
-    type Target = ManualEventReader<Entity>;
+    type Target = ManualEventReader<RemovedComponentEntity>;
     fn deref(&self) -> &Self::Target {
         &self.reader
     }
@@ -52,11 +63,11 @@ impl<T: Component> DerefMut for RemovedComponentReader<T> {
     }
 }
 
-/// Wrapper around a map of components to [`Events<Entity>`].
+/// Wrapper around a map of components to [`Events<RemovedComponentEntity>`].
 /// So that we can find the events without naming the type directly.
 #[derive(Default, Debug)]
 pub struct RemovedComponentEvents {
-    event_sets: SparseSet<ComponentId, Events<Entity>>,
+    event_sets: SparseSet<ComponentId, Events<RemovedComponentEntity>>,
 }
 
 impl RemovedComponentEvents {
@@ -70,14 +81,17 @@ impl RemovedComponentEvents {
         }
     }
 
-    pub fn get(&self, component_id: impl Into<ComponentId>) -> Option<&Events<Entity>> {
+    pub fn get(
+        &self,
+        component_id: impl Into<ComponentId>,
+    ) -> Option<&Events<RemovedComponentEntity>> {
         self.event_sets.get(component_id.into())
     }
 
     pub fn send(&mut self, component_id: impl Into<ComponentId>, entity: Entity) {
         self.event_sets
             .get_or_insert_with(component_id.into(), Default::default)
-            .send(entity);
+            .send(RemovedComponentEntity(entity));
     }
 }
 
@@ -122,8 +136,10 @@ pub struct RemovedComponents<'w, 's, T: Component> {
 /// Iterator over entities that had a specific component removed.
 ///
 /// See [`RemovedComponents`].
-pub type RemovedIter<'a> =
-    iter::Flatten<option::IntoIter<iter::Cloned<ManualEventIterator<'a, Entity>>>>;
+pub type RemovedIter<'a> = iter::Map<
+    iter::Flatten<option::IntoIter<iter::Cloned<ManualEventIterator<'a, RemovedComponentEntity>>>>,
+    fn(RemovedComponentEntity) -> Entity,
+>;
 
 impl<'w, 's, T: Component> RemovedComponents<'w, 's, T> {
     pub fn iter(&mut self) -> RemovedIter<'_> {
@@ -132,6 +148,7 @@ impl<'w, 's, T: Component> RemovedComponents<'w, 's, T> {
             .map(|events| self.reader.iter(events).cloned())
             .into_iter()
             .flatten()
+            .map(RemovedComponentEntity::into)
     }
 }
 
diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs
index 981accba96f87..ced43f4002094 100644
--- a/crates/bevy_ecs/src/world/mod.rs
+++ b/crates/bevy_ecs/src/world/mod.rs
@@ -788,6 +788,7 @@ impl World {
             .map(|removed| removed.iter_current_update_events().cloned())
             .into_iter()
             .flatten()
+            .map(|e| e.into())
     }
 
     /// Initializes a new resource and returns the [`ComponentId`] created for it.