use crate::{Extract, ExtractSchedule, Render, RenderApp, RenderSet};
use bevy_app::{App, Plugin};
use bevy_asset::{Asset, AssetEvent, Assets, Handle};
use bevy_derive::{Deref, DerefMut};
use bevy_ecs::{
prelude::*,
system::{StaticSystemParam, SystemParam, SystemParamItem},
};
use bevy_utils::{HashMap, HashSet};
use std::marker::PhantomData;
pub enum PrepareAssetError<E: Send + Sync + 'static> {
RetryNextUpdate(E),
}
pub trait RenderAsset: Asset {
type ExtractedAsset: Send + Sync + 'static;
type PreparedAsset: Send + Sync + 'static;
type Param: SystemParam;
fn extract_asset(&self) -> Self::ExtractedAsset;
fn prepare_asset(
extracted_asset: Self::ExtractedAsset,
param: &mut SystemParamItem<Self::Param>,
) -> Result<Self::PreparedAsset, PrepareAssetError<Self::ExtractedAsset>>;
}
#[derive(Clone, Hash, Debug, Default, PartialEq, Eq, SystemSet)]
pub enum PrepareAssetSet {
PreAssetPrepare,
#[default]
AssetPrepare,
PostAssetPrepare,
}
pub struct RenderAssetPlugin<A: RenderAsset> {
prepare_asset_set: PrepareAssetSet,
phantom: PhantomData<fn() -> A>,
}
impl<A: RenderAsset> RenderAssetPlugin<A> {
pub fn with_prepare_asset_set(prepare_asset_set: PrepareAssetSet) -> Self {
Self {
prepare_asset_set,
phantom: PhantomData,
}
}
}
impl<A: RenderAsset> Default for RenderAssetPlugin<A> {
fn default() -> Self {
Self {
prepare_asset_set: Default::default(),
phantom: PhantomData,
}
}
}
impl<A: RenderAsset> Plugin for RenderAssetPlugin<A> {
fn build(&self, app: &mut App) {
if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
render_app
.init_resource::<ExtractedAssets<A>>()
.init_resource::<RenderAssets<A>>()
.init_resource::<PrepareNextFrameAssets<A>>()
.add_systems(ExtractSchedule, extract_render_asset::<A>)
.configure_sets(
Render,
(
PrepareAssetSet::PreAssetPrepare,
PrepareAssetSet::AssetPrepare,
PrepareAssetSet::PostAssetPrepare,
)
.chain()
.in_set(RenderSet::Prepare),
)
.add_systems(
Render,
prepare_assets::<A>.in_set(self.prepare_asset_set.clone()),
);
}
}
}
#[derive(Resource)]
pub struct ExtractedAssets<A: RenderAsset> {
extracted: Vec<(Handle<A>, A::ExtractedAsset)>,
removed: Vec<Handle<A>>,
}
impl<A: RenderAsset> Default for ExtractedAssets<A> {
fn default() -> Self {
Self {
extracted: Default::default(),
removed: Default::default(),
}
}
}
#[derive(Resource, Deref, DerefMut)]
pub struct RenderAssets<A: RenderAsset>(HashMap<Handle<A>, A::PreparedAsset>);
impl<A: RenderAsset> Default for RenderAssets<A> {
fn default() -> Self {
Self(Default::default())
}
}
fn extract_render_asset<A: RenderAsset>(
mut commands: Commands,
mut events: Extract<EventReader<AssetEvent<A>>>,
assets: Extract<Res<Assets<A>>>,
) {
let mut changed_assets = HashSet::default();
let mut removed = Vec::new();
for event in events.iter() {
match event {
AssetEvent::Created { handle } | AssetEvent::Modified { handle } => {
changed_assets.insert(handle.clone_weak());
}
AssetEvent::Removed { handle } => {
changed_assets.remove(handle);
removed.push(handle.clone_weak());
}
}
}
let mut extracted_assets = Vec::new();
for handle in changed_assets.drain() {
if let Some(asset) = assets.get(&handle) {
extracted_assets.push((handle, asset.extract_asset()));
}
}
commands.insert_resource(ExtractedAssets {
extracted: extracted_assets,
removed,
});
}
#[derive(Resource)]
pub struct PrepareNextFrameAssets<A: RenderAsset> {
assets: Vec<(Handle<A>, A::ExtractedAsset)>,
}
impl<A: RenderAsset> Default for PrepareNextFrameAssets<A> {
fn default() -> Self {
Self {
assets: Default::default(),
}
}
}
pub fn prepare_assets<R: RenderAsset>(
mut extracted_assets: ResMut<ExtractedAssets<R>>,
mut render_assets: ResMut<RenderAssets<R>>,
mut prepare_next_frame: ResMut<PrepareNextFrameAssets<R>>,
param: StaticSystemParam<<R as RenderAsset>::Param>,
) {
let mut param = param.into_inner();
let queued_assets = std::mem::take(&mut prepare_next_frame.assets);
for (handle, extracted_asset) in queued_assets {
match R::prepare_asset(extracted_asset, &mut param) {
Ok(prepared_asset) => {
render_assets.insert(handle, prepared_asset);
}
Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
prepare_next_frame.assets.push((handle, extracted_asset));
}
}
}
for removed in std::mem::take(&mut extracted_assets.removed) {
render_assets.remove(&removed);
}
for (handle, extracted_asset) in std::mem::take(&mut extracted_assets.extracted) {
match R::prepare_asset(extracted_asset, &mut param) {
Ok(prepared_asset) => {
render_assets.insert(handle, prepared_asset);
}
Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
prepare_next_frame.assets.push((handle, extracted_asset));
}
}
}
}