From 23576728810dd9543d669afac52bfbf6cecde3bc Mon Sep 17 00:00:00 2001
From: Peepo-Juice <102474367+Peepo-Juice@users.noreply.github.com>
Date: Fri, 3 May 2024 16:48:07 -0400
Subject: [PATCH] Run a one-shot system upon completion of Tween or Delay
 (#111)

Co-authored-by: Peepo-Juice <houssein.dev@gmail.com>
---
 src/plugin.rs    |   4 +
 src/tweenable.rs | 214 +++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 200 insertions(+), 18 deletions(-)

diff --git a/src/plugin.rs b/src/plugin.rs
index 46afce3..e03e0a7 100644
--- a/src/plugin.rs
+++ b/src/plugin.rs
@@ -88,6 +88,7 @@ pub fn component_animator_system<T: Component>(
     time: Res<Time>,
     mut query: Query<(Entity, &mut T, &mut Animator<T>)>,
     events: ResMut<Events<TweenCompleted>>,
+    mut commands: Commands,
 ) {
     let mut events: Mut<Events<TweenCompleted>> = events.into();
     for (entity, target, mut animator) in query.iter_mut() {
@@ -99,6 +100,7 @@ pub fn component_animator_system<T: Component>(
                 &mut target,
                 entity,
                 &mut events,
+                &mut commands,
             );
         }
     }
@@ -116,6 +118,7 @@ pub fn asset_animator_system<T: Asset>(
     mut assets: ResMut<Assets<T>>,
     mut query: Query<(Entity, &Handle<T>, &mut AssetAnimator<T>)>,
     events: ResMut<Events<TweenCompleted>>,
+    mut commands: Commands,
 ) {
     let mut events: Mut<Events<TweenCompleted>> = events.into();
     let mut target = AssetTarget::new(assets.reborrow());
@@ -131,6 +134,7 @@ pub fn asset_animator_system<T: Asset>(
                 &mut target,
                 entity,
                 &mut events,
+                &mut commands,
             );
         }
     }
diff --git a/src/tweenable.rs b/src/tweenable.rs
index b497fd3..de83b3d 100644
--- a/src/tweenable.rs
+++ b/src/tweenable.rs
@@ -3,7 +3,7 @@ use std::{
     time::Duration,
 };
 
-use bevy::prelude::*;
+use bevy::{ecs::system::SystemId, prelude::*};
 
 use crate::{EaseMethod, Lens, RepeatCount, RepeatStrategy, TweeningDirection};
 
@@ -25,6 +25,7 @@ use crate::{EaseMethod, Lens, RepeatCount, RepeatStrategy, TweeningDirection};
 /// implement [`From`]:
 /// ```no_run
 /// # use std::time::Duration;
+/// # use bevy::ecs::system::Commands;
 /// # use bevy::prelude::{Entity, Events, Mut, Transform};
 /// # use bevy_tweening::{BoxedTweenable, Sequence, Tweenable, TweenCompleted, TweenState, Targetable, TotalDuration};
 /// #
@@ -34,7 +35,7 @@ use crate::{EaseMethod, Lens, RepeatCount, RepeatStrategy, TweeningDirection};
 /// #     fn total_duration(&self) -> TotalDuration  { unimplemented!() }
 /// #     fn set_elapsed(&mut self, elapsed: Duration)  { unimplemented!() }
 /// #     fn elapsed(&self) -> Duration  { unimplemented!() }
-/// #     fn tick<'a>(&mut self, delta: Duration, target: &'a mut dyn Targetable<Transform>, entity: Entity, events: &mut Mut<Events<TweenCompleted>>) -> TweenState  { unimplemented!() }
+/// #     fn tick<'a>(&mut self, delta: Duration, target: &'a mut dyn Targetable<Transform>, entity: Entity, events: &mut Mut<Events<TweenCompleted>>, commands: &mut Commands) -> TweenState  { unimplemented!() }
 /// #     fn rewind(&mut self) { unimplemented!() }
 /// # }
 ///
