From c440de06f1387dfaee43e15a559b312a27b44283 Mon Sep 17 00:00:00 2001
From: Joseph <21144246+JoJoJet@users.noreply.github.com>
Date: Mon, 28 Aug 2023 11:55:59 -0700
Subject: [PATCH] Add a variant of `Events::update` that returns the removed
 events (#9542)

# Objective

Every frame, `Events::update` gets called, which clears out any old
events from the buffer. There should be a way of taking ownership of
these old events instead of throwing them away. My use-case is dumping
old events into a debug menu so they can be inspected later.

One potential workaround is to just have a system that clones any
incoming events and stores them in a list -- however, this requires the
events to implement `Clone`.

## Solution

Add `Events::update_drain`, which returns an iterator of the events that
were removed from the buffer.
---
 crates/bevy_ecs/src/event.rs | 53 +++++++++++++++++++++++++++++-------
 1 file changed, 43 insertions(+), 10 deletions(-)

diff --git a/crates/bevy_ecs/src/event.rs b/crates/bevy_ecs/src/event.rs
index 9693bc18fa..25ae9ebb72 100644
--- a/crates/bevy_ecs/src/event.rs
+++ b/crates/bevy_ecs/src/event.rs
@@ -228,14 +228,27 @@ impl<E: Event> Events<E> {
 
     /// Swaps the event buffers and clears the oldest event buffer. In general, this should be
     /// called once per frame/update.
+    ///
+    /// If you need access to the events that were removed, consider using [`Events::update_drain`].
     pub fn update(&mut self) {
+        let _ = self.update_drain();
+    }
+
+    /// Swaps the event buffers and drains the oldest event buffer, returning an iterator
+    /// of all events that were removed. In general, this should be called once per frame/update.
+    ///
+    /// If you do not need to take ownership of the removed events, use [`Events::update`] instead.
+    #[must_use = "If you do not need the returned events, call .update() instead."]
+    pub fn update_drain(&mut self) -> impl Iterator<Item = E> + '_ {
         std::mem::swap(&mut self.events_a, &mut self.events_b);
-        self.events_b.clear();
+        let iter = self.events_b.events.drain(..);
         self.events_b.start_event_count = self.event_count;
         debug_assert_eq!(
             self.events_a.start_event_count + self.events_a.len(),
             self.events_b.start_event_count
         );
+
+        iter.map(|e| e.event)
     }
 
     /// A system that calls [`Events::update`] once per frame.
@@ -725,7 +738,7 @@ impl<'a, E: Event> ExactSizeIterator for EventIteratorWithId<'a, E> {
 
 #[cfg(test)]
 mod tests {
-    use crate::{prelude::World, system::SystemState};
+    use crate::system::assert_is_read_only_system;
 
     use super::*;
 
@@ -982,6 +995,31 @@ mod tests {
         assert!(is_empty, "EventReader should be empty");
     }
 
+    #[test]
+    fn test_update_drain() {
+        let mut events = Events::<TestEvent>::default();
+        let mut reader = events.get_reader();
+
+        events.send(TestEvent { i: 0 });
+        events.send(TestEvent { i: 1 });
+        assert_eq!(reader.iter(&events).count(), 2);
+
+        let mut old_events = Vec::from_iter(events.update_drain());
+        assert!(old_events.is_empty());
+
+        events.send(TestEvent { i: 2 });
+        assert_eq!(reader.iter(&events).count(), 1);
+
+        old_events.extend(events.update_drain());
+        assert_eq!(old_events.len(), 2);
+
+        old_events.extend(events.update_drain());
+        assert_eq!(
+            old_events,
+            &[TestEvent { i: 0 }, TestEvent { i: 1 }, TestEvent { i: 2 }]
+        );
+    }
+
     #[allow(clippy::iter_nth_zero)]
     #[test]
     fn test_event_iter_nth() {
@@ -1053,13 +1091,8 @@ mod tests {
 
     #[test]
     fn ensure_reader_readonly() {
-        fn read_for<E: Event>() {
-            let mut world = World::new();
-            world.init_resource::<Events<E>>();
-            let mut state = SystemState::<EventReader<E>>::new(&mut world);
-            // This can only work if EventReader only reads the world
-            let _reader = state.get(&world);
-        }
-        read_for::<EmptyTestEvent>();
+        fn reader_system(_: EventReader<EmptyTestEvent>) {}
+
+        assert_is_read_only_system(reader_system);
     }
 }