bevy/examples/asset/asset_decompression.rs
Joseph 5876352206
Optimize common usages of AssetReader (#14082)
# Objective

The `AssetReader` trait allows customizing the behavior of fetching
bytes for an `AssetPath`, and expects implementors to return `dyn
AsyncRead + AsyncSeek`. This gives implementors of `AssetLoader` great
flexibility to tightly integrate their asset loading behavior with the
asynchronous task system.

However, almost all implementors of `AssetLoader` don't use the async
functionality at all, and just call `AsyncReadExt::read_to_end(&mut
Vec<u8>)`. This is incredibly inefficient, as this method repeatedly
calls `poll_read` on the trait object, filling the vector 32 bytes at a
time. At my work we have assets that are hundreds of megabytes which
makes this a meaningful overhead.

## Solution

Turn the `Reader` type alias into an actual trait, with a provided
method `read_to_end`. This provided method should be more efficient than
the existing extension method, as the compiler will know the underlying
type of `Reader` when generating this function, which removes the
repeated dynamic dispatches and allows the compiler to make further
optimizations after inlining. Individual implementors are able to
override the provided implementation -- for simple asset readers that
just copy bytes from one buffer to another, this allows removing a large
amount of overhead from the provided implementation.

Now that `Reader` is an actual trait, I also improved the ergonomics for
implementing `AssetReader`. Currently, implementors are expected to box
their reader and return it as a trait object, which adds unnecessary
boilerplate to implementations. This PR changes that trait method to
return a pseudo trait alias, which allows implementors to return `impl
Reader` instead of `Box<dyn Reader>`. Now, the boilerplate for boxing
occurs in `ErasedAssetReader`.

## Testing

I made identical changes to my company's fork of bevy. Our app, which
makes heavy use of `read_to_end` for asset loading, still worked
properly after this. I am not aware if we have a more systematic way of
testing asset loading for correctness.

---

## Migration Guide

The trait method `bevy_asset::io::AssetReader::read` (and `read_meta`)
now return an opaque type instead of a boxed trait object. Implementors
of these methods should change the type signatures appropriately

```rust
impl AssetReader for MyReader {
    // Before
    async fn read<'a>(&'a self, path: &'a Path) -> Result<Box<Reader<'a>>, AssetReaderError> {
        let reader = // construct a reader
        Box::new(reader) as Box<Reader<'a>>
    }

    // After
    async fn read<'a>(&'a self, path: &'a Path) -> Result<impl Reader + 'a, AssetReaderError> {
        // create a reader
    }
}
```

`bevy::asset::io::Reader` is now a trait, rather than a type alias for a
trait object. Implementors of `AssetLoader::load` will need to adjust
the method signature accordingly

```rust
impl AssetLoader for MyLoader {
    async fn load<'a>(
        &'a self,
        // Before:
        reader: &'a mut bevy::asset::io::Reader,
        // After:
        reader: &'a mut dyn bevy::asset::io::Reader,
        _: &'a Self::Settings,
        load_context: &'a mut LoadContext<'_>,
    ) -> Result<Self::Asset, Self::Error> {
}
```

Additionally, implementors of `AssetReader` that return a type
implementing `futures_io::AsyncRead` and `AsyncSeek` might need to
explicitly implement `bevy::asset::io::Reader` for that type.

```rust
impl bevy::asset::io::Reader for MyAsyncReadAndSeek {}
```
2024-07-01 19:59:42 +00:00

138 lines
3.8 KiB
Rust

//! Implements loader for a Gzip compressed asset.
use bevy::{
asset::{
io::{Reader, VecReader},
AssetLoader, ErasedLoadedAsset, LoadContext, LoadDirectError,
},
prelude::*,
reflect::TypePath,
};
use flate2::read::GzDecoder;
use std::io::prelude::*;
use std::marker::PhantomData;
use thiserror::Error;
#[derive(Asset, TypePath)]
struct GzAsset {
uncompressed: ErasedLoadedAsset,
}
#[derive(Default)]
struct GzAssetLoader;
/// Possible errors that can be produced by [`GzAssetLoader`]
#[non_exhaustive]
#[derive(Debug, Error)]
enum GzAssetLoaderError {
/// An [IO](std::io) Error
#[error("Could not load asset: {0}")]
Io(#[from] std::io::Error),
/// An error caused when the asset path cannot be used to determine the uncompressed asset type.
#[error("Could not determine file path of uncompressed asset")]
IndeterminateFilePath,
/// An error caused by the internal asset loader.
#[error("Could not load contained asset: {0}")]
LoadDirectError(#[from] LoadDirectError),
}
impl AssetLoader for GzAssetLoader {
type Asset = GzAsset;
type Settings = ();
type Error = GzAssetLoaderError;
async fn load<'a>(
&'a self,
reader: &'a mut dyn Reader,
_settings: &'a (),
load_context: &'a mut LoadContext<'_>,
) -> Result<Self::Asset, Self::Error> {
let compressed_path = load_context.path();
let file_name = compressed_path
.file_name()
.ok_or(GzAssetLoaderError::IndeterminateFilePath)?
.to_string_lossy();
let uncompressed_file_name = file_name
.strip_suffix(".gz")
.ok_or(GzAssetLoaderError::IndeterminateFilePath)?;
let contained_path = compressed_path.join(uncompressed_file_name);
let mut bytes_compressed = Vec::new();
reader.read_to_end(&mut bytes_compressed).await?;
let mut decoder = GzDecoder::new(bytes_compressed.as_slice());
let mut bytes_uncompressed = Vec::new();
decoder.read_to_end(&mut bytes_uncompressed)?;
// Now that we have decompressed the asset, let's pass it back to the
// context to continue loading
let mut reader = VecReader::new(bytes_uncompressed);
let uncompressed = load_context
.loader()
.direct()
.with_reader(&mut reader)
.untyped()
.load(contained_path)
.await?;
Ok(GzAsset { uncompressed })
}
fn extensions(&self) -> &[&str] {
&["gz"]
}
}
#[derive(Component, Default)]
struct Compressed<T> {
compressed: Handle<GzAsset>,
_phantom: PhantomData<T>,
}
fn main() {
App::new()
.add_plugins(DefaultPlugins)
.init_asset::<GzAsset>()
.init_asset_loader::<GzAssetLoader>()
.add_systems(Startup, setup)
.add_systems(Update, decompress::<Image>)
.run();
}
fn setup(mut commands: Commands, asset_server: Res<AssetServer>) {
commands.spawn(Camera2dBundle::default());
commands.spawn((
Compressed::<Image> {
compressed: asset_server.load("data/compressed_image.png.gz"),
..default()
},
Sprite::default(),
TransformBundle::default(),
VisibilityBundle::default(),
));
}
fn decompress<A: Asset>(
mut commands: Commands,
asset_server: Res<AssetServer>,
mut compressed_assets: ResMut<Assets<GzAsset>>,
query: Query<(Entity, &Compressed<A>)>,
) {
for (entity, Compressed { compressed, .. }) in query.iter() {
let Some(GzAsset { uncompressed }) = compressed_assets.remove(compressed) else {
continue;
};
let uncompressed = uncompressed.take::<A>().unwrap();
commands
.entity(entity)
.remove::<Compressed<A>>()
.insert(asset_server.add(uncompressed));
}
}