@@ -354,6 +355,7 @@ pub trait Tweenable<T>: Send + Sync {
         target: &mut dyn Targetable<T>,
         entity: Entity,
         events: &mut Mut<Events<TweenCompleted>>,
+        commands: &mut Commands,
     ) -> TweenState;
 
     /// Rewind the animation to its starting state.
@@ -434,6 +436,7 @@ pub struct Tween<T> {
     lens: Box<dyn Lens<T> + Send + Sync + 'static>,
     on_completed: Option<Box<CompletedCallback<Tween<T>>>>,
     event_data: Option<u64>,
+    system_id: Option<SystemId>,
 }
 
 impl<T: 'static> Tween<T> {
@@ -498,6 +501,7 @@ impl<T> Tween<T> {
             lens: Box::new(lens),
             on_completed: None,
             event_data: None,
+            system_id: None,
         }
     }
 
@@ -577,6 +581,45 @@ impl<T> Tween<T> {
         self
     }
 
+    /// Enable running a one-shot system upon completion.
+    ///
+    /// If enabled, the tween will run a system via a provided [`SystemId`] when the
+    /// animation completes. This is similar to the [`with_completed()`],
+    /// but uses a system registered by [`register_system()`] instead of a callback.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use bevy_tweening::{lens::*, *};
+    /// # use bevy::{ecs::event::EventReader, math::Vec3, ecs::world::World, ecs::system::Query, ecs::entity::Entity, ecs::query::With};
+    /// # use std::time::Duration;
+    /// let mut world = World::new();  
+    /// let test_system_system_id = world.register_system(test_system);
+    /// let tween = Tween::new(
+    ///     // [...]
+    /// #    EaseFunction::QuadraticInOut,
+    /// #    Duration::from_secs(1),
+    /// #    TransformPositionLens {
+    /// #        start: Vec3::ZERO,
+    /// #        end: Vec3::new(3.5, 0., 0.),
+    /// #    },
+    /// )
+    /// .with_completed_system(test_system_system_id);
+    ///
+    /// fn test_system(query: Query<Entity>) {
+    ///    for entity in query.iter() {
+    ///       println!("Found an entity!");
+    ///    }
+    /// }
+    /// ```
+    /// [`with_completed()`]: Tween::with_completed
+    /// [`register_system()`]: bevy::ecs::world::World::register_system
+    #[must_use]
+    pub fn with_completed_system(mut self, system_id: SystemId) -> Self {
+        self.system_id = Some(system_id);
+        self
+    }
+
     /// Set the playback direction of the tween.
     ///
     /// The playback direction influences the mapping of the progress ratio (in
