Fix AssetTransformer breaking LabeledAssets (#11626)

# Objective - `AssetTransformer` provides an input asset, and output an asset, but provides no access to the `LabeledAsset`'s created by the `AssetLoader`. Labeled sub assets are an extremely important piece of many assets, Gltf in particular, and without them the amount of transformation on an asset is limited. In order for `AssetTransformer`'s to be useful, they need to have access to these sub assets. - LabeledAsset's loaded by `AssetLoader`s are provided to `AssetSaver`s in the `LoadAndSave` process, but the `LoadTransformAndSave` process drops these values in the transform stage, and so `AssetSaver` is given none. - Fixes #11606 Ideally the AssetTransformer should not ignore labeled sub assets, and they should be kept at least for the AssetSaver ## Solution - I created a new struct similar to `SavedAsset` named `TransformedAsset` which holds the input asset, and the HashMap of `LabeledAsset`s. The transform function now takes as input a `TransformedAsset`, and returns a `TransformedAsset::<AssetOutput>`. This gives the transform function access to the labeled sub assets created by the `AssetLoader`. - I also created `TransformedSubAsset` which holds mutable references to a sub asset and that sub assets HashMap of `LabeledAsset`s. This allows you to travers the Tree of `LabeledAsset`s by reference relatively easily. - The `LoadTransformAndSave` processor was then reworked to use the new structs, stopping the `LabeledAsset`s from being dropped. --- ## Changelog - Created TransformedAsset struct and TransformedSubAsset struct. - Changed `get_untyped_handle` to return a `UntypedHandle` directly rather than a reference and added `get_handle` as a typed variant in SavedAsset and TransformedAsset - Added `SavedAsset::from_transformed` as a constructor from a `TransformedAsset` - Switched LoadTransformAndSave process code to work with new `TransformedAsset` type - Added a `ProcessError` for `AssetTransformer` in process.rs - Switched `AssetTransformer::transform` to use `TransformedAsset` as input and output. - Switched `AssetTransformer` to use a `BoxedFuture` like `AssetLoader` and `AssetSaver` to allow for async transformation code. - Updated AssetTransformer example to use new structure.
2024-02-02 09:57:31 -05:00 · 2024-02-02 09:57:31 -05:00 · 8866c61161
commit 8866c61161
parent 176223b406
4 changed files with 272 additions and 25 deletions
--- a/crates/bevy_asset/src/processor/process.rs
+++ b/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
--- a/crates/bevy_asset/src/saver.rs
+++ b/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)
+        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
    }

    /// Iterate over all labels for "labeled assets" in the loaded asset
--- a/crates/bevy_asset/src/transformer.rs
+++ b/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)
+    }
 }
--- a/examples/asset/processing/asset_processing.rs
+++ b/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)
        })
    }
 }