Fix nested immediate asset loading.

crates/bevy_asset/src/lib.rs

@@ -719,15 +719,15 @@ mod tests {
             let mut ron: CoolTextRon = ron::de::from_bytes(&bytes)?;
             let mut embedded = String::new();
             for dep in ron.embedded_dependencies {
-                let loaded = load_context
+                let (_, loaded) = load_context
                     .loader()
                     .immediate()
                     .load::<CoolText>(&dep)
                     .await
                     .map_err(|_| Self::Error::CannotLoadDependency {
                         dependency: dep.into(),
                     })?;
-                let cool = loaded.get_asset().get();
+                let cool = loaded.get_asset();
                 embedded.push_str(&cool.text);
             }
             Ok(CoolText {
@@ -1007,6 +1007,20 @@ mod tests {
         // Re-open a and b gates to allow c to load embedded deps (gates are closed after each load)
         gate_opener.open(a_path);
         gate_opener.open(b_path);
+
+        // Wait for the C-load task to finish. If we don't do this, it's possible for nested asset A
+        // to be sent in one frame and for the nested asset B to be sent in the next frame. This
+        // causes an alternative order of events (but one that is just as valid), which would make
+        // this test flaky.
+        while !asset_server
+            .data
+            .infos
+            .read()
+            .pending_tasks
+            .iter()
+            .any(|(_, task)| task.is_finished())
+        {}
+
         run_app_until(&mut app, |world| {
             let a_text = get::<CoolText>(world, a_id)?;
             let (a_load, a_deps, a_rec_deps) = asset_server.get_load_states(a_id).unwrap();
@@ -1131,9 +1145,19 @@ mod tests {
             AssetEvent::Added {
                 id: id_results.b_id,
             },
+            // This extra LoadedWithDependencies happens because asset B got replaced when C was
+            // loaded (and we only see the modify after since events on `Assets` only merge in a
+            // system, so they look like they happen at the wrong time).
+            AssetEvent::LoadedWithDependencies {
+                id: id_results.b_id,
+            },
             AssetEvent::Added {
                 id: id_results.c_id,
             },
+            AssetEvent::Modified { id: a_id },
+            AssetEvent::Modified {
+                id: id_results.b_id,
+            },
             AssetEvent::LoadedWithDependencies {
                 id: id_results.d_id,
             },

crates/bevy_asset/src/loader.rs

@@ -3,8 +3,8 @@ use crate::{
     loader_builders::{Deferred, NestedLoader, StaticTyped},
     meta::{AssetHash, AssetMeta, AssetMetaDyn, ProcessedInfoMinimal, Settings},
     path::AssetPath,
-    Asset, AssetLoadError, AssetServer, AssetServerMode, Assets, Handle, UntypedAssetId,
-    UntypedHandle,
+    Asset, AssetLoadError, AssetServer, AssetServerMode, Assets, Handle, NestedAssets,
+    UntypedAssetId, UntypedHandle,
 };
 use alloc::{
     boxed::Box,
@@ -15,8 +15,12 @@ use atomicow::CowArc;
 use bevy_ecs::world::World;
 use bevy_platform_support::collections::{HashMap, HashSet};
 use bevy_tasks::{BoxedFuture, ConditionalSendFuture};
-use core::any::{Any, TypeId};
+use core::{
+    any::{Any, TypeId},
+    ops::Deref,
+};
 use downcast_rs::{impl_downcast, Downcast};
+use parking_lot::{RwLock, RwLockReadGuard};
 use ron::error::SpannedError;
 use serde::{Deserialize, Serialize};
 use std::path::{Path, PathBuf};
@@ -383,6 +387,9 @@ pub enum DeserializeMetaError {
 /// Any asset state accessed by [`LoadContext`] will be tracked and stored for use in dependency events and asset preprocessing.
 pub struct LoadContext<'a> {
     pub(crate) asset_server: &'a AssetServer,
+    /// The nested assets that have been directly loaded across the entire loader. We use a
+    /// [`RwLock`] so that all [`LoadContext`]'s based on this one share these nested loaded assets.
+    pub(crate) nested_direct_loaded_assets: &'a RwLock<NestedAssets>,
     pub(crate) should_load_dependencies: bool,
     populate_hashes: bool,
     asset_path: AssetPath<'static>,
@@ -396,12 +403,14 @@ impl<'a> LoadContext<'a> {
     /// Creates a new [`LoadContext`] instance.
     pub(crate) fn new(
         asset_server: &'a AssetServer,
+        nested_direct_loaded_assets: &'a RwLock<NestedAssets>,
         asset_path: AssetPath<'static>,
         should_load_dependencies: bool,
         populate_hashes: bool,
     ) -> Self {
         Self {
             asset_server,
+            nested_direct_loaded_assets,
             asset_path,
             populate_hashes,
             should_load_dependencies,
@@ -443,6 +452,7 @@ impl<'a> LoadContext<'a> {
     pub fn begin_labeled_asset(&self) -> LoadContext {
         LoadContext::new(
             self.asset_server,
+            self.nested_direct_loaded_assets,
             self.asset_path.clone(),
             self.should_load_dependencies,
             self.populate_hashes,
@@ -613,7 +623,7 @@ impl<'a> LoadContext<'a> {
         meta: &dyn AssetMetaDyn,
         loader: &dyn ErasedAssetLoader,
         reader: &mut dyn Reader,
-    ) -> Result<CompleteErasedLoadedAsset, LoadDirectError> {
+    ) -> Result<NestedErasedAssetRef<'_>, LoadDirectError> {
         let complete_asset = self
             .asset_server
             .load_with_meta_loader_and_reader(
@@ -623,6 +633,7 @@ impl<'a> LoadContext<'a> {
                 reader,
                 false,
                 self.populate_hashes,
+                self.nested_direct_loaded_assets,
             )
             .await
             .map_err(|error| LoadDirectError::LoadError {
@@ -631,8 +642,25 @@ impl<'a> LoadContext<'a> {
             })?;
         let info = meta.processed_info().as_ref();
         let hash = info.map(|i| i.full_hash).unwrap_or_default();
-        self.loader_dependencies.insert(path, hash);
-        Ok(complete_asset)
+        self.loader_dependencies.insert(path.clone(), hash);
+        self.nested_direct_loaded_assets
+            .write()
+            .insert(path.clone(), complete_asset);
+        Ok(
+            NestedErasedAssetRef::new(self.nested_direct_loaded_assets.read(), &path)
+                .expect("asset path is loaded, since we just inserted it"),
+        )
+    }
+
+    /// Gets a previously direct-loaded nested asset by its path.
+    ///
+    /// Returns None if the asset either hasn't been loaded yet or hasn't finished loading. See
+    /// [`Self::loader`] for starting nested asset loads.
+    pub fn get_direct_nested_asset<'p>(
+        &self,
+        path: impl Into<AssetPath<'p>>,
+    ) -> Option<NestedErasedAssetRef> {
+        NestedErasedAssetRef::new(self.nested_direct_loaded_assets.read(), &path.into())
     }
 
     /// Create a builder for loading nested assets in this context.
@@ -654,6 +682,93 @@ impl<'a> LoadContext<'a> {
     }
 }
 
+/// A reference to a nested, directly-loaded asset. This is similar to
+/// [`CompleteErasedLoadedAsset`].
+pub struct NestedErasedAssetRef<'context> {
+    /// The lock that we hold to access the nested asset.
+    _lock: RwLockReadGuard<'context, NestedAssets>,
+    /// A pointer to the asset we are referencing. We store a pointer since otherwise, we'd need to
+    /// look up the asset every time we access the asset (and store the asset path here as well).
+    asset: *const CompleteErasedLoadedAsset,
+}
+
+impl Deref for NestedErasedAssetRef<'_> {
+    type Target = CompleteErasedLoadedAsset;
+
+    fn deref(&self) -> &Self::Target {
+        #[expect(
+            unsafe_code,
+            reason = "Not using raw pointers requires looking up the asset every time, and including the asset path in this struct"
+        )]
+        // SAFETY: We got this pointer from a reference, and we hold `self.lock`, so no mutable
+        // references alias this pointer.
+        unsafe {
+            &*self.asset
+        }
+    }
+}
+
+impl<'context> NestedErasedAssetRef<'context> {
+    fn new(
+        lock: RwLockReadGuard<'context, NestedAssets>,
+        asset_path: &AssetPath<'_>,
+    ) -> Option<Self> {
+        lock.get(asset_path)
+            .map(|asset| &raw const *asset)
+            .map(|asset| Self { _lock: lock, asset })
+    }
+
+    pub fn downcast<A: Asset>(self) -> Result<NestedAssetRef<'context, A>, Self> {
+        let Some(asset) = self.get_asset().get().map(|asset| &raw const *asset) else {
+            return Err(self);
+        };
+        Ok(NestedAssetRef {
+            erased_ref: self,
+            asset,
+        })
+    }
+}
+
+/// A reference to a nested, directly-loaded asset. This is similar to [`CompleteLoadedAsset`].
+pub struct NestedAssetRef<'context, A: Asset> {
+    /// The type-erased asset reference.
+    erased_ref: NestedErasedAssetRef<'context>,
+    /// A pointer to the typed asset we are referencing. We store a pointer since otherwise, we'd
+    /// need to do the cast every time (even though we've verified this is the correct type).
+    asset: *const A,
+}
+
+impl<A: Asset> NestedAssetRef<'_, A> {
+    /// Returns the stored asset.
+    pub fn get_asset(&self) -> &A {
+        #[expect(
+            unsafe_code,
+            reason = "Not using raw pointers requires casting the asset every time"
+        )]
+        // SAFETY: We got this pointer from a reference, and we hold `self.erased_ref`, so no
+        // mutable references alias this pointer.
+        unsafe {
+            &*self.asset
+        }
+    }
+
+    /// Returns the [`ErasedLoadedAsset`] for the given label, if it exists.
+    pub fn get_labeled(
+        &self,
+        label: impl Into<CowArc<'static, str>>,
+    ) -> Option<&ErasedLoadedAsset> {
+        self.erased_ref
+            .labeled_assets
+            .get(&label.into())
+            .map(|a| &a.asset)
+    }
+
+    /// Iterate over all labels for "labeled assets" in the loaded asset
+    pub fn iter_labels(&self) -> impl Iterator<Item = &str> {
+        self.erased_ref.labeled_assets.keys().map(|s| &**s)
+    }
+}
+
 /// An error produced when calling [`LoadContext::read_asset_bytes`]
 #[derive(Error, Debug)]
 pub enum ReadAssetBytesError {