Remove unnecessary type for iterator

bevyengine · Jan 5, 2023 · df404c2 · df404c2
1 parent 09d1814
commit df404c2
Showing 1 changed file with 9 additions and 21 deletions.
diff --git a/crates/bevy_ecs/src/removal_detection.rs b/crates/bevy_ecs/src/removal_detection.rs
@@ -12,7 +12,7 @@ use bevy_ecs_macros::SystemParam;
 
 use std::{
     fmt::Debug,
-    iter::Cloned,
+    iter::{Cloned, Flatten},
     marker::PhantomData,
     ops::{Deref, DerefMut},
 };
@@ -117,27 +117,15 @@ pub struct RemovedComponents<'w, 's, T: Component> {
     event_sets: &'w RemovedComponentEvents,
 }
 
-pub enum MaybeIter<I: Iterator> {
-    Iter(I),
-    Empty,
-}
+type RemovedIter<'a> = Flatten<std::option::IntoIter<Cloned<ManualEventIterator<'a, Entity>>>>;
 
-impl<I: Iterator> Iterator for MaybeIter<I> {
-    type Item = I::Item;
-    fn next(&mut self) -> Option<Self::Item> {
-        match self {
-            MaybeIter::Iter(iter) => iter.next(),
-            MaybeIter::Empty => None,
-        }
-    }
-}
 impl<'w, 's, T: Component> RemovedComponents<'w, 's, T> {
-    pub fn iter(&mut self) -> MaybeIter<Cloned<ManualEventIterator<'_, Entity>>> {
-        if let Some(events) = self.event_sets.get(**self.component_id) {
-            MaybeIter::Iter(self.reader.iter(events).cloned())
-        } else {
-            MaybeIter::Empty
-        }
+    pub fn iter(&mut self) -> RemovedIter<'_> {
+        self.event_sets
+            .get(**self.component_id)
+            .map(|events| self.reader.iter(events).cloned())
+            .into_iter()
+            .flatten()
     }
 }
 
@@ -146,7 +134,7 @@ where
     T: Component,
 {
     type Item = Entity;
-    type IntoIter = MaybeIter<Cloned<ManualEventIterator<'a, Entity>>>;
+    type IntoIter = RemovedIter<'a>;
     fn into_iter(self) -> Self::IntoIter {
         self.iter()
     }