Fix AssetTransformer breaking LabeledAssets

crates/bevy_asset/src/processor/process.rs

@@ -6,8 +6,8 @@ use crate::{
     meta::{AssetAction, AssetMeta, AssetMetaDyn, ProcessDependencyInfo, ProcessedInfo, Settings},
     processor::AssetProcessor,
     saver::{AssetSaver, SavedAsset},
-    transformer::AssetTransformer,
-    AssetLoadError, AssetLoader, AssetPath, DeserializeMetaError, ErasedLoadedAsset, LoadedAsset,
+    transformer::{AssetTransformer, TransformedAsset},
+    AssetLoadError, AssetLoader, AssetPath, DeserializeMetaError, ErasedLoadedAsset,
     MissingAssetLoaderForExtensionError, MissingAssetLoaderForTypeNameError,
 };
 use bevy_utils::BoxedFuture;
@@ -158,6 +158,8 @@ pub enum ProcessError {
     WrongMetaType,
     #[error("Encountered an error while saving the asset: {0}")]
     AssetSaveError(#[from] Box<dyn std::error::Error + Send + Sync + 'static>),
+    #[error("Encountered an error while transforming the asset: {0}")]
+    AssetTransformError(Box<dyn std::error::Error + Send + Sync + 'static>),
     #[error("Assets without extensions are not supported.")]
     ExtensionRequired,
 }
@@ -185,18 +187,19 @@ impl<
                 loader: std::any::type_name::<Loader>().to_string(),
                 settings: settings.loader_settings,
             });
-            let loaded_asset = context
-                .load_source_asset(loader_meta)
-                .await?
-                .take::<Loader::Asset>()
-                .expect("Asset type is known");
-            let transformed_asset = self
+            let pre_transformed_asset = TransformedAsset::<Loader::Asset>::from_loaded(
+                context.load_source_asset(loader_meta).await?,
+            )
+            .unwrap();
+
+            let post_transformed_asset = self
                 .transformer
-                .transform(loaded_asset, &settings.transformer_settings)?;
-            let loaded_transformed_asset =
-                ErasedLoadedAsset::from(LoadedAsset::from(transformed_asset));
+                .transform(pre_transformed_asset, &settings.transformer_settings)
+                .await
+                .map_err(|err| ProcessError::AssetTransformError(err.into()))?;
+
             let saved_asset =
-                SavedAsset::<T::AssetOutput>::from_loaded(&loaded_transformed_asset).unwrap();
+                SavedAsset::<T::AssetOutput>::from_transformed(&post_transformed_asset);
 
             let output_settings = self
                 .saver

crates/bevy_asset/src/saver.rs

@@ -1,5 +1,6 @@
+use crate::transformer::TransformedAsset;
 use crate::{io::Writer, meta::Settings, Asset, ErasedLoadedAsset};
-use crate::{AssetLoader, LabeledAsset, UntypedHandle};
+use crate::{AssetLoader, Handle, LabeledAsset, UntypedHandle};
 use bevy_utils::{BoxedFuture, CowArc, HashMap};
 use serde::{Deserialize, Serialize};
 use std::{borrow::Borrow, hash::Hash, ops::Deref};
@@ -88,6 +89,14 @@ impl<'a, A: Asset> SavedAsset<'a, A> {
         })
     }
 
+    /// Creates a new [`SavedAsset`] from the a [`TransformedAsset`]
+    pub fn from_transformed(asset: &'a TransformedAsset<A>) -> Self {
+        Self {
+            value: &asset.value,
+            labeled_assets: &asset.labeled_assets,
+        }
+    }
+
     /// Retrieves the value of this asset.
     #[inline]
     pub fn get(&self) -> &'a A {
@@ -119,13 +128,26 @@ impl<'a, A: Asset> SavedAsset<'a, A> {
     }
 
     /// Returns the [`UntypedHandle`] of the labeled asset with the provided 'label', if it exists.
