Fix modifying listener from inside a callback (#18)

* Add test for modifying listener

* Add test and fix for updating listeners from a listener

* Bump version

CHANGELOG.md

@@ -1,3 +1,8 @@
+# 0.6.2
+
+- Fixed: `On<E>` event listeners that mutate themselves inside a callback were being overwritten
+  during cleanup.
+
 # 0.6.1
 
 - Fixed: event listeners adding extra delay when processing events.

Cargo.toml

@@ -6,7 +6,7 @@ resolver = "2"
 
 [package]
 name = "bevy_eventlistener"
-version = "0.6.1"
+version = "0.6.2"
 edition = "2021"
 description = "Event listeners and callbacks for bevy"
 license = "MIT OR Apache-2.0"

crates/bevy_eventlistener_core/src/callbacks.rs

@@ -24,6 +24,10 @@ impl CallbackSystem {
         system.apply_deferred(world);
         *self = CallbackSystem::Initialized(system);
     }
+
+    pub(crate) fn is_empty(&self) -> bool {
+        matches!(self, CallbackSystem::Empty)
+    }
 }
 
 /// A [`SystemParam`](bevy_ecs::system::SystemParam) used to get immutable access the the

crates/bevy_eventlistener_core/src/event_dispatcher.rs

@@ -75,7 +75,10 @@ impl<E: EntityEvent> EventDispatcher<E> {
     pub fn cleanup(mut listeners: Query<&mut On<E>>, mut callbacks: ResMut<EventDispatcher<E>>) {
         for (entity, (callback, _)) in callbacks.listener_graph.drain() {
             if let Ok(mut listener) = listeners.get_mut(entity) {
-                listener.callback = callback;
+                // Do not restore the callback if it has been replaced by the event handler.
+                if listener.callback.is_empty() {
+                    listener.callback = callback;
+                }
             }
         }
     }

src/lib.rs

@@ -92,3 +92,75 @@ pub mod prelude {
     };
     pub use bevy_eventlistener_derive::EntityEvent;
 }
+
+#[test]
+fn replace_listener() {
+    use crate::prelude::*;
+    use bevy::prelude::*;
+
+    #[derive(Clone, Event, EntityEvent)]
+    struct Foo {
+        #[target]
+        target: Entity,
+    }
+
+    let (tx, rx) = std::sync::mpsc::channel();
+    let mut app = App::new();
+    let entity = app.world.spawn_empty().id();
+    app.add_plugins(MinimalPlugins)
+        .add_plugins(EventListenerPlugin::<Foo>::default())
+        .add_systems(Update, move |mut event: EventWriter<Foo>| {
+            event.send(Foo { target: entity })
+        })
+        .update();
+
+    let sender = tx.clone();
+    let callback = On::<Foo>::run(move || sender.send("one").unwrap());
+    app.world.entity_mut(entity).insert(callback);
+    app.update();
+
+    let sender = tx.clone();
+    let callback = On::<Foo>::run(move || sender.send("two").unwrap());
+    app.world.entity_mut(entity).insert(callback);
+    app.update();
+
+    assert_eq!(rx.recv(), Ok("one"));
+    assert_eq!(rx.recv(), Ok("two"));
+}
+
+#[test]
+fn replace_listener_in_callback() {
+    use crate::prelude::*;
+    use bevy::ecs::system::EntityCommands;
+    use bevy::prelude::*;
+
+    #[derive(Clone, Event, EntityEvent)]
+    struct Foo {
+        #[target]
+        target: Entity,
+    }
+
+    let (sender, receiver) = std::sync::mpsc::channel();
+    let mut app = App::new();
+    let entity = app.world.spawn_empty().id();
+    app.add_plugins(MinimalPlugins)
+        .add_plugins(EventListenerPlugin::<Foo>::default())
+        .add_systems(Update, move |mut event: EventWriter<Foo>| {
+            event.send(Foo { target: entity })
+        })
+        .update();
+
+    let callback = On::<Foo>::listener_commands_mut(
+        move |_: &ListenerInput<Foo>, commands: &mut EntityCommands| {
+            sender.send("one").unwrap();
+            let sender2 = sender.clone();
+            commands.insert(On::<Foo>::run(move || sender2.send("two").unwrap()));
+        },
+    );
+    app.world.entity_mut(entity).insert(callback);
+    app.update();
+    app.update();
+
+    assert_eq!(receiver.recv(), Ok("one"));
+    assert_eq!(receiver.recv(), Ok("two"));
+}