feature: add current and next state matches to on_enter and on_exit

src/machine.rs

@@ -159,8 +159,8 @@ pub struct StateMachine {
     /// each StateMetadata would mean that e.g. we'd have to check every AnyState trigger before any
     /// state-specific trigger or vice versa.
     transitions: Vec<(fn(TypeId) -> bool, Box<dyn Transition>)>,
-    on_exit: Vec<(fn(TypeId) -> bool, OnEvent)>,
-    on_enter: Vec<(fn(TypeId) -> bool, OnEvent)>,
+    on_exit: Vec<(fn(TypeId) -> bool, fn(TypeId) -> bool, OnEvent)>,
+    on_enter: Vec<(fn(TypeId) -> bool, fn(TypeId) -> bool, OnEvent)>,
     /// Transitions must be initialized whenever a transition is added or a transition occurs
     init_transitions: bool,
     /// If true, all transitions are logged at info level
@@ -236,51 +236,96 @@ impl StateMachine {
         self
     }
 
-    /// Adds an on-enter event to the state machine. Whenever the state machine transitions into the
-    /// given state, it will run the event.
-    pub fn on_enter<S: EntityState>(
+    /// Adds an on-enter event to the state machine. Whenever the state machine transitions
+    /// from the given previous state to the given next state, it will run the event.
+    pub fn on_enter_to<Prev: EntityState, Next: EntityState>(
         mut self,
         on_enter: impl 'static + Fn(&mut EntityCommands) + Send + Sync,
     ) -> Self {
-        self.on_enter
-            .push((S::matches, OnEvent::Entity(Box::new(on_enter))));
+        self.on_enter.push((
+            Prev::matches,
+            Next::matches,
+            OnEvent::Entity(Box::new(on_enter)),
+        ));
 
         self
     }
 
-    /// Adds an on-exit event to the state machine. Whenever the state machine transitions from the
-    /// given state, it will run the event.
-    pub fn on_exit<S: EntityState>(
+    /// Adds an on-enter event to the state machine. Whenever the state machine transitions
+    /// from any previous state to the given next state, it will run the event.
+    pub fn on_enter<Next: EntityState>(
+        self,
+        on_enter: impl 'static + Fn(&mut EntityCommands) + Send + Sync,
+    ) -> Self {
+        self.on_enter_to::<AnyState, Next>(on_enter)
+    }
+
+    /// Adds an on-enter event to the state machine. Whenever the state machine transitions
+    /// from the given previous state to the given next state, it will run the event.
+    pub fn on_exit_from<Prev: EntityState, Next: EntityState>(
         mut self,
         on_exit: impl 'static + Fn(&mut EntityCommands) + Send + Sync,
     ) -> Self {
-        self.on_exit
-            .push((S::matches, OnEvent::Entity(Box::new(on_exit))));
+        self.on_exit.push((
+            Prev::matches,
+            Next::matches,
+            OnEvent::Entity(Box::new(on_exit)),
+        ));
 
         self
     }
 
-    /// Adds an on-enter command to the state machine. Whenever the state machine transitions into
-    /// the given state, it will run the command.
-    pub fn command_on_enter<S: EntityState>(
+    /// Adds an on-exit event to the state machine. Whenever the state machine transitions
+    /// from the given previous state to any next state, it will run the event.
+    pub fn on_exit<Prev: EntityState>(
+        self,
+        on_exit: impl 'static + Fn(&mut EntityCommands) + Send + Sync,
+    ) -> Self {
+        self.on_exit_from::<Prev, AnyState>(on_exit)
+    }
+
+    /// Adds an on-enter command to the state machine. Whenever the state machine transitions
+    /// from the given previous state to the given next state, it will run the command.
+    pub fn command_on_enter_to<Prev: EntityState, Next: EntityState>(
         mut self,
         command: impl Clone + Command + Sync,
     ) -> Self {
-        self.on_enter
-            .push((S::matches, OnEvent::Command(Box::new(command))));
+        self.on_enter.push((
+            Prev::matches,
+            Next::matches,
+            OnEvent::Command(Box::new(command)),
+        ));
 
         self
     }
 
-    /// Adds an on-exit command to the state machine. Whenever the state machine transitions from
-    /// the given state, it will run the command.
-    pub fn command_on_exit<S: EntityState>(mut self, command: impl Clone + Command + Sync) -> Self {
-        self.on_exit
-            .push((S::matches, OnEvent::Command(Box::new(command))));
+    /// Adds an on-enter command to the state machine. Whenever the state machine transitions
+    /// from any previous state to the given next state, it will run the command.
+    pub fn command_on_enter<Next: EntityState>(self, command: impl Clone + Command + Sync) -> Self {
+        self.command_on_enter_to::<AnyState, Next>(command)
+    }
+
+    /// Adds an on-exit command to the state machine. Whenever the state machine transitions
+    /// from the given previous state to the given next state, it will run the command.
+    pub fn command_on_exit_from<Prev: EntityState, Next: EntityState>(
+        mut self,
+        command: impl Clone + Command + Sync,
+    ) -> Self {
+        self.on_exit.push((
+            Prev::matches,
+            Next::matches,
+            OnEvent::Command(Box::new(command)),
+        ));
 
         self
     }
 
+    /// Adds an on-exit command to the state machine. Whenever the state machine transitions
+    /// from the given previous state to any next state, it will run the command.
+    pub fn command_on_exit<Prev: EntityState>(self, command: impl Clone + Command + Sync) -> Self {
+        self.command_on_exit_from::<Prev, AnyState>(command)
+    }
+
     /// Sets whether transitions are logged to the console
     pub fn set_trans_logging(mut self, log_transitions: bool) -> Self {
         self.log_transitions = log_transitions;
@@ -331,16 +376,16 @@ impl StateMachine {
         };
         let to = &self.states[&next_state];
 
-        for (matches, event) in &self.on_exit {
-            if matches(current) {
+        for (matches_current, matches_next, event) in &self.on_exit {
+            if matches_current(current) && matches_next(next_state) {
                 event.trigger(entity, &mut world.commands());
             }
         }
 
         trans(world, current);
 
-        for (matches, event) in &self.on_enter {
-            if matches(next_state) {
+        for (matches_current, matches_next, event) in &self.on_enter {
+            if matches_current(current) && matches_next(next_state) {
                 event.trigger(entity, &mut world.commands());
             }
         }
@@ -551,4 +596,117 @@ mod tests {
         assert!(test.on_b, "on_b should be true");
         assert!(test.on_any, "on_any should be true");
     }
+
+    #[test]
+    fn test_event_matches() {
+        #[derive(Component, Default)]
+        struct InB;
+
+        #[derive(Component, Clone)]
+        struct A;
+
+        #[derive(Component, Clone)]
+        struct B1;
+
+        #[derive(Component, Clone)]
+        struct B2;
+
+        #[derive(Component, Clone)]
+        struct C;
+
+        let mut app = App::new();
+        app.add_plugins((
+            MinimalPlugins,
+            bevy::log::LogPlugin::default(),
+            StateMachinePlugin::default(),
+        ));
+        app.update();
+
+        let machine = StateMachine::default()
+            .trans::<A, _>(always, B1)
+            .trans::<B1, _>(always, B2)
+            .trans::<B2, _>(always, C)
+            .on_enter_to::<NotState<OneOfState<(B1, B2)>>, OneOfState<(B1, B2)>>(|ec| {
+                ec.insert(InB);
+            })
+            .on_exit_from::<OneOfState<(B1, B2)>, NotState<OneOfState<(B1, B2)>>>(|ec| {
+                ec.remove::<InB>();
+            })
+            .set_trans_logging(true);
+
+        let id = app.world_mut().spawn((A, machine)).id();
+        app.update();
+        assert!(app.world().get::<InB>(id).is_some());
+
+        app.update();
+        assert!(app.world().get::<InB>(id).is_some());
+
+        app.update();
+        assert!(app.world().get::<InB>(id).is_none());
+    }
+
+    #[test]
+    fn test_command_event_matches() {
+        #[derive(Resource, Default)]
+        struct InBResource;
+
+        #[derive(Clone, Debug)]
+        struct InitInBResourceCommand;
+
+        impl Command for InitInBResourceCommand {
+            fn apply(self, world: &mut World) {
+                world.init_resource::<InBResource>();
+            }
+        }
+
+        #[derive(Clone, Debug)]
+        struct RemoveInBResourceCommand;
+
+        impl Command for RemoveInBResourceCommand {
+            fn apply(self, world: &mut World) {
+                world.remove_resource::<InBResource>();
+            }
+        }
+
+        #[derive(Component, Clone)]
+        struct A;
+
+        #[derive(Component, Clone)]
+        struct B1;
+
+        #[derive(Component, Clone)]
+        struct B2;
+
+        #[derive(Component, Clone)]
+        struct C;
+
+        let mut app = App::new();
+        app.add_plugins((
+            MinimalPlugins,
+            bevy::log::LogPlugin::default(),
+            StateMachinePlugin::default(),
+        ));
+        app.update();
+
+        let machine = StateMachine::default()
+            .trans::<A, _>(always, B1)
+            .trans::<B1, _>(always, B2)
+            .trans::<B2, _>(always, C)
+            .command_on_enter_to::<NotState<OneOfState<(B1, B2)>>, OneOfState<(B1, B2)>>(
+                InitInBResourceCommand,
+            )
+            .command_on_exit_from::<OneOfState<(B1, B2)>, NotState<OneOfState<(B1, B2)>>>(
+                RemoveInBResourceCommand,
+            )
+            .set_trans_logging(true);
+
+        app.world_mut().spawn((A, machine));
+
+        app.update();
+        assert!(app.world().contains_resource::<InBResource>());
+
+        app.update();
+        app.update();
+        assert!(!app.world().contains_resource::<InBResource>());
+    }
 }