-    pub fn get_untyped_handle<Q>(&self, label: &Q) -> Option<&UntypedHandle>
+    pub fn get_untyped_handle<Q>(&self, label: &Q) -> Option<UntypedHandle>
+    where
+        CowArc<'static, str>: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        let labeled = self.labeled_assets.get(label)?;
+        Some(labeled.handle.clone())
+    }
+
+    /// Returns the [`Handle`] of the labeled asset with the provided 'label', if it exists and is an asset of type `B`
+    pub fn get_handle<Q, B: Asset>(&self, label: &Q) -> Option<Handle<B>>
     where
         CowArc<'static, str>: Borrow<Q>,
         Q: ?Sized + Hash + Eq,
     {
         let labeled = self.labeled_assets.get(label)?;
-        Some(&labeled.handle)
+        if let Ok(handle) = labeled.handle.clone().try_typed::<B>() {
+            return Some(handle);
+        }
+        None
     }
 
     /// Iterate over all labels for "labeled assets" in the loaded asset

crates/bevy_asset/src/transformer.rs

@@ -1,5 +1,11 @@
-use crate::{meta::Settings, Asset};
+use crate::{meta::Settings, Asset, ErasedLoadedAsset, Handle, LabeledAsset, UntypedHandle};
+use bevy_utils::{BoxedFuture, CowArc, HashMap};
 use serde::{Deserialize, Serialize};
+use std::{
+    borrow::Borrow,
+    hash::Hash,
+    ops::{Deref, DerefMut},
+};
 
 /// Transforms an [`Asset`] of a given [`AssetTransformer::AssetInput`] type to an [`Asset`] of [`AssetTransformer::AssetOutput`] type.
 pub trait AssetTransformer: Send + Sync + 'static {
@@ -9,12 +15,228 @@ pub trait AssetTransformer: Send + Sync + 'static {
     type AssetOutput: Asset;
     /// The settings type used by this [`AssetTransformer`].
     type Settings: Settings + Default + Serialize + for<'a> Deserialize<'a>;
-    /// The type of [error](`std::error::Error`) which could be encountered by this saver.
+    /// The type of [error](`std::error::Error`) which could be encountered by this transformer.
     type Error: Into<Box<dyn std::error::Error + Send + Sync + 'static>>;
 
+    /// Transformes the given [`TransformedAsset`] to [`AssetTransformer::AssetOutput`].
+    /// The [`TransformedAsset`]'s `labeled_assets` can be altered to add new Labeled Sub-Assets
+    /// The passed in `settings` can influence how the `asset` is transformed
     fn transform<'a>(
         &'a self,
-        asset: Self::AssetInput,
+        asset: TransformedAsset<Self::AssetInput>,
         settings: &'a Self::Settings,
