Repository: NiklasEi/bevy_asset_loader (rust)

Issue:
Incorrect behaviour when using two states
# System configuration

**OS**: Windows 11
**Bevy**:  0.10.1 (from crates.io)

`rustc --version`: `rustc 1.68.2 (9eb3afe9e 2023-03-27)`
`rustup default`: `stable-x86_64-pc-windows-msvc`

# What versions the problems happens on

This problem happens with `bevy_asset_loader 0.16` and the latest version on the github (when I fetched it the latest commit was 5cd90672eb5b9a8a2127bd225c3f9d3a57c963d4)

# The issue

I have been trying to implement some multistep loading in my game prototype by having two states -- one that tracks loading and the other tracks the game state. 

It seems that `bevy_asset_loader` does something incorrect in such configurations and the assets don't get loaded. 

# Minimal examples

Program 1 (Single probing)

```rust
use bevy::prelude::*;
use bevy_asset_loader::prelude::{AssetCollection, LoadingStateAppExt, LoadingState};

#[derive(Clone, Copy, Debug, States, Default, PartialEq, Eq, Hash)]
enum Game {
    #[default]
    Booting,
    Loading,
    Play,
}

#[derive(Clone, Copy, Debug, States, Default, PartialEq, Eq, Hash)]
enum Loading {
    #[default]
    Done,
    Loading,
    Finalize,
}

#[derive(Resource, Default)]
pub struct Probe;

#[derive(Resource, AssetCollection)]
pub struct Collection1 {
    #[asset(path = "bevy1.png")]
    _img: Handle<Image>,
}


#[derive(Resource, AssetCollection)]
pub struct Collection2 {
    #[asset(path = "bevy2.png")]
    _img: Handle<Image>,
}

fn probe(_res: Res<Collection2>) {}

fn transition1(
    mut next_l: ResMut<NextState<Loading>>,
) {
    next_l.0 = Some(Loading::Loading);
}

fn transition2(
    mut next_g: ResMut<NextState<Game>>,
    mut next_l: ResMut<NextState<Loading>>,
) {
    next_g.0 = Some(Game::Play);
    next_l.0 = Some(Loading::Done);
}

fn main() {
    App::new()
    .add_plugins(DefaultPlugins)
    .add_state::<Loading>()
    .add_state::<Game>()
    .add_loading_state(
        LoadingState::new(Game::Booting)
        .continue_to_state(Game::Loading)
    )
    .add_collection_to_loading_state::<_, Collection1>(Game::Booting)
    .add_loading_state(
        LoadingState::new(Loading::Loading)
        .continue_to_state(Loading::Finalize)
    )
    .add_collection_to_loading_state::<_, Collection2>(Loading::Loading)
    .add_system(transition1.in_schedule(OnEnter(Game::Loading)))
    .add_system(transition2.in_schedule(OnEnter(Loading::Finalize)))
    .add_system(probe.in_schedule(OnEnter(Game::Play)))
    .run();
}
```

Program 2 (repeated probing)

```rust
use bevy::prelude::*;
use bevy_asset_loader::prelude::{AssetCollection, LoadingStateAppExt, LoadingState};

#[derive(Clone, Copy, Debug, States, Default, PartialEq, Eq, Hash)]
enum Game {
    #[default]
    Booting,
    Loading,
    Play,
}

#[derive(Clone, Copy, Debug, States, Default, PartialEq, Eq, Hash)]
enum Loading {
    #[default]
    Done,
    Loading,
    Finalize,
}

#[derive(Resource, Default)]
pub struct Probe;

#[derive(Resource, AssetCollection)]
pub struct Collection1 {
    #[asset(path = "bevy1.png")]
    _img: Handle<Image>,
}


#[derive(Resource, AssetCollection)]
pub struct Collection2 {
    #[asset(path = "bevy2.png")]
    _img: Handle<Image>,
}

fn probe(res: Option<Res<Collection2>>) {
    info!("probe: {}", res.is_some());
}

fn transition1(
    mut next_l: ResMut<NextState<Loading>>,
) {
    next_l.0 = Some(Loading::Loading);
}

fn transition2(
    mut next_g: ResMut<NextState<Game>>,
    mut next_l: ResMut<NextState<Loading>>,
) {
    next_g.0 = Some(Game::Play);
    next_l.0 = Some(Loading::Done);
}

fn main() {
    App::new()
    .add_plugins(DefaultPlugins)
    .add_state::<Loading>()
    .add_state::<Game>()
    .add_loading_state(
        LoadingState::new(Game::Booting)
        .continue_to_state(Game::Loading)
    )
    .add_collection_to_loading_state::<_, Collection1>(Game::Booting)
    .add_loading_state(
        LoadingState::new(Loading::Loading)
        .continue_to_state(Loading::Finalize)
    )
    .add_collection_to_loading_state::<_, Collection2>(Loading::Loading)
    .add_system(transition1.in_schedule(OnEnter(Game::Loading)))
    .add_system(transition2.in_schedule(OnEnter(Loading::Finalize)))
    .add_system(probe.in_set(OnUpdate(Game::Play)))
    .run();
}
```

# What you expected to happen

The first program shouldn't crash. The second program should have printed `probe: true`

# What happened

The first program crashed

```
     Running `target\debug\bevy_ald.exe`
2023-04-10T12:07:27.081547Z  INFO bevy_asset_loader::loading_state::systems: Loading state 'Loading' is done
thread 'Compute Task Pool (7)' panicked at 'Resource requested by bevy_ald::probe does not exist: bevy_ald::Collection', bevy_ecs-0.10.1\src\system\system_param.rs:460:17
```

The second program not only prints `false`, but also never prints `true`

```
     Running `target\debug\bevy_ald.exe`
2023-04-10T12:06:12.621932Z  INFO bevy_asset_loader::loading_state::systems: Loading state 'Loading' is done
2023-04-10T12:06:12.623797Z  INFO bevy_ald: probe: false
2023-04-10T12:06:12.625324Z  INFO bevy_ald: probe: false
2023-04-10T12:06:12.627005Z  INFO bevy_ald: probe: false
2023-04-10T12:06:12.628147Z  INFO bevy_ald: probe: false
2023-04-10T12:06:12.629495Z  INFO bevy_ald: probe: false
2023-04-10T12:06:12.630807Z  INFO bevy_ald: probe: false
2023-04-10T12:06:12.632035Z  INFO bevy_ald: probe: false
2023-04-10T12:06:12.633344Z  INFO bevy_ald: probe: false
2023-04-10T12:06:12.634580Z  INFO bevy_ald: probe: false
2023-04-10T12:06:12.635828Z  INFO bevy_ald: probe: false
2023-04-10T12:06:12.637343Z  INFO bevy_ald: probe: false
2023-04-10T12:06:12.638524Z  INFO bevy_ald: probe: false
```

Pull request description:
Allow usage in multiple states
Resolves #125

Patch:
diff --git a/bevy_asset_loader/src/loading_state.rs b/bevy_asset_loader/src/loading_state.rs
index f3edf94..6a02d95 100644
--- a/bevy_asset_loader/src/loading_state.rs
+++ b/bevy_asset_loader/src/loading_state.rs
@@ -12,6 +12,7 @@ use bevy::ecs::system::Resource;
 use bevy::ecs::world::FromWorld;
 use bevy::utils::{default, HashMap, HashSet};
 use std::any::TypeId;
+use std::array::IntoIter;
 use std::marker::PhantomData;
 
 use crate::asset_collection::AssetCollection;
@@ -155,7 +156,7 @@ where
         }
     }
 
-    /// The [`LoadingState`] will set this Bevy [`State`](::bevy::ecs::schedule::State) after all asset collections
+    /// The [`LoadingState`] will set this Bevy [`State`](State) after all asset collections
     /// are loaded and inserted as resources.
     /// ```edition2021
     /// # use bevy_asset_loader::prelude::*;
@@ -202,7 +203,7 @@ where
         self
     }
 
-    /// The [`LoadingState`] will set this Bevy [`State`](::bevy::ecs::schedule::State) if an asset fails to load.
+    /// The [`LoadingState`] will set this Bevy [`State`](State) if an asset fails to load.
     /// ```edition2021
     /// # use bevy_asset_loader::prelude::*;
     /// # use bevy::prelude::*;
@@ -337,8 +338,8 @@ where
                 .state_configurations
                 .insert(self.loading_state.clone(), loading_config);
         }
-        app.init_resource::<State<InternalLoadingState>>();
-        app.init_resource::<NextState<InternalLoadingState>>();
+        app.init_resource::<State<InternalLoadingState<S>>>();
+        app.init_resource::<NextState<InternalLoadingState<S>>>();
 
         app.init_resource::<DynamicAssetCollections<S>>();
         #[cfg(feature = "standard_dynamic_assets")]
@@ -375,7 +376,7 @@ where
                 resume_to_loading_asset_collections::<S>
                     .in_schedule(loading_state_schedule.clone())
                     .in_set(InternalLoadingStateSet::ResumeDynamicAssetCollections),
-                initialize_loading_state
+                initialize_loading_state::<S>
                     .in_schedule(loading_state_schedule.clone())
                     .in_set(InternalLoadingStateSet::Initialize),
                 resume_to_finalize::<S>
@@ -385,7 +386,7 @@ where
                     .in_schedule(loading_state_schedule.clone())
                     .in_set(InternalLoadingStateSet::Finalize),
             ))
-            .add_system(reset_loading_state.in_schedule(OnEnter(self.loading_state.clone())))
+            .add_system(reset_loading_state::<S>.in_schedule(OnEnter(self.loading_state.clone())))
             .configure_set(
                 LoadingStateSet(self.loading_state.clone())
                     .after(CoreSet::StateTransitions)
@@ -396,27 +397,27 @@ where
             loading_state_schedule
                 .configure_set(
                     InternalLoadingStateSet::Initialize
-                        .run_if(in_state(InternalLoadingState::Initialize)),
+                        .run_if(in_state(InternalLoadingState::<S>::Initialize)),
                 )
                 .configure_set(
                     InternalLoadingStateSet::CheckDynamicAssetCollections.run_if(in_state(
-                        InternalLoadingState::LoadingDynamicAssetCollections,
+                        InternalLoadingState::<S>::LoadingDynamicAssetCollections,
                     )),
                 )
                 .configure_set(
                     InternalLoadingStateSet::ResumeDynamicAssetCollections
                         .after(InternalLoadingStateSet::CheckDynamicAssetCollections)
                         .run_if(in_state(
-                            InternalLoadingState::LoadingDynamicAssetCollections,
+                            InternalLoadingState::<S>::LoadingDynamicAssetCollections,
                         )),
                 )
                 .configure_set(
                     InternalLoadingStateSet::CheckAssets
-                        .run_if(in_state(InternalLoadingState::LoadingAssets)),
+                        .run_if(in_state(InternalLoadingState::<S>::LoadingAssets)),
                 )
                 .configure_set(
                     InternalLoadingStateSet::Finalize
-                        .run_if(in_state(InternalLoadingState::Finalize)),
+                        .run_if(in_state(InternalLoadingState::<S>::Finalize)),
                 );
 
             #[cfg(feature = "progress_tracking")]
@@ -455,12 +456,12 @@ pub(crate) enum InternalLoadingStateSet {
 }
 
 #[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
-pub(crate) struct OnEnterInternalLoadingState<S: States>(pub S, pub InternalLoadingState);
+pub(crate) struct OnEnterInternalLoadingState<S: States>(pub S, pub InternalLoadingState<S>);
 #[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
 pub(crate) struct LoadingStateSchedule<S: States>(pub S);
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, States)]
-pub(crate) enum InternalLoadingState {
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
+pub(crate) enum InternalLoadingState<S: States> {
     /// Starting point. Here it will be decided whether or not dynamic asset collections need to be loaded.
     #[default]
     Initialize,
@@ -471,7 +472,22 @@ pub(crate) enum InternalLoadingState {
     /// All collections are loaded and inserted. Time to e.g. run custom [insert_resource](bevy_asset_loader::AssetLoader::insert_resource).
     Finalize,
     /// A 'parking' state in case no next state is defined
-    Done,
+    Done(PhantomData<S>),
+}
+
+impl<S: States> States for InternalLoadingState<S> {
+    type Iter = IntoIter<Self, 5>;
+
+    fn variants() -> Self::Iter {
+        [
+            Self::Initialize,
+            Self::LoadingDynamicAssetCollections,
+            Self::LoadingAssets,
+            Self::Finalize,
+            Self::Done(PhantomData::default()),
+        ]
+        .into_iter()
+    }
 }
 
 /// This resource is used for handles from asset collections and loading dynamic asset collection files.
diff --git a/bevy_asset_loader/src/loading_state/dynamic_asset_systems.rs b/bevy_asset_loader/src/loading_state/dynamic_asset_systems.rs
index 3a37216..e0aee6a 100644
--- a/bevy_asset_loader/src/loading_state/dynamic_asset_systems.rs
+++ b/bevy_asset_loader/src/loading_state/dynamic_asset_systems.rs
@@ -91,7 +91,7 @@ pub(crate) fn check_dynamic_asset_collections<S: States, C: DynamicAssetCollecti
 
 pub(crate) fn resume_to_loading_asset_collections<S: States>(
     state: Res<State<S>>,
-    mut loading_state: ResMut<NextState<InternalLoadingState>>,
+    mut loading_state: ResMut<NextState<InternalLoadingState<S>>>,
     asset_loader_config: Res<AssetLoaderConfiguration<S>>,
 ) {
     let config = asset_loader_config
diff --git a/bevy_asset_loader/src/loading_state/systems.rs b/bevy_asset_loader/src/loading_state/systems.rs
index 24bcfc1..5bbeb4e 100644
--- a/bevy_asset_loader/src/loading_state/systems.rs
+++ b/bevy_asset_loader/src/loading_state/systems.rs
@@ -121,7 +121,7 @@ fn count_loaded_handles<S: States, Assets: AssetCollection>(cell: WorldCell) ->
 
 pub(crate) fn resume_to_finalize<S: States>(
     loader_configuration: Res<AssetLoaderConfiguration<S>>,
-    mut internal_state: ResMut<NextState<InternalLoadingState>>,
+    mut internal_state: ResMut<NextState<InternalLoadingState<S>>>,
     user_state: Res<State<S>>,
     mut next_user_state: ResMut<NextState<S>>,
 ) {
@@ -138,8 +138,8 @@ pub(crate) fn resume_to_finalize<S: States>(
     }
 }
 
-pub(crate) fn initialize_loading_state(
-    mut loading_state: ResMut<NextState<InternalLoadingState>>,
+pub(crate) fn initialize_loading_state<S: States>(
+    mut loading_state: ResMut<NextState<InternalLoadingState<S>>>,
     #[cfg(feature = "progress_tracking")] mut progress_counter: ResMut<ProgressCounter>,
 ) {
     #[cfg(feature = "progress_tracking")]
@@ -151,7 +151,7 @@ pub(crate) fn finish_loading_state<S: States>(
     state: Res<State<S>>,
     mut next_state: ResMut<NextState<S>>,
     #[cfg(feature = "progress_tracking")] mut progress_counter: ResMut<ProgressCounter>,
-    mut loading_state: ResMut<NextState<InternalLoadingState>>,
+    mut loading_state: ResMut<NextState<InternalLoadingState<S>>>,
     asset_loader_configuration: Res<AssetLoaderConfiguration<S>>,
 ) {
     #[cfg(feature = "progress_tracking")]
@@ -167,10 +167,10 @@ pub(crate) fn finish_loading_state<S: States>(
         }
     }
 
-    loading_state.set(InternalLoadingState::Done);
+    loading_state.set(InternalLoadingState::Done(PhantomData::default()));
 }
 
-pub(crate) fn reset_loading_state(mut state: ResMut<State<InternalLoadingState>>) {
+pub(crate) fn reset_loading_state<S: States>(mut state: ResMut<State<InternalLoadingState<S>>>) {
     state.0 = InternalLoadingState::Initialize;
 }
 
@@ -182,25 +182,28 @@ pub(crate) fn run_loading_state<S: States>(world: &mut World) {
 pub fn apply_internal_state_transition<S: States>(world: &mut World) {
     let state = world.resource::<State<S>>().0.clone();
     if world
-        .resource::<NextState<InternalLoadingState>>()
+        .resource::<NextState<InternalLoadingState<S>>>()
         .0
         .is_some()
     {
         let entered_state = world
-            .resource_mut::<NextState<InternalLoadingState>>()
+            .resource_mut::<NextState<InternalLoadingState<S>>>()
             .0
             .take()
             .unwrap();
         let exited_state = mem::replace(
-            &mut world.resource_mut::<State<InternalLoadingState>>().0,
-            entered_state,
+            &mut world.resource_mut::<State<InternalLoadingState<S>>>().0,
+            entered_state.clone(),
         );
         trace!(
             "Switching internal state of loading state from {exited_state:?} to {entered_state:?}"
         );
         if world
             .resource::<Schedules>()
-            .contains(&OnEnterInternalLoadingState(state.clone(), entered_state))
+            .contains(&OnEnterInternalLoadingState(
+                state.clone(),
+                entered_state.clone(),
+            ))
         {
             world.run_schedule(OnEnterInternalLoadingState(state, entered_state));
         }


Test patch:
diff --git a/bevy_asset_loader/tests/multiple_states.rs b/bevy_asset_loader/tests/multiple_states.rs
new file mode 100644
index 0000000..89ee4b4
--- /dev/null
+++ b/bevy_asset_loader/tests/multiple_states.rs
@@ -0,0 +1,87 @@
+#![allow(dead_code, unused_imports)]
+
+use bevy::app::AppExit;
+use bevy::audio::AudioPlugin;
+use bevy::prelude::*;
+use bevy_asset_loader::prelude::{AssetCollection, LoadingState, LoadingStateAppExt};
+
+#[cfg(all(
+    not(feature = "2d"),
+    not(feature = "3d"),
+    not(feature = "progress_tracking")
+))]
+#[test]
+fn main() {
+    App::new()
+        .add_plugins(MinimalPlugins)
+        .add_plugin(AssetPlugin::default())
+        .add_plugin(AudioPlugin::default())
+        .add_state::<Loading>()
+        .add_state::<Game>()
+        .add_loading_state(LoadingState::new(Game::Booting).continue_to_state(Game::Loading))
+        .add_collection_to_loading_state::<_, GameStateCollection>(Game::Booting)
+        .add_loading_state(LoadingState::new(Loading::Loading).continue_to_state(Loading::Finalize))
+        .add_collection_to_loading_state::<_, LoadingStateCollection>(Loading::Loading)
+        .add_system(go_to_loading_loading.in_schedule(OnEnter(Game::Loading)))
+        .add_system(probe_game_state.in_schedule(OnEnter(Game::Loading)))
+        .add_system(go_to_game_play_loading_done.in_schedule(OnEnter(Loading::Finalize)))
+        .add_system(probe_loading_state.in_schedule(OnEnter(Game::Play)))
+        .add_system(quit.in_set(OnUpdate(Game::Play)))
+        .add_system(timeout)
+        .run();
+}
+
+#[derive(Clone, Copy, Debug, States, Default, PartialEq, Eq, Hash)]
+enum Game {
+    #[default]
+    Booting,
+    Loading,
+    Play,
+}
+
+#[derive(Clone, Copy, Debug, States, Default, PartialEq, Eq, Hash)]
+enum Loading {
+    #[default]
+    Done,
+    Loading,
+    Finalize,
+}
+
+#[derive(Resource, AssetCollection)]
+pub struct GameStateCollection {
+    #[asset(path = "audio/background.ogg")]
+    background: Handle<AudioSource>,
+}
+
+fn probe_game_state(_res: Res<GameStateCollection>) {}
+
+#[derive(Resource, AssetCollection)]
+pub struct LoadingStateCollection {
+    #[asset(path = "audio/yipee.ogg")]
+    yipee: Handle<AudioSource>,
+}
+
+fn probe_loading_state(_res: Res<LoadingStateCollection>) {}
+
+fn go_to_loading_loading(mut state: ResMut<NextState<Loading>>) {
+    state.set(Loading::Loading);
+}
+
+fn go_to_game_play_loading_done(
+    mut game_state: ResMut<NextState<Game>>,
+    mut loading_state: ResMut<NextState<Loading>>,
+) {
+    game_state.set(Game::Play);
+    loading_state.set(Loading::Done);
+}
+
+fn timeout(time: Res<Time>) {
+    if time.elapsed_seconds_f64() > 30. {
+        panic!("The app did not finish in 30 seconds");
+    }
+}
+
+fn quit(mut exit: EventWriter<AppExit>) {
+    info!("Everything fine, quitting the app");
+    exit.send(AppExit);
+}