From 5ed296ff0323ae349865fc9196d9b4379ce1bc39 Mon Sep 17 00:00:00 2001
From: MiniaczQ <xnetroidpl@gmail.com>
Date: Sat, 15 Jun 2024 00:22:59 +0200
Subject: [PATCH] Restore overwrite capabilities of `insert_state` (#13848)

# Objective

- Fixes #13844
- Warn user when initializing state multiple times

## Solution

- `insert_state` will overwrite previously initialized state value,
reset transition events and re-insert it's own transition event.
- `init_state`, `add_sub_state`, `add_computed_state` are idempotent, so
calling them multiple times will emit a warning.

## Testing

- 2 tests confirming overwrite works.
- Given the example from #13844
```rs
use bevy::prelude::*;

fn main() {
    App::new()
        .add_plugins(DefaultPlugins)
        .insert_state(AppState::A)
        .insert_state(AppState::B)
        .add_systems(OnEnter(AppState::A), setup_a)
        .add_systems(OnEnter(AppState::B), setup_b)
        .add_systems(OnExit(AppState::A), cleanup_a)
        .add_systems(OnExit(AppState::B), cleanup_b)
        .run();
}

#[derive(States, Debug, Clone, PartialEq, Eq, Hash)]
enum AppState {
    A,
    B,
}

fn setup_a() {
    info!("setting up A");
}

fn setup_b() {
    info!("setting up B");
}

fn cleanup_a() {
    info!("cleaning up A");
}

fn cleanup_b() {
    info!("cleaning up B");
}
```

We get the following result:
```
INFO states: setting up B
```
which matches our expectations.
---
 crates/bevy_state/src/app.rs | 79 ++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)

diff --git a/crates/bevy_state/src/app.rs b/crates/bevy_state/src/app.rs
index 955c957edee8d..33a141e8d6be9 100644
--- a/crates/bevy_state/src/app.rs
+++ b/crates/bevy_state/src/app.rs
@@ -4,6 +4,7 @@ use bevy_ecs::{
     schedule::{IntoSystemConfigs, ScheduleLabel},
     world::FromWorld,
 };
+use bevy_utils::tracing::warn;
 
 use crate::state::{
     setup_state_transitions_in_world, ComputedStates, FreelyMutableState, NextState, State,
@@ -70,6 +71,9 @@ impl AppExtStates for SubApp {
                 exited: None,
                 entered: Some(state),
             });
+        } else {
+            let name = std::any::type_name::<S>();
+            warn!("State {} is already initialized.", name);
         }
 
         self
@@ -87,6 +91,16 @@ impl AppExtStates for SubApp {
                 exited: None,
                 entered: Some(state),
             });
+        } else {
+            // Overwrite previous state and initial event
+            self.insert_resource::<State<S>>(State::new(state.clone()));
+            self.world_mut()
+                .resource_mut::<Events<StateTransitionEvent<S>>>()
+                .clear();
+            self.world_mut().send_event(StateTransitionEvent {
+                exited: None,
+                entered: Some(state),
+            });
         }
 
         self
@@ -109,6 +123,9 @@ impl AppExtStates for SubApp {
                 exited: None,
                 entered: state,
             });
+        } else {
+            let name = std::any::type_name::<S>();
+            warn!("Computed state {} is already initialized.", name);
         }
 
         self
@@ -132,6 +149,9 @@ impl AppExtStates for SubApp {
                 exited: None,
                 entered: state,
             });
+        } else {
+            let name = std::any::type_name::<S>();
+            warn!("Sub state {} is already initialized.", name);
         }
 
         self
@@ -192,3 +212,62 @@ impl Plugin for StatesPlugin {
         schedule.insert_after(PreUpdate, StateTransition);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        self as bevy_state,
+        state::{State, StateTransition, StateTransitionEvent},
+    };
+    use bevy_app::App;
+    use bevy_ecs::event::Events;
+    use bevy_state_macros::States;
+
+    use super::AppExtStates;
+
+    #[derive(States, Default, PartialEq, Eq, Hash, Debug, Clone)]
+    enum TestState {
+        #[default]
+        A,
+        B,
+        C,
+    }
+
+    #[test]
+    fn insert_state_can_overwrite_init_state() {
+        let mut app = App::new();
+
+        app.init_state::<TestState>();
+        app.insert_state(TestState::B);
+
+        let world = app.world_mut();
+        world.run_schedule(StateTransition);
+
+        assert_eq!(world.resource::<State<TestState>>().0, TestState::B);
+        let events = world.resource::<Events<StateTransitionEvent<TestState>>>();
+        assert_eq!(events.len(), 1);
+        let mut reader = events.get_reader();
+        let last = reader.read(events).last().unwrap();
+        assert_eq!(last.exited, None);
+        assert_eq!(last.entered, Some(TestState::B));
+    }
+
+    #[test]
+    fn insert_state_can_overwrite_insert_state() {
+        let mut app = App::new();
+
+        app.insert_state(TestState::B);
+        app.insert_state(TestState::C);
+
+        let world = app.world_mut();
+        world.run_schedule(StateTransition);
+
+        assert_eq!(world.resource::<State<TestState>>().0, TestState::C);
+        let events = world.resource::<Events<StateTransitionEvent<TestState>>>();
+        assert_eq!(events.len(), 1);
+        let mut reader = events.get_reader();
+        let last = reader.read(events).last().unwrap();
+        assert_eq!(last.exited, None);
+        assert_eq!(last.entered, Some(TestState::C));
+    }
+}