bevy_render/
erased_render_asset.rs

1use crate::{
2    render_resource::AsBindGroupError, ExtractSchedule, MainWorld, Render, RenderApp,
3    RenderSystems, Res,
4};
5use bevy_app::{App, Plugin, SubApp};
6use bevy_asset::RenderAssetUsages;
7use bevy_asset::{Asset, AssetEvent, AssetId, Assets, UntypedAssetId};
8use bevy_ecs::{
9    prelude::{Commands, IntoScheduleConfigs, MessageReader, ResMut, Resource},
10    schedule::{ScheduleConfigs, SystemSet},
11    system::{ScheduleSystem, StaticSystemParam, SystemParam, SystemParamItem, SystemState},
12    world::{FromWorld, Mut},
13};
14use bevy_platform::collections::{HashMap, HashSet};
15use bevy_render::render_asset::RenderAssetBytesPerFrameLimiter;
16use core::marker::PhantomData;
17use thiserror::Error;
18use tracing::{debug, error};
19
20#[derive(Debug, Error)]
21pub enum PrepareAssetError<E: Send + Sync + 'static> {
22    #[error("Failed to prepare asset")]
23    RetryNextUpdate(E),
24    #[error("Failed to build bind group: {0}")]
25    AsBindGroupError(AsBindGroupError),
26}
27
28/// The system set during which we extract modified assets to the render world.
29#[derive(SystemSet, Clone, PartialEq, Eq, Debug, Hash)]
30pub struct AssetExtractionSystems;
31
32/// Deprecated alias for [`AssetExtractionSystems`].
33#[deprecated(since = "0.17.0", note = "Renamed to `AssetExtractionSystems`.")]
34pub type ExtractAssetsSet = AssetExtractionSystems;
35
36/// Describes how an asset gets extracted and prepared for rendering.
37///
38/// In the [`ExtractSchedule`] step the [`ErasedRenderAsset::SourceAsset`] is transferred
39/// from the "main world" into the "render world".
40///
41/// After that in the [`RenderSystems::PrepareAssets`] step the extracted asset
42/// is transformed into its GPU-representation of type [`ErasedRenderAsset`].
43pub trait ErasedRenderAsset: Send + Sync + 'static {
44    /// The representation of the asset in the "main world".
45    type SourceAsset: Asset + Clone;
46    /// The target representation of the asset in the "render world".
47    type ErasedAsset: Send + Sync + 'static + Sized;
48
49    /// Specifies all ECS data required by [`ErasedRenderAsset::prepare_asset`].
50    ///
51    /// For convenience use the [`lifetimeless`](bevy_ecs::system::lifetimeless) [`SystemParam`].
52    type Param: SystemParam;
53
54    /// Whether or not to unload the asset after extracting it to the render world.
55    #[inline]
56    fn asset_usage(_source_asset: &Self::SourceAsset) -> RenderAssetUsages {
57        RenderAssetUsages::default()
58    }
59
60    /// Size of the data the asset will upload to the gpu. Specifying a return value
61    /// will allow the asset to be throttled via [`RenderAssetBytesPerFrameLimiter`].
62    #[inline]
63    #[expect(
64        unused_variables,
65        reason = "The parameters here are intentionally unused by the default implementation; however, putting underscores here will result in the underscores being copied by rust-analyzer's tab completion."
66    )]
67    fn byte_len(erased_asset: &Self::SourceAsset) -> Option<usize> {
68        None
69    }
70
71    /// Prepares the [`ErasedRenderAsset::SourceAsset`] for the GPU by transforming it into a [`ErasedRenderAsset`].
72    ///
73    /// ECS data may be accessed via `param`.
74    fn prepare_asset(
75        source_asset: Self::SourceAsset,
76        asset_id: AssetId<Self::SourceAsset>,
77        param: &mut SystemParamItem<Self::Param>,
78    ) -> Result<Self::ErasedAsset, PrepareAssetError<Self::SourceAsset>>;
79
80    /// Called whenever the [`ErasedRenderAsset::SourceAsset`] has been removed.
81    ///
82    /// You can implement this method if you need to access ECS data (via
83    /// `_param`) in order to perform cleanup tasks when the asset is removed.
84    ///
85    /// The default implementation does nothing.
86    fn unload_asset(
87        _source_asset: AssetId<Self::SourceAsset>,
88        _param: &mut SystemParamItem<Self::Param>,
89    ) {
90    }
91}
92
93/// This plugin extracts the changed assets from the "app world" into the "render world"
94/// and prepares them for the GPU. They can then be accessed from the [`ErasedRenderAssets`] resource.
95///
96/// Therefore it sets up the [`ExtractSchedule`] and
97/// [`RenderSystems::PrepareAssets`] steps for the specified [`ErasedRenderAsset`].
98///
99/// The `AFTER` generic parameter can be used to specify that `A::prepare_asset` should not be run until
100/// `prepare_assets::<AFTER>` has completed. This allows the `prepare_asset` function to depend on another
101/// prepared [`ErasedRenderAsset`], for example `Mesh::prepare_asset` relies on `ErasedRenderAssets::<GpuImage>` for morph
102/// targets, so the plugin is created as `ErasedRenderAssetPlugin::<RenderMesh, GpuImage>::default()`.
103pub struct ErasedRenderAssetPlugin<
104    A: ErasedRenderAsset,
105    AFTER: ErasedRenderAssetDependency + 'static = (),
106> {
107    phantom: PhantomData<fn() -> (A, AFTER)>,
108}
109
110impl<A: ErasedRenderAsset, AFTER: ErasedRenderAssetDependency + 'static> Default
111    for ErasedRenderAssetPlugin<A, AFTER>
112{
113    fn default() -> Self {
114        Self {
115            phantom: Default::default(),
116        }
117    }
118}
119
120impl<A: ErasedRenderAsset, AFTER: ErasedRenderAssetDependency + 'static> Plugin
121    for ErasedRenderAssetPlugin<A, AFTER>
122{
123    fn build(&self, app: &mut App) {
124        app.init_resource::<CachedExtractErasedRenderAssetSystemState<A>>();
125    }
126
127    fn finish(&self, app: &mut App) {
128        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
129            render_app
130                .init_resource::<ExtractedAssets<A>>()
131                .init_resource::<ErasedRenderAssets<A::ErasedAsset>>()
132                .init_resource::<PrepareNextFrameAssets<A>>()
133                .add_systems(
134                    ExtractSchedule,
135                    extract_erased_render_asset::<A>.in_set(AssetExtractionSystems),
136                );
137            AFTER::register_system(
138                render_app,
139                prepare_erased_assets::<A>.in_set(RenderSystems::PrepareAssets),
140            );
141        }
142    }
143}
144
145// helper to allow specifying dependencies between render assets
146pub trait ErasedRenderAssetDependency {
147    fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>);
148}
149
150impl ErasedRenderAssetDependency for () {
151    fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
152        render_app.add_systems(Render, system);
153    }
154}
155
156impl<A: ErasedRenderAsset> ErasedRenderAssetDependency for A {
157    fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
158        render_app.add_systems(Render, system.after(prepare_erased_assets::<A>));
159    }
160}
161
162/// Temporarily stores the extracted and removed assets of the current frame.
163#[derive(Resource)]
164pub struct ExtractedAssets<A: ErasedRenderAsset> {
165    /// The assets extracted this frame.
166    ///
167    /// These are assets that were either added or modified this frame.
168    pub extracted: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
169
170    /// IDs of the assets that were removed this frame.
171    ///
172    /// These assets will not be present in [`ExtractedAssets::extracted`].
173    pub removed: HashSet<AssetId<A::SourceAsset>>,
174
175    /// IDs of the assets that were modified this frame.
176    pub modified: HashSet<AssetId<A::SourceAsset>>,
177
178    /// IDs of the assets that were added this frame.
179    pub added: HashSet<AssetId<A::SourceAsset>>,
180}
181
182impl<A: ErasedRenderAsset> Default for ExtractedAssets<A> {
183    fn default() -> Self {
184        Self {
185            extracted: Default::default(),
186            removed: Default::default(),
187            modified: Default::default(),
188            added: Default::default(),
189        }
190    }
191}
192
193/// Stores all GPU representations ([`ErasedRenderAsset`])
194/// of [`ErasedRenderAsset::SourceAsset`] as long as they exist.
195#[derive(Resource)]
196pub struct ErasedRenderAssets<ERA>(HashMap<UntypedAssetId, ERA>);
197
198impl<ERA> Default for ErasedRenderAssets<ERA> {
199    fn default() -> Self {
200        Self(Default::default())
201    }
202}
203
204impl<ERA> ErasedRenderAssets<ERA> {
205    pub fn get(&self, id: impl Into<UntypedAssetId>) -> Option<&ERA> {
206        self.0.get(&id.into())
207    }
208
209    pub fn get_mut(&mut self, id: impl Into<UntypedAssetId>) -> Option<&mut ERA> {
210        self.0.get_mut(&id.into())
211    }
212
213    pub fn insert(&mut self, id: impl Into<UntypedAssetId>, value: ERA) -> Option<ERA> {
214        self.0.insert(id.into(), value)
215    }
216
217    pub fn remove(&mut self, id: impl Into<UntypedAssetId>) -> Option<ERA> {
218        self.0.remove(&id.into())
219    }
220
221    pub fn iter(&self) -> impl Iterator<Item = (UntypedAssetId, &ERA)> {
222        self.0.iter().map(|(k, v)| (*k, v))
223    }
224
225    pub fn iter_mut(&mut self) -> impl Iterator<Item = (UntypedAssetId, &mut ERA)> {
226        self.0.iter_mut().map(|(k, v)| (*k, v))
227    }
228}
229
230#[derive(Resource)]
231struct CachedExtractErasedRenderAssetSystemState<A: ErasedRenderAsset> {
232    state: SystemState<(
233        MessageReader<'static, 'static, AssetEvent<A::SourceAsset>>,
234        ResMut<'static, Assets<A::SourceAsset>>,
235    )>,
236}
237
238impl<A: ErasedRenderAsset> FromWorld for CachedExtractErasedRenderAssetSystemState<A> {
239    fn from_world(world: &mut bevy_ecs::world::World) -> Self {
240        Self {
241            state: SystemState::new(world),
242        }
243    }
244}
245
246/// This system extracts all created or modified assets of the corresponding [`ErasedRenderAsset::SourceAsset`] type
247/// into the "render world".
248pub(crate) fn extract_erased_render_asset<A: ErasedRenderAsset>(
249    mut commands: Commands,
250    mut main_world: ResMut<MainWorld>,
251) {
252    main_world.resource_scope(
253        |world, mut cached_state: Mut<CachedExtractErasedRenderAssetSystemState<A>>| {
254            let (mut events, mut assets) = cached_state.state.get_mut(world);
255
256            let mut needs_extracting = <HashSet<_>>::default();
257            let mut removed = <HashSet<_>>::default();
258            let mut modified = <HashSet<_>>::default();
259
260            for event in events.read() {
261                #[expect(
262                    clippy::match_same_arms,
263                    reason = "LoadedWithDependencies is marked as a TODO, so it's likely this will no longer lint soon."
264                )]
265                match event {
266                    AssetEvent::Added { id } => {
267                        needs_extracting.insert(*id);
268                    }
269                    AssetEvent::Modified { id } => {
270                        needs_extracting.insert(*id);
271                        modified.insert(*id);
272                    }
273                    AssetEvent::Removed { .. } => {
274                        // We don't care that the asset was removed from Assets<T> in the main world.
275                        // An asset is only removed from ErasedRenderAssets<T> when its last handle is dropped (AssetEvent::Unused).
276                    }
277                    AssetEvent::Unused { id } => {
278                        needs_extracting.remove(id);
279                        modified.remove(id);
280                        removed.insert(*id);
281                    }
282                    AssetEvent::LoadedWithDependencies { .. } => {
283                        // TODO: handle this
284                    }
285                }
286            }
287
288            let mut extracted_assets = Vec::new();
289            let mut added = <HashSet<_>>::default();
290            for id in needs_extracting.drain() {
291                if let Some(asset) = assets.get(id) {
292                    let asset_usage = A::asset_usage(asset);
293                    if asset_usage.contains(RenderAssetUsages::RENDER_WORLD) {
294                        if asset_usage == RenderAssetUsages::RENDER_WORLD {
295                            if let Some(asset) = assets.remove(id) {
296                                extracted_assets.push((id, asset));
297                                added.insert(id);
298                            }
299                        } else {
300                            extracted_assets.push((id, asset.clone()));
301                            added.insert(id);
302                        }
303                    }
304                }
305            }
306
307            commands.insert_resource(ExtractedAssets::<A> {
308                extracted: extracted_assets,
309                removed,
310                modified,
311                added,
312            });
313            cached_state.state.apply(world);
314        },
315    );
316}
317
318// TODO: consider storing inside system?
319/// All assets that should be prepared next frame.
320#[derive(Resource)]
321pub struct PrepareNextFrameAssets<A: ErasedRenderAsset> {
322    assets: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
323}
324
325impl<A: ErasedRenderAsset> Default for PrepareNextFrameAssets<A> {
326    fn default() -> Self {
327        Self {
328            assets: Default::default(),
329        }
330    }
331}
332
333/// This system prepares all assets of the corresponding [`ErasedRenderAsset::SourceAsset`] type
334/// which where extracted this frame for the GPU.
335pub fn prepare_erased_assets<A: ErasedRenderAsset>(
336    mut extracted_assets: ResMut<ExtractedAssets<A>>,
337    mut render_assets: ResMut<ErasedRenderAssets<A::ErasedAsset>>,
338    mut prepare_next_frame: ResMut<PrepareNextFrameAssets<A>>,
339    param: StaticSystemParam<<A as ErasedRenderAsset>::Param>,
340    bpf: Res<RenderAssetBytesPerFrameLimiter>,
341) {
342    let mut wrote_asset_count = 0;
343
344    let mut param = param.into_inner();
345    let queued_assets = core::mem::take(&mut prepare_next_frame.assets);
346    for (id, extracted_asset) in queued_assets {
347        if extracted_assets.removed.contains(&id) || extracted_assets.added.contains(&id) {
348            // skip previous frame's assets that have been removed or updated
349            continue;
350        }
351
352        let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
353            // we could check if available bytes > byte_len here, but we want to make some
354            // forward progress even if the asset is larger than the max bytes per frame.
355            // this way we always write at least one (sized) asset per frame.
356            // in future we could also consider partial asset uploads.
357            if bpf.exhausted() {
358                prepare_next_frame.assets.push((id, extracted_asset));
359                continue;
360            }
361            size
362        } else {
363            0
364        };
365
366        match A::prepare_asset(extracted_asset, id, &mut param) {
367            Ok(prepared_asset) => {
368                render_assets.insert(id, prepared_asset);
369                bpf.write_bytes(write_bytes);
370                wrote_asset_count += 1;
371            }
372            Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
373                prepare_next_frame.assets.push((id, extracted_asset));
374            }
375            Err(PrepareAssetError::AsBindGroupError(e)) => {
376                error!(
377                    "{} Bind group construction failed: {e}",
378                    core::any::type_name::<A>()
379                );
380            }
381        }
382    }
383
384    for removed in extracted_assets.removed.drain() {
385        render_assets.remove(removed);
386        A::unload_asset(removed, &mut param);
387    }
388
389    for (id, extracted_asset) in extracted_assets.extracted.drain(..) {
390        // we remove previous here to ensure that if we are updating the asset then
391        // any users will not see the old asset after a new asset is extracted,
392        // even if the new asset is not yet ready or we are out of bytes to write.
393        render_assets.remove(id);
394
395        let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
396            if bpf.exhausted() {
397                prepare_next_frame.assets.push((id, extracted_asset));
398                continue;
399            }
400            size
401        } else {
402            0
403        };
404
405        match A::prepare_asset(extracted_asset, id, &mut param) {
406            Ok(prepared_asset) => {
407                render_assets.insert(id, prepared_asset);
408                bpf.write_bytes(write_bytes);
409                wrote_asset_count += 1;
410            }
411            Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
412                prepare_next_frame.assets.push((id, extracted_asset));
413            }
414            Err(PrepareAssetError::AsBindGroupError(e)) => {
415                error!(
416                    "{} Bind group construction failed: {e}",
417                    core::any::type_name::<A>()
418                );
419            }
420        }
421    }
422
423    if bpf.exhausted() && !prepare_next_frame.assets.is_empty() {
424        debug!(
425            "{} write budget exhausted with {} assets remaining (wrote {})",
426            core::any::type_name::<A>(),
427            prepare_next_frame.assets.len(),
428            wrote_asset_count
429        );
430    }
431}