-    ) -> Result<Self::AssetOutput, Box<dyn std::error::Error + Send + Sync + 'static>>;
+    ) -> BoxedFuture<'a, Result<TransformedAsset<Self::AssetOutput>, Self::Error>>;
+}
+
+/// An [`Asset`] (and any "sub assets") intended to be transformed
+pub struct TransformedAsset<A: Asset> {
+    pub(crate) value: A,
+    pub(crate) labeled_assets: HashMap<CowArc<'static, str>, LabeledAsset>,
+}
+
+impl<A: Asset> Deref for TransformedAsset<A> {
+    type Target = A;
+    fn deref(&self) -> &Self::Target {
+        &self.value
+    }
+}
+
+impl<A: Asset> DerefMut for TransformedAsset<A> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.value
+    }
+}
+
+impl<A: Asset> TransformedAsset<A> {
+    /// Creates a new [`TransformedAsset`] from `asset` if its internal value matches `A`.
+    pub fn from_loaded(asset: ErasedLoadedAsset) -> Option<Self> {
+        if let Ok(value) = asset.value.downcast::<A>() {
+            return Some(TransformedAsset {
+                value: *value,
+                labeled_assets: asset.labeled_assets,
+            });
+        }
+        None
+    }
+    /// Creates a new [`TransformedAsset`] from `asset`, transfering the `labeled_assets` from this [`TransformedAsset`] to the new one
+    pub fn replace_asset<B: Asset>(self, asset: B) -> TransformedAsset<B> {
+        TransformedAsset {
+            value: asset,
+            labeled_assets: self.labeled_assets,
+        }
+    }
+    /// Takes the labeled assets from `labeled_source` and places them in this [`TransformedAsset`]
+    pub fn take_labeled_assets<B: Asset>(&mut self, labeled_source: TransformedAsset<B>) {
+        self.labeled_assets = labeled_source.labeled_assets;
+    }
+    /// Retrieves the value of this asset.
+    #[inline]
+    pub fn get(&self) -> &A {
+        &self.value
+    }
+    /// Mutably retrieves the value of this asset.
+    #[inline]
+    pub fn get_mut(&mut self) -> &mut A {
+        &mut self.value
+    }
+    /// Returns the labeled asset, if it exists and matches this type.
+    pub fn get_labeled<B: Asset, Q>(&mut self, label: &Q) -> Option<TransformedSubAsset<B>>
+    where
+        CowArc<'static, str>: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        let labeled = self.labeled_assets.get_mut(label)?;
+        let value = labeled.asset.value.downcast_mut::<B>()?;
+        Some(TransformedSubAsset {
+            value,
+            labeled_assets: &mut labeled.asset.labeled_assets,
+        })
+    }
+    /// Returns the type-erased labeled asset, if it exists and matches this type.
+    pub fn get_erased_labeled<Q>(&self, label: &Q) -> Option<&ErasedLoadedAsset>
+    where
+        CowArc<'static, str>: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        let labeled = self.labeled_assets.get(label)?;
+        Some(&labeled.asset)
+    }
+    /// Returns the [`UntypedHandle`] of the labeled asset with the provided 'label', if it exists.
+    pub fn get_untyped_handle<Q>(&self, label: &Q) -> Option<UntypedHandle>
+    where
+        CowArc<'static, str>: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        let labeled = self.labeled_assets.get(label)?;
+        Some(labeled.handle.clone())
+    }
+    /// Returns the [`Handle`] of the labeled asset with the provided 'label', if it exists and is an asset of type `B`
+    pub fn get_handle<Q, B: Asset>(&self, label: &Q) -> Option<Handle<B>>
+    where
+        CowArc<'static, str>: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        let labeled = self.labeled_assets.get(label)?;
+        if let Ok(handle) = labeled.handle.clone().try_typed::<B>() {
+            return Some(handle);
+        }
+        None
+    }
+    /// Adds `asset` as a labeled sub asset using `label` and `handle`
+    pub fn insert_labeled(
+        &mut self,
+        label: impl Into<CowArc<'static, str>>,
+        handle: impl Into<UntypedHandle>,
+        asset: impl Into<ErasedLoadedAsset>,
+    ) {
+        let labeled = LabeledAsset {
+            asset: asset.into(),
+            handle: handle.into(),
+        };
+        self.labeled_assets.insert(label.into(), labeled);
+    }
+    /// Iterate over all labels for "labeled assets" in the loaded asset
+    pub fn iter_labels(&self) -> impl Iterator<Item = &str> {
+        self.labeled_assets.keys().map(|s| &**s)
+    }
+}
+
+/// A labeled sub-asset of [`TransformedAsset`]
+pub struct TransformedSubAsset<'a, A: Asset> {
+    value: &'a mut A,
+    labeled_assets: &'a mut HashMap<CowArc<'static, str>, LabeledAsset>,
+}
+
+impl<'a, A: Asset> Deref for TransformedSubAsset<'a, A> {
+    type Target = A;
+    fn deref(&self) -> &Self::Target {
+        self.value
+    }
+}
+
+impl<'a, A: Asset> DerefMut for TransformedSubAsset<'a, A> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.value
+    }
+}
+
+impl<'a, A: Asset> TransformedSubAsset<'a, A> {
+    /// Creates a new [`TransformedSubAsset`] from `asset` if its internal value matches `A`.
+    pub fn from_loaded(asset: &'a mut ErasedLoadedAsset) -> Option<Self> {
+        let value = asset.value.downcast_mut::<A>()?;
+        Some(TransformedSubAsset {
+            value,
+            labeled_assets: &mut asset.labeled_assets,
+        })
+    }
+    /// Retrieves the value of this asset.
+    #[inline]
+    pub fn get(&self) -> &A {
+        self.value
+    }
+    /// Mutably retrieves the value of this asset.
+    #[inline]
+    pub fn get_mut(&mut self) -> &mut A {
+        self.value
+    }
+    /// Returns the labeled asset, if it exists and matches this type.
+    pub fn get_labeled<B: Asset, Q>(&mut self, label: &Q) -> Option<TransformedSubAsset<B>>
+    where
+        CowArc<'static, str>: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        let labeled = self.labeled_assets.get_mut(label)?;
+        let value = labeled.asset.value.downcast_mut::<B>()?;
+        Some(TransformedSubAsset {
+            value,
+            labeled_assets: &mut labeled.asset.labeled_assets,
+        })
+    }
+    /// Returns the type-erased labeled asset, if it exists and matches this type.
+    pub fn get_erased_labeled<Q>(&self, label: &Q) -> Option<&ErasedLoadedAsset>
+    where
+        CowArc<'static, str>: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        let labeled = self.labeled_assets.get(label)?;
+        Some(&labeled.asset)
+    }
+    /// Returns the [`UntypedHandle`] of the labeled asset with the provided 'label', if it exists.
+    pub fn get_untyped_handle<Q>(&self, label: &Q) -> Option<UntypedHandle>
+    where
+        CowArc<'static, str>: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        let labeled = self.labeled_assets.get(label)?;
+        Some(labeled.handle.clone())
+    }
+    /// Returns the [`Handle`] of the labeled asset with the provided 'label', if it exists and is an asset of type `B`
+    pub fn get_handle<Q, B: Asset>(&self, label: &Q) -> Option<Handle<B>>
+    where
+        CowArc<'static, str>: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        let labeled = self.labeled_assets.get(label)?;
+        if let Ok(handle) = labeled.handle.clone().try_typed::<B>() {
+            return Some(handle);
+        }
+        None
+    }
+    /// Adds `asset` as a labeled sub asset using `label` and `handle`
+    pub fn insert_labeled(
+        &mut self,
+        label: impl Into<CowArc<'static, str>>,
+        handle: impl Into<UntypedHandle>,
+        asset: impl Into<ErasedLoadedAsset>,
+    ) {
+        let labeled = LabeledAsset {
+            asset: asset.into(),
+            handle: handle.into(),
+        };
+        self.labeled_assets.insert(label.into(), labeled);
+    }
+    /// Iterate over all labels for "labeled assets" in the loaded asset
+    pub fn iter_labels(&self) -> impl Iterator<Item = &str> {
+        self.labeled_assets.keys().map(|s| &**s)
+    }
 }

examples/asset/processing/asset_processing.rs

@@ -7,7 +7,7 @@ use bevy::{
         processor::LoadTransformAndSave,
         ron,
         saver::{AssetSaver, SavedAsset},
-        transformer::AssetTransformer,
+        transformer::{AssetTransformer, TransformedAsset},
         AssetLoader, AsyncReadExt, AsyncWriteExt, LoadContext,
     },
     prelude::*,
@@ -186,12 +186,12 @@ impl AssetTransformer for CoolTextTransformer {
 
     fn transform<'a>(
         &'a self,
-        asset: Self::AssetInput,
+        mut asset: TransformedAsset<Self::AssetInput>,
         settings: &'a Self::Settings,
-    ) -> Result<Self::AssetOutput, Box<dyn std::error::Error + Send + Sync + 'static>> {
-        Ok(CoolText {
-            text: format!("{}{}", asset.text, settings.appended),
-            dependencies: asset.dependencies.clone(),
+    ) -> BoxedFuture<'a, Result<TransformedAsset<Self::AssetOutput>, Self::Error>> {
+        Box::pin(async move {
+            asset.text = format!("{}{}", asset.text, settings.appended);
+            Ok(asset)
         })
     }
 }