wde_renderer/assets/
asset.rs

1use wde_logger::prelude::*;
2
3use bevy::{
4    app::{App, Plugin},
5    ecs::{
6        system::{StaticSystemParam, SystemParam, SystemParamItem, SystemState},
7        world
8    },
9    platform::collections::{HashMap, HashSet},
10    prelude::*
11};
12use thiserror::Error;
13
14use crate::core::{Extract, MainWorld, Render, RenderApp, RenderSet};
15
16#[derive(Debug, Error)]
17pub enum PrepareAssetError<E: Send + Sync + 'static> {
18    #[error("Failed to prepare asset. Retry next frame: {0}.")]
19    RetryNextUpdate(E),
20    #[error("Fatal error preparing asset: {0}.")]
21    Fatal(String)
22}
23/// Trait that describes a GPU asset extracted from a CPU asset that implements the [bevy::prelude::Asset] trait.
24/// The GPU asset is prepared from the CPU asset using the render-world system params, and can fail with a retry or fatal error.
25pub trait RenderAsset: Send + Sync + 'static + Sized {
26    type SourceAsset: Asset + Clone;
27    type Params: SystemParam;
28
29    /// Prepare the GPU asset from the CPU source [bevy::prelude::Asset] using the render-world system params.
30    fn prepare(
31        asset: Self::SourceAsset,
32        params: &mut SystemParamItem<Self::Params>
33    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>>;
34    fn label(&self) -> &str {
35        std::any::type_name::<Self>()
36    }
37}
38
39/// Stores all assets of a given GPU [RenderAsset] type, indexed by the ID of their source CPU asset.
40#[derive(Resource)]
41pub struct RenderAssets<A: RenderAsset>(HashMap<AssetId<A::SourceAsset>, A>);
42impl<A: RenderAsset> Default for RenderAssets<A> {
43    fn default() -> Self {
44        Self(Default::default())
45    }
46}
47impl<A: RenderAsset> RenderAssets<A> {
48    pub fn get(&self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<&A> {
49        self.0.get(&id.into())
50    }
51    pub fn get_mut(&mut self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<&mut A> {
52        self.0.get_mut(&id.into())
53    }
54    pub fn insert(&mut self, id: impl Into<AssetId<A::SourceAsset>>, value: A) -> Option<A> {
55        self.0.insert(id.into(), value)
56    }
57    pub fn remove(&mut self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<A> {
58        self.0.remove(&id.into())
59    }
60    pub fn iter(&self) -> impl Iterator<Item = (&AssetId<A::SourceAsset>, &A)> {
61        self.0.iter()
62    }
63    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&AssetId<A::SourceAsset>, &mut A)> {
64        self.0.iter_mut()
65    }
66}
67
68/// Plugin that adds the systems and resources to extract, prepare and store a given type of GPU [RenderAsset] from their source CPU [bevy::prelude::Asset].
69/// To use, simply add `RenderAssetsPlugin::<YourGpuAssetType>::default()` to your app, and make sure to implement the [RenderAsset] trait for your GPU asset type.
70pub struct RenderAssetsPlugin<A: RenderAsset> {
71    _phantom: std::marker::PhantomData<fn() -> A>
72}
73impl<A: RenderAsset> Default for RenderAssetsPlugin<A> {
74    fn default() -> Self {
75        Self {
76            _phantom: Default::default()
77        }
78    }
79}
80impl<A: RenderAsset> Plugin for RenderAssetsPlugin<A> {
81    fn build(&self, app: &mut App) {
82        // Create the cached for extracting assets from the main world
83        app.init_resource::<CachedExtractAssetsState<A>>();
84
85        // Add the extract system to the renderer app
86        let renderer_app = app.get_sub_app_mut(RenderApp).unwrap();
87        renderer_app
88            .init_resource::<PrepareNextFrameAssets<A>>()
89            .init_resource::<ExtractedAssets<A>>()
90            .init_resource::<RenderAssets<A>>()
91            .add_systems(Extract, extract_render_assets::<A>);
92
93        // Add the prepare system to the renderer app
94        renderer_app.add_systems(Render, prepare_assets::<A>.in_set(RenderSet::Prepare));
95    }
96}
97
98/// Stores the list of assets extracted from the main world AssetServer for the current frame, with their IDs and added/removed status.
99#[allow(clippy::type_complexity)]
100#[derive(Resource)]
101struct CachedExtractAssetsState<A: RenderAsset> {
102    state: SystemState<(
103        MessageReader<'static, 'static, AssetEvent<A::SourceAsset>>,
104        ResMut<'static, Assets<A::SourceAsset>>
105    )>
106}
107impl<A: RenderAsset> FromWorld for CachedExtractAssetsState<A> {
108    fn from_world(world: &mut world::World) -> Self {
109        Self {
110            state: SystemState::new(world)
111        }
112    }
113}
114
115/// Resource that stores the assets that failed to prepare in the previous frame and should be retried in the next frame.
116#[derive(Resource)]
117struct PrepareNextFrameAssets<A: RenderAsset> {
118    assets: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>
119}
120impl<A: RenderAsset> Default for PrepareNextFrameAssets<A> {
121    fn default() -> Self {
122        Self {
123            assets: Default::default()
124        }
125    }
126}
127
128/// Resource that stores the extracted assets from the main world AssetServer for the current frame, with their IDs and added/removed status.
129#[derive(Resource)]
130struct ExtractedAssets<A: RenderAsset> {
131    /// List of IDs of the assets added this frame.
132    pub added: HashSet<AssetId<A::SourceAsset>>,
133    /// List of IDs of the assets removed this frame.
134    pub removed: HashSet<AssetId<A::SourceAsset>>,
135    /// The pair (id, CPU asset) of the added assets extracted this frame.
136    pub extracted: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>
137}
138impl<A: RenderAsset> Default for ExtractedAssets<A> {
139    fn default() -> Self {
140        Self {
141            extracted: Default::default(),
142            removed: Default::default(),
143            added: Default::default()
144        }
145    }
146}
147
148/// Extract the modified assets instructions from the main world AssetServer and load them to the renderer AssetServer.
149fn extract_render_assets<A: RenderAsset>(
150    mut commands: Commands,
151    mut main_world: ResMut<MainWorld>
152) {
153    main_world.resource_scope(
154        |main_world, mut cached_state: Mut<CachedExtractAssetsState<A>>| {
155            let (mut events, mut assets) = cached_state.state.get_mut(main_world);
156            let mut changed_assets: HashSet<AssetId<<A as RenderAsset>::SourceAsset>> =
157                HashSet::default();
158            let mut removed = HashSet::default();
159
160            // Read all asset events and track the changed assets by their ID
161            for event in events.read() {
162                match event {
163                    AssetEvent::Added { id } | AssetEvent::Modified { id } => {
164                        changed_assets.insert(*id);
165                    }
166                    AssetEvent::Unused { id } => {
167                        changed_assets.remove(id);
168                        removed.insert(*id);
169                    }
170                    AssetEvent::Removed { .. } => {}
171                    AssetEvent::LoadedWithDependencies { .. } => {}
172                }
173            }
174
175            // Add the changed assets to the extracted assets list
176            let mut extracted_assets = Vec::new();
177            let mut added = HashSet::new();
178            for id in changed_assets.drain() {
179                // Remove the asset from the main world AssetServer to avoid it being used by other systems while we prepare it for the GPU, and add it to the extracted assets list if it was present
180                if let Some(asset) = assets.remove(id) {
181                    extracted_assets.push((id, asset));
182                    added.insert(id);
183                }
184            }
185            commands.insert_resource(ExtractedAssets::<A> {
186                extracted: extracted_assets,
187                removed,
188                added
189            });
190
191            // Apply all queued asset events
192            cached_state.state.apply(main_world);
193        }
194    );
195}
196
197/// Load and unload the assets from the renderer based on the extracted assets.
198fn prepare_assets<A: RenderAsset>(
199    mut extracted_assets: ResMut<ExtractedAssets<A>>,
200    mut render_assets: ResMut<RenderAssets<A>>,
201    mut prepare_next_frame: ResMut<PrepareNextFrameAssets<A>>,
202    params: StaticSystemParam<<A as RenderAsset>::Params>
203) {
204    let mut params = params.into_inner();
205    let queued_assets = std::mem::take(&mut prepare_next_frame.assets);
206
207    // Initialize the render assets from the previous frame that have not been finalized yet
208    for (id, extracted_asset) in queued_assets {
209        // Skip previous frame's assets removed or updated
210        if extracted_assets.removed.contains(&id) || extracted_assets.added.contains(&id) {
211            continue;
212        }
213
214        // Load the asset to the GPU from the CPU
215        match A::prepare(extracted_asset, &mut params) {
216            Ok(prepared_asset) => {
217                // Add the asset to the render world
218                render_assets.insert(id, prepared_asset);
219            }
220            Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
221                // Try again next frame
222                prepare_next_frame.assets.push((id, extracted_asset));
223            }
224            Err(PrepareAssetError::Fatal(error)) => {
225                // Skip the asset
226                error!("Fatal error preparing asset of id {}: {:?}.", id, error);
227                extracted_assets.removed.insert(id);
228            }
229        }
230    }
231
232    // Remove assets
233    for removed in extracted_assets.removed.drain() {
234        let label = match render_assets.get(removed) {
235            Some(asset) => asset.label(),
236            None => "(asset not loaded)"
237        };
238        debug!(
239            "Removing asset {} of type {}.",
240            label,
241            std::any::type_name::<A::SourceAsset>()
242        );
243        render_assets.remove(removed);
244    }
245
246    // Update changed assets
247    for (id, extracted_asset) in extracted_assets.extracted.drain(..) {
248        render_assets.remove(id);
249
250        // Load the asset to the GPU from the CPU
251        match A::prepare(extracted_asset, &mut params) {
252            Ok(prepared_asset) => {
253                // Add the asset to the render world
254                render_assets.insert(id, prepared_asset);
255            }
256            Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
257                // Try again next frame
258                prepare_next_frame.assets.push((id, extracted_asset));
259            }
260            Err(PrepareAssetError::Fatal(error)) => {
261                // Skip the asset
262                error!("Fatal error preparing asset of id {}: {:?}", id, error);
263            }
264        }
265    }
266}