@@ -670,6 +713,27 @@ impl<T> Tween<T> {
     pub fn clear_completed_event(&mut self) {
         self.event_data = None;
     }
+
+    /// Enable running a one-shot system upon completion.
+    ///
+    /// If enabled, the tween will run a system via a provided [`SystemId`] when the
+    /// animation completes. This is similar to the [`with_completed()`],
+    /// but uses a system registered by [`register_system()`] instead of a callback.
+    ///
+    /// [`with_completed()`]: Tween::with_completed  
+    /// [`register_system()`]: bevy::ecs::world::World::register_system
+    pub fn set_completed_system(&mut self, user_data: SystemId) {
+        self.system_id = Some(user_data);
+    }
+
+    /// Clear the system that will execute when the animation completes.
+    ///
+    /// See also [`set_completed_system()`].
+    ///
+    /// [`set_completed_system()`]: Tween::set_completed_system
+    pub fn clear_completed_system(&mut self) {
+        self.system_id = None;
+    }
 }
 
 impl<T> Tweenable<T> for Tween<T> {
@@ -695,6 +759,7 @@ impl<T> Tweenable<T> for Tween<T> {
         target: &mut dyn Targetable<T>,
         entity: Entity,
         events: &mut Mut<Events<TweenCompleted>>,
+        commands: &mut Commands,
     ) -> TweenState {
         if self.clock.state() == TweenState::Completed {
             return TweenState::Completed;
@@ -732,6 +797,9 @@ impl<T> Tweenable<T> for Tween<T> {
             if let Some(cb) = &self.on_completed {
                 cb(entity, self);
             }
+            if let Some(system_id) = &self.system_id {
+                commands.run_system(*system_id);
+            }
         }
 
         state
@@ -873,12 +941,13 @@ impl<T> Tweenable<T> for Sequence<T> {
         target: &mut dyn Targetable<T>,
         entity: Entity,
         events: &mut Mut<Events<TweenCompleted>>,
+        commands: &mut Commands,
     ) -> TweenState {
         self.elapsed = self.elapsed.saturating_add(delta).min(self.duration);
         while self.index < self.tweens.len() {
             let tween = &mut self.tweens[self.index];
             let tween_remaining = tween.duration() - tween.elapsed();
-            if let TweenState::Active = tween.tick(delta, target, entity, events) {
+            if let TweenState::Active = tween.tick(delta, target, entity, events, commands) {
                 return TweenState::Active;
             }
 
@@ -954,11 +1023,12 @@ impl<T> Tweenable<T> for Tracks<T> {
         target: &mut dyn Targetable<T>,
         entity: Entity,
         events: &mut Mut<Events<TweenCompleted>>,
+        commands: &mut Commands,
     ) -> TweenState {
         self.elapsed = self.elapsed.saturating_add(delta).min(self.duration);
         let mut any_active = false;
         for tweenable in &mut self.tracks {
-            let state = tweenable.tick(delta, target, entity, events);
+            let state = tweenable.tick(delta, target, entity, events, commands);
             any_active = any_active || (state == TweenState::Active);
         }
         if any_active {
@@ -986,6 +1056,7 @@ pub struct Delay<T> {
     timer: Timer,
     on_completed: Option<Box<CompletedCallback<Delay<T>>>>,
     event_data: Option<u64>,
+    system_id: Option<SystemId>,
 }
 
 impl<T: 'static> Delay<T> {
@@ -1010,6 +1081,7 @@ impl<T> Delay<T> {
             timer: Timer::new(duration, TimerMode::Once),
             on_completed: None,
             event_data: None,
+            system_id: None,
         }
     }
 
@@ -1082,6 +1154,46 @@ impl<T> Delay<T> {
         self
     }
 
+    /// Enable running a one-shot system upon completion.
+    ///
+    /// If enabled, the tween will run a system via a provided [`SystemId`] when the
+    /// animation completes. This is similar to the [`with_completed()`],
+    /// but uses a system registered by [`register_system()`] instead.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use bevy_tweening::{lens::*, *};
+    /// # use bevy::{ecs::event::EventReader, math::Vec3, ecs::world::World, ecs::system::Query, ecs::entity::Entity};
+    /// # use std::time::Duration;
+    /// let mut world = World::new();
+    /// let test_system_system_id = world.register_system(test_system);
+    /// let tween = Tween::new(
+    ///     // [...]
+    /// #    EaseFunction::QuadraticInOut,
+    /// #    Duration::from_secs(1),
+    /// #    TransformPositionLens {
+    /// #        start: Vec3::ZERO,
+    /// #        end: Vec3::new(3.5, 0., 0.),
+    /// #    },
+    /// )
+    /// .with_completed_system(test_system_system_id);
+    ///
+    /// fn test_system(query: Query<Entity>) {
+    ///    for entity in query.iter() {
+    ///       println!("Found an Entity!");
+    ///    }
+    /// }
+    /// ```
+    ///
+    /// [`with_completed()`]: Tween::with_completed
+    /// [`register_system()`]: bevy::ecs::world::World::register_system
+    #[must_use]
+    pub fn with_completed_system(mut self, system_id: SystemId) -> Self {
+        self.system_id = Some(system_id);
+        self
+    }
+
     /// Check if the delay completed.
     pub fn is_completed(&self) -> bool {
         self.timer.finished()
@@ -1141,6 +1253,29 @@ impl<T> Delay<T> {
     pub fn clear_completed_event(&mut self) {
         self.event_data = None;
     }
+
+    /// Enable running a one-shot system upon completion.
+    ///
+    /// If enabled, the tween will run a system via a provided [`SystemId`] when the
+    /// animation completes. This is similar to the [`with_completed()`],
+    /// but uses a system registered by [`register_system()`] instead of a callback.
+    ///
+    /// See [`with_completed_system()`] for details.
+    ///
+    /// [`with_completed()`]: Tween::with_completed
+    /// [`register_system()`]: bevy::ecs::world::World::register_system
+    pub fn set_completed_system(&mut self, system_id: SystemId) {
+        self.system_id = Some(system_id);
+    }
+
+    /// Clear the system that will execute when the animation completes.
+    ///
+    /// See also [`set_completed_system()`].
+    ///
+    /// [`clear_completed_system()`]: Tween::clear_completed_system
+    pub fn clear_completed_system(&mut self) {
+        self.system_id = None;
+    }
 }
 
 impl<T> Tweenable<T> for Delay<T> {
@@ -1170,6 +1305,7 @@ impl<T> Tweenable<T> for Delay<T> {
         _target: &mut dyn Targetable<T>,
         entity: Entity,
         events: &mut Mut<Events<TweenCompleted>>,
+        commands: &mut Commands,
     ) -> TweenState {
         let was_completed = self.is_completed();
 
@@ -1188,6 +1324,9 @@ impl<T> Tweenable<T> for Delay<T> {
             if let Some(cb) = &self.on_completed {
                 cb(entity, self);
             }
+            if let Some(system_id) = &self.system_id {
+                commands.run_system(*system_id);
+            }
         }
 
         state
@@ -1202,7 +1341,11 @@ impl<T> Tweenable<T> for Delay<T> {
 mod tests {
     use std::sync::{Arc, Mutex};
 
-    use bevy::ecs::{component::Tick, system::SystemState};
+    use bevy::ecs::{
+        component::Tick,
+        event::Events,
+        system::{CommandQueue, SystemState},
+    };
 
     use super::*;
     use crate::{lens::*, test_utils::*};
@@ -1226,13 +1369,17 @@ mod tests {
     }
 
     /// Utility to create a test environment to tick a tween.
-    fn make_test_env() -> (World, Entity) {
+    fn make_test_env() -> (World, Entity, SystemId) {
         let mut world = World::new();
         world.init_resource::<Events<TweenCompleted>>();
         let entity = world.spawn(Transform::default()).id();
-        (world, entity)
+        let system_id = world.register_system(oneshot_test);
+        (world, entity, system_id)
     }
 
+    /// one-shot system to be used for testing
+    fn oneshot_test() {}
+
     /// Manually tick a test tweenable targeting a component.
     fn manual_tick_component<T: Component>(
         duration: Duration,
@@ -1244,7 +1391,16 @@ mod tests {
             |world: &mut World, mut events: Mut<Events<TweenCompleted>>| {
                 let transform = world.get_mut::<T>(entity).unwrap();
                 let mut target = ComponentTarget::new(transform);
-                tween.tick(duration, &mut target, entity, &mut events)
+                // let command_queue = &mut CommandQueue::default(); // todo
+                // let mut asd = Commands::new(command_queue, world);
+                tween.tick(
+                    duration,
+                    &mut target,
+                    entity,
+                    &mut events,
+                    // passing dummy values to let things compile
+                    &mut Commands::new(&mut CommandQueue::default(), &World::default()),
+                )
             },
         )
     }
@@ -1369,7 +1525,7 @@ mod tests {
                 assert!(tween.on_completed.is_none());
                 assert!(tween.event_data.is_none());
 
-                let (mut world, entity) = make_test_env();
+                let (mut world, entity, system_id) = make_test_env();
                 let mut event_reader_system_state: SystemState<EventReader<TweenCompleted>> =
                     SystemState::new(&mut world);
 
@@ -1393,6 +1549,11 @@ mod tests {
                 assert!(tween.event_data.is_some());
                 assert_eq!(tween.event_data.unwrap(), USER_DATA);
 
+                // Activate oneshot system
+                tween.set_completed_system(system_id);
+                assert!(tween.system_id.is_some());
+                assert_eq!(tween.system_id.unwrap(), system_id);
+
                 // Loop over 2.2 seconds, so greater than one ping-pong loop
                 let tick_duration = Duration::from_millis(200);
                 for i in 1..=11 {
@@ -1575,6 +1736,10 @@ mod tests {
                 // Clear event sending
                 tween.clear_completed_event();
                 assert!(tween.event_data.is_none());
+
+                // Clear oneshot system
+                tween.clear_completed_system();
+                assert!(tween.system_id.is_none());
             }
         }
     }
@@ -1606,7 +1771,7 @@ mod tests {
         // progress is independent of direction
         assert_approx_eq!(tween.progress(), 0.3);
 
-        let (mut world, entity) = make_test_env();
+        let (mut world, entity, _system_id) = make_test_env();
 
         // Progress always increases alongside the current direction
         tween.set_direction(TweeningDirection::Backward);
@@ -1661,7 +1826,7 @@ mod tests {
         );
         let mut seq = tween1.then(tween2);
 
-        let (mut world, entity) = make_test_env();
+        let (mut world, entity, _system_id) = make_test_env();
 
         for i in 1..=16 {
             let state =
@@ -1703,7 +1868,7 @@ mod tests {
             .with_repeat_count(RepeatCount::Finite(1))
         }));
 
-        let (mut world, entity) = make_test_env();
+        let (mut world, entity, _system_id) = make_test_env();
 
         // Tick halfway through the first tween, then in one tick:
         // - Finish the first tween
@@ -1815,7 +1980,7 @@ mod tests {
         let mut tracks = Tracks::new([tween1, tween2]);
         assert_eq!(tracks.duration(), Duration::from_secs(1)); // max(1., 0.8)
 
-        let (mut world, entity) = make_test_env();
+        let (mut world, entity, _system_id) = make_test_env();
 
         for i in 1..=6 {
             let state =
@@ -1884,13 +2049,20 @@ mod tests {
     #[test]
     fn delay_tick() {
         let duration = Duration::from_secs(1);
+        // Dummy world and registered oneshot system ID
+        let (mut world, entity, system_id) = make_test_env();
 
         const USER_DATA: u64 = 42;
-        let mut delay = Delay::new(duration).with_completed_event(USER_DATA);
+        let mut delay = Delay::new(duration)
+            .with_completed_event(USER_DATA)
+            .with_completed_system(system_id);
 
         assert!(delay.event_data.is_some());
         assert_eq!(delay.event_data.unwrap(), USER_DATA);
 
+        assert!(delay.system_id.is_some());
+        assert_eq!(delay.system_id.unwrap(), system_id);
+
         delay.clear_completed_event();
         assert!(delay.event_data.is_none());
 
@@ -1898,6 +2070,13 @@ mod tests {
         assert!(delay.event_data.is_some());
         assert_eq!(delay.event_data.unwrap(), USER_DATA);
 
+        delay.clear_completed_system();
+        assert!(delay.system_id.is_none());
+
+        delay.set_completed_system(system_id);
+        assert!(delay.system_id.is_some());
+        assert_eq!(delay.system_id.unwrap(), system_id);
+
         {
             let tweenable: &dyn Tweenable<Transform> = &delay;
             assert_eq!(tweenable.duration(), duration);
@@ -1905,8 +2084,7 @@ mod tests {
             assert_eq!(tweenable.elapsed(), Duration::ZERO);
         }
 
-        // Dummy world and event writer
-        let (mut world, entity) = make_test_env();
+        // Dummy event writer
         let mut event_reader_system_state: SystemState<EventReader<TweenCompleted>> =
             SystemState::new(&mut world);
 
@@ -2041,7 +2219,7 @@ mod tests {
 
         assert_approx_eq!(tween.progress(), 0.);
 
-        let (mut world, entity) = make_test_env();
+        let (mut world, entity, _system_id) = make_test_env();
 
         // 10%
         let state =
@@ -2088,7 +2266,7 @@ mod tests {
 
         assert_approx_eq!(tween.progress(), 0.);
 
-        let (mut world, entity) = make_test_env();
+        let (mut world, entity, _system_id) = make_test_env();
 
         // 10%
         let state =