Skip to main content

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