feat: asset loading check

src/assets.rs

@@ -1,3 +1,10 @@
+use std::{
+    any::TypeId,
+    sync::{LazyLock, Mutex},
+};
+
+use bevy::reflect::{GetTypeRegistration, ReflectFromPtr};
+
 use crate::prelude::*;
 
 #[cfg(feature = "embedded")]
@@ -7,8 +14,11 @@ mod meta;
 mod music;
 mod sound;
 
+static ASSET_MAP: LazyLock<Mutex<Vec<TypeId>>> = LazyLock::new(|| Mutex::new(vec![]));
+
 pub(super) fn plugin(app: &mut App) {
-    app.add_plugins((fonts::plugin, meta::plugin, music::plugin, sound::plugin));
+    app.add_plugins((fonts::plugin, meta::plugin, music::plugin, sound::plugin))
+        .add_systems(Update, check_loaded.run_if(in_state(GameState::Startup)));
 }
 
 pub mod prelude {
@@ -17,18 +27,23 @@ pub mod prelude {
         meta::MetaAssetKey,
         music::MusicAssetKey,
         sound::SoundAssetKey,
+        AssetExt,
         AssetKey,
         AssetMap,
     };
 }
 
 /// Represents a handle to any asset type
-pub trait AssetKey: Sized + Eq + std::hash::Hash {
+pub trait AssetKey:
+    Sized + Eq + std::hash::Hash + Reflect + FromReflect + TypePath + GetTypeRegistration
+{
     type Asset: Asset;
 }
 
 /// A resource that holds asset `Handle`s for a particular type of `AssetKey`
-#[derive(Resource, Deref, DerefMut)]
+/// Easy to access on any system using `Res<AssetMap<...>>`
+#[derive(Resource, Reflect, Deref, DerefMut)]
+#[reflect(AssetsLoaded)]
 pub struct AssetMap<K: AssetKey>(HashMap<K, Handle<K::Asset>>);
 
 impl<K: AssetKey, T> From<T> for AssetMap<K>
@@ -41,14 +56,104 @@ where
 }
 
 impl<K: AssetKey> AssetMap<K> {
+    /// Returns a weak clone of the asset handle
+    pub fn get(&self, key: &K) -> Handle<K::Asset> {
+        self[key].clone_weak()
+    }
+}
+
+#[reflect_trait]
+trait AssetsLoaded {
     /// Check if all of the assets are loaded
-    pub fn all_loaded(&self, asset_server: &AssetServer) -> bool {
+    fn all_loaded(&self, asset_server: &AssetServer) -> bool;
+}
+
+impl<K: AssetKey> AssetsLoaded for AssetMap<K> {
+    fn all_loaded(&self, asset_server: &AssetServer) -> bool {
         self.values()
             .all(|x| asset_server.is_loaded_with_dependencies(x))
     }
+}
 
-    /// Returns a weak clone of the asset handle
-    pub fn get(&self, key: &K) -> Handle<K::Asset> {
-        self[key].clone_weak()
+pub trait AssetExt {
+    fn load_asset<K: AssetKey>(&mut self) -> &mut Self
+    where
+        AssetMap<K>: FromWorld;
+}
+
+impl AssetExt for App {
+    fn load_asset<K: AssetKey>(&mut self) -> &mut Self
+    where
+        AssetMap<K>: FromWorld,
+    {
+        ASSET_MAP.lock().unwrap().push(TypeId::of::<AssetMap<K>>());
+        self.init_resource::<AssetMap<K>>()
+            .register_type::<AssetMap<K>>()
+    }
+}
+
+fn check_loaded(world: &mut World) {
+    let mut map = ASSET_MAP.lock().unwrap();
+    let mut loaded = vec![];
+
+    for id in map.iter() {
+        match is_resource_loaded(*id, world) {
+            Ok(l) if !l => continue,
+            Err(e) => warn!("{}", e),
+            _ => {},
+        }
+        loaded.push(*id);
+    }
+
+    (*map).retain(|x| !loaded.contains(x));
+
+    if map.len() == 0 {
+        let mut next_state = world
+            .get_resource_mut::<NextState<GameState>>()
+            .expect("NextState should exist");
+        next_state.set(GameState::Menu);
     }
 }
+
+fn is_resource_loaded(id: TypeId, world: &World) -> Result<bool> {
+    // Get world resources
+    let asset_server = world
+        .get_resource::<AssetServer>()
+        .expect("Bevy's asset server should exist");
+    let registry = world
+        .get_resource::<AppTypeRegistry>()
+        .expect("Bevy's type registry should exist");
+    let registry = registry.read();
+
+    // Get the AssetMap component id and raw pointer
+    let registration = registry.get(id).context("Id is not registered")?;
+    let component_id = world
+        .components()
+        .get_resource_id(registration.type_id())
+        .context("Couldn't get the component id of the resource")?;
+    let ptr = world
+        .get_resource_by_id(component_id)
+        .context("The resource is not registred")?;
+
+    // Convert the pointer into a Reflect trait
+    // This should be safe
+    let reflect_from_ptr = registration
+        .data::<ReflectFromPtr>()
+        .context("Type registration should exist")?;
+    // SAFETY: from the context it is known that `ReflectFromPtr` was made for the
+    // type of the `MutUntyped`
+    let resource: &dyn Reflect = unsafe { reflect_from_ptr.as_reflect(ptr) };
+
+    // Get the LoadedAsset trait registration
+    let loaded_trait = registry
+        .get_type_data::<ReflectAssetsLoaded>(id)
+        .context("The AssetLoaded trait is not registered")?;
+
+    // Convert the AssetMap dyn Reflect object into a dyn LoadedAsset trait
+    let resource = loaded_trait
+        .get(resource)
+        .context("The resource doesn't implement the LoadedAsset trait")?;
+
+    // Check if all of the assets are loaded
+    Ok(resource.all_loaded(asset_server))
+}

src/assets/fonts.rs

@@ -2,7 +2,7 @@ use crate::prelude::*;
 
 /// Preloads the font assets when the game starts
 pub(super) fn plugin(app: &mut App) {
-    app.init_resource::<AssetMap<FontAssetKey>>();
+    app.load_asset::<FontAssetKey>();
 }
 
 /// Defines all of the font assets

src/assets/meta.rs

@@ -2,11 +2,10 @@ use crate::prelude::*;
 
 /// Preloads the meta assets when the game starts
 pub(super) fn plugin(app: &mut App) {
-    app.init_resource::<AssetMap<MetaAssetKey>>();
+    app.load_asset::<MetaAssetKey>();
 }
 
 /// Defines all of the meta assets
-/// Easy to access on any system using `Res<AssetMap<MetaAssetKey>>`
 #[asset_key(Image)]
 pub enum MetaAssetKey {
     #[asset = "meta/bevy.png"]

src/assets/music.rs

@@ -2,7 +2,7 @@ use crate::prelude::*;
 
 /// Preloads the music assets when the game starts
 pub(super) fn plugin(app: &mut App) {
-    app.init_resource::<AssetMap<MusicAssetKey>>();
+    app.load_asset::<MusicAssetKey>();
 }
 
 /// Defines all of the musical assets

src/assets/sound.rs

@@ -2,7 +2,7 @@ use crate::prelude::*;
 
 /// Preloads the sound assets when the game starts
 pub(super) fn plugin(app: &mut App) {
-    app.init_resource::<AssetMap<SoundAssetKey>>();
+    app.load_asset::<SoundAssetKey>();
 }
 
 /// Defines all of the sound effects

src/base/states.rs

@@ -16,12 +16,8 @@ pub(super) fn plugin(app: &mut App) {
 pub enum GameState {
     /// The game starts on the `Startup` state.
     /// It runs before *anything*, including the `Startup` schedule.
-    /// It inmediately transitions to `Loading`.
     #[default]
     Startup,
-    /// Handles splash screens and assets.
-    /// The game stays here until all of the assets are ready.
-    Loading,
     /// The main menu of the game. All of the game systems are paused.
     Menu,
     /// Main state representing the actual gameplay.

src/components/camera.rs

@@ -15,13 +15,7 @@ pub struct GameCamera;
 pub struct FinalCamera;
 
 /// Spawn the main cameras
-fn init(mut cmd: Commands, meta_assets: Res<AssetMap<MetaAssetKey>>) {
+fn init(mut cmd: Commands) {
     let camera_bundle = Camera2dBundle::default();
     cmd.spawn((camera_bundle, GameCamera, FinalCamera));
-
-    // Test logo, delete
-    cmd.spawn(SpriteBundle {
-        texture: meta_assets.get(&MetaAssetKey::BevyLogo),
-        ..default()
-    });
 }

src/ui.rs

@@ -1,28 +1,19 @@
 use crate::prelude::*;
 
+mod menu;
 mod navigation;
 mod widgets;
 
 pub(super) fn plugin(app: &mut App) {
-    app.add_plugins((navigation::plugin, widgets::plugin));
-    app.add_systems(OnEnter(GameState::Startup), init);
+    app.add_plugins((menu::plugin, navigation::plugin, widgets::plugin));
 }
 
 pub mod prelude {
     pub use bevy_trait_query::RegisterExt;
 
     pub use super::{
+        menu::MenuState,
         navigation::{NavBundle, NavContainer, Navigable},
         widgets::{Container, NavigableExt, Stylable, Widget},
     };
 }
-
-fn init(mut cmd: Commands) {
-    let mut root = cmd.ui_root();
-    root.with_children(|node| {
-        node.button("hey");
-        node.button("hi").no_nav();
-        node.button("hello");
-    })
-    .nav_container();
-}

src/ui/menu.rs

@@ -0,0 +1,21 @@
+use crate::prelude::*;
+
+mod main;
+
+pub(super) fn plugin(app: &mut App) {
+    app.add_sub_state::<MenuState>()
+        .enable_state_scoped_entities::<MenuState>()
+        .add_plugins(main::plugin);
+}
+
+#[derive(SubStates, Std!, Default)]
+#[source(GameState = GameState::Menu)]
+pub enum MenuState {
+    /// Main menu screen, used to play or exit the game and access other options
+    #[default]
+    Main,
+    /// Menu screen to customize game options
+    Options,
+    /// Menu screen to view keys assigned to actions
+    Mappings,
+}

src/ui/menu/main.rs

@@ -0,0 +1,16 @@
+use crate::prelude::*;
+
+pub(super) fn plugin(app: &mut App) {
+    app.add_systems(OnEnter(MenuState::Main), init);
+}
+
+fn init(mut cmd: Commands) {
+    cmd.ui_root()
+        .with_children(|root| {
+            root.button("hey");
+            root.button("hi");
+            root.button("hello");
+        })
+        .nav_container()
+        .insert(StateScoped(MenuState::Main));
+}