Skip to main content

bevy_render/
render_asset.rs

1use crate::{
2    render_resource::AsBindGroupError, Extract, ExtractSchedule, MainWorld, Render, RenderApp,
3    RenderStartup, RenderSystems, Res,
4};
5use bevy_app::{App, Plugin, SubApp};
6use bevy_asset::{Asset, AssetEvent, AssetId, Assets, RenderAssetUsages};
7use bevy_ecs::{
8    prelude::{Commands, IntoScheduleConfigs, Local, MessageReader, ResMut, Resource},
9    schedule::{ScheduleConfigs, SystemSet},
10    system::{ScheduleSystem, StaticSystemParam, SystemParam, SystemParamItem, SystemState},
11    world::{FromWorld, Mut},
12};
13use bevy_log::{debug, error};
14use bevy_platform::collections::{HashMap, HashSet};
15use core::marker::PhantomData;
16use core::sync::atomic::{AtomicUsize, Ordering};
17use thiserror::Error;
18
19#[derive(Debug, Error)]
20pub enum PrepareAssetError<E: Send + Sync + 'static> {
21    #[error("Failed to prepare asset")]
22    RetryNextUpdate(E),
23    #[error("Failed to build bind group: {0}")]
24    AsBindGroupError(AsBindGroupError),
25}
26
27/// The system set during which we extract modified assets to the render world.
28#[derive(SystemSet, Clone, PartialEq, Eq, Debug, Hash)]
29pub struct AssetExtractionSystems;
30
31/// Error returned when an asset due for extraction has already been extracted
32#[derive(Debug, Error)]
33pub enum AssetExtractionError {
34    #[error("The asset has already been extracted")]
35    AlreadyExtracted,
36    #[error("The asset type does not support extraction. To clone the asset to the renderworld, use `RenderAssetUsages::default()`")]
37    NoExtractionImplementation,
38}
39
40/// Describes how an asset gets extracted and prepared for rendering.
41///
42/// In the [`ExtractSchedule`] step the [`RenderAsset::SourceAsset`] is transferred
43/// from the "main world" into the "render world".
44///
45/// After that in the [`RenderSystems::PrepareAssets`] step the extracted asset
46/// is transformed into its GPU-representation of type [`RenderAsset`].
47pub trait RenderAsset: Send + Sync + 'static + Sized {
48    /// The representation of the asset in the "main world".
49    type SourceAsset: Asset + Clone;
50
51    /// Specifies all ECS data required by [`RenderAsset::prepare_asset`].
52    ///
53    /// For convenience use the [`lifetimeless`](bevy_ecs::system::lifetimeless) [`SystemParam`].
54    type Param: SystemParam;
55
56    /// Whether or not to unload the asset after extracting it to the render world.
57    #[inline]
58    fn asset_usage(_source_asset: &Self::SourceAsset) -> RenderAssetUsages {
59        RenderAssetUsages::default()
60    }
61
62    /// Size of the data the asset will upload to the gpu. Specifying a return value
63    /// will allow the asset to be throttled via [`RenderAssetBytesPerFrame`].
64    #[inline]
65    #[expect(
66        unused_variables,
67        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."
68    )]
69    fn byte_len(source_asset: &Self::SourceAsset) -> Option<usize> {
70        None
71    }
72
73    /// Prepares the [`RenderAsset::SourceAsset`] for the GPU by transforming it into a [`RenderAsset`].
74    ///
75    /// ECS data may be accessed via `param`.
76    fn prepare_asset(
77        source_asset: Self::SourceAsset,
78        asset_id: AssetId<Self::SourceAsset>,
79        param: &mut SystemParamItem<Self::Param>,
80        previous_asset: Option<&Self>,
81    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>>;
82
83    /// Called whenever the [`RenderAsset::SourceAsset`] has been removed.
84    ///
85    /// You can implement this method if you need to access ECS data (via
86    /// `_param`) in order to perform cleanup tasks when the asset is removed.
87    ///
88    /// The default implementation does nothing.
89    fn unload_asset(
90        _source_asset: AssetId<Self::SourceAsset>,
91        _param: &mut SystemParamItem<Self::Param>,
92    ) {
93    }
94
95    /// Make a copy of the asset to be moved to the `RenderWorld` / gpu. Heavy internal data (pixels, vertex attributes)
96    /// should be moved into the copy, leaving this asset with only metadata.
97    /// An error may be returned to indicate that the asset has already been extracted, and should not
98    /// have been modified on the CPU side (as it cannot be transferred to GPU again).
99    /// The previous GPU asset is also provided, which can be used to check if the modification is valid.
100    fn take_gpu_data(
101        _source: &mut Self::SourceAsset,
102        _previous_gpu_asset: Option<&Self>,
103    ) -> Result<Self::SourceAsset, AssetExtractionError> {
104        Err(AssetExtractionError::NoExtractionImplementation)
105    }
106}
107
108/// This plugin extracts the changed assets from the "app world" into the "render world"
109/// and prepares them for the GPU. They can then be accessed from the [`RenderAssets`] resource.
110///
111/// Therefore it sets up the [`ExtractSchedule`] and
112/// [`RenderSystems::PrepareAssets`] steps for the specified [`RenderAsset`].
113///
114/// The `AFTER` generic parameter can be used to specify that `A::prepare_asset` should not be run until
115/// `prepare_assets::<AFTER>` has completed. This allows the `prepare_asset` function to depend on another
116/// prepared [`RenderAsset`], for example `Mesh::prepare_asset` relies on `RenderAssets::<GpuImage>` for morph
117/// targets, so the plugin is created as `RenderAssetPlugin::<RenderMesh, GpuImage>::default()`.
118pub struct RenderAssetPlugin<A: RenderAsset, AFTER: RenderAssetDependency + 'static = ()> {
119    phantom: PhantomData<fn() -> (A, AFTER)>,
120}
121
122impl<A: RenderAsset, AFTER: RenderAssetDependency + 'static> Default
123    for RenderAssetPlugin<A, AFTER>
124{
125    fn default() -> Self {
126        Self {
127            phantom: Default::default(),
128        }
129    }
130}
131
132impl<A: RenderAsset, AFTER: RenderAssetDependency + 'static> Plugin
133    for RenderAssetPlugin<A, AFTER>
134{
135    fn build(&self, app: &mut App) {
136        app.init_resource::<CachedExtractRenderAssetSystemState<A>>();
137        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
138            render_app
139                .init_resource::<ExtractedAssets<A>>()
140                .init_resource::<RenderAssets<A>>()
141                .allow_ambiguous_resource::<RenderAssets<A>>()
142                .init_resource::<PrepareNextFrameAssets<A>>()
143                .add_systems(RenderStartup, collect_render_assets_to_reextract::<A>)
144                .add_systems(
145                    ExtractSchedule,
146                    extract_render_asset::<A>.in_set(AssetExtractionSystems),
147                );
148            AFTER::register_system(
149                render_app,
150                prepare_assets::<A>.in_set(RenderSystems::PrepareAssets),
151            );
152        }
153    }
154}
155
156// helper to allow specifying dependencies between render assets
157pub trait RenderAssetDependency {
158    fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>);
159}
160
161impl RenderAssetDependency for () {
162    fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
163        render_app.add_systems(Render, system);
164    }
165}
166
167impl<A: RenderAsset> RenderAssetDependency for A {
168    fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
169        render_app.add_systems(Render, system.after(prepare_assets::<A>));
170    }
171}
172
173impl<A, B> RenderAssetDependency for (A, B)
174where
175    A: RenderAsset,
176    B: RenderAsset,
177{
178    fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
179        render_app.add_systems(
180            Render,
181            system.after(prepare_assets::<A>).after(prepare_assets::<B>),
182        );
183    }
184}
185
186/// Temporarily stores the extracted and removed assets of the current frame.
187#[derive(Resource)]
188pub struct ExtractedAssets<A: RenderAsset> {
189    /// The assets extracted this frame.
190    ///
191    /// These are assets that were either added or modified this frame.
192    pub extracted: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
193
194    /// IDs of the assets that were removed this frame.
195    ///
196    /// These assets will not be present in [`ExtractedAssets::extracted`].
197    pub removed: HashSet<AssetId<A::SourceAsset>>,
198
199    /// IDs of the assets that were modified this frame.
200    pub modified: HashSet<AssetId<A::SourceAsset>>,
201
202    /// IDs of the assets that were added this frame.
203    pub added: HashSet<AssetId<A::SourceAsset>>,
204}
205
206impl<A: RenderAsset> Default for ExtractedAssets<A> {
207    fn default() -> Self {
208        Self {
209            extracted: Default::default(),
210            removed: Default::default(),
211            modified: Default::default(),
212            added: Default::default(),
213        }
214    }
215}
216
217/// Stores all GPU representations ([`RenderAsset`])
218/// of [`RenderAsset::SourceAsset`] as long as they exist.
219#[derive(Resource)]
220pub struct RenderAssets<A: RenderAsset>(HashMap<AssetId<A::SourceAsset>, A>);
221
222impl<A: RenderAsset> Default for RenderAssets<A> {
223    fn default() -> Self {
224        Self(Default::default())
225    }
226}
227
228impl<A: RenderAsset> RenderAssets<A> {
229    pub fn get(&self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<&A> {
230        self.0.get(&id.into())
231    }
232
233    pub fn get_mut(&mut self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<&mut A> {
234        self.0.get_mut(&id.into())
235    }
236
237    pub fn insert(&mut self, id: impl Into<AssetId<A::SourceAsset>>, value: A) -> Option<A> {
238        self.0.insert(id.into(), value)
239    }
240
241    pub fn remove(&mut self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<A> {
242        self.0.remove(&id.into())
243    }
244
245    pub fn iter(&self) -> impl Iterator<Item = (AssetId<A::SourceAsset>, &A)> {
246        self.0.iter().map(|(k, v)| (*k, v))
247    }
248
249    pub fn iter_mut(&mut self) -> impl Iterator<Item = (AssetId<A::SourceAsset>, &mut A)> {
250        self.0.iter_mut().map(|(k, v)| (*k, v))
251    }
252}
253
254#[derive(Resource)]
255struct CachedExtractRenderAssetSystemState<A: RenderAsset> {
256    state: SystemState<(
257        MessageReader<'static, 'static, AssetEvent<A::SourceAsset>>,
258        ResMut<'static, Assets<A::SourceAsset>>,
259        Option<Res<'static, RenderAssets<A>>>,
260    )>,
261}
262
263impl<A: RenderAsset> FromWorld for CachedExtractRenderAssetSystemState<A> {
264    fn from_world(world: &mut bevy_ecs::world::World) -> Self {
265        Self {
266            state: SystemState::new(world),
267        }
268    }
269}
270
271/// Resource inserted during [`RenderStartup`] containing asset IDs that need
272/// re-extraction from the main world after device recovery.
273#[derive(Resource)]
274pub(crate) struct RenderAssetsToReExtract<A: RenderAsset> {
275    ids: Vec<AssetId<A::SourceAsset>>,
276}
277
278/// Drains all asset IDs from [`RenderAssets<A>`] to mark for re-extraction.
279fn collect_render_assets_to_reextract<A: RenderAsset>(
280    mut commands: Commands,
281    mut render_assets: ResMut<RenderAssets<A>>,
282    mut prepare_next_frame: ResMut<PrepareNextFrameAssets<A>>,
283) {
284    let ids: Vec<_> = render_assets.0.drain().map(|(id, _)| id).collect();
285    prepare_next_frame.assets.clear();
286    if !ids.is_empty() {
287        commands.insert_resource(RenderAssetsToReExtract::<A> { ids });
288    }
289}
290
291/// Extracts all created or modified assets of the corresponding [`RenderAsset::SourceAsset`] type
292/// into the "render world", including any assets invalidated by device recovery.
293pub(crate) fn extract_render_asset<A: RenderAsset>(
294    mut to_reextract: Option<ResMut<RenderAssetsToReExtract<A>>>,
295    mut extracted_assets: ResMut<ExtractedAssets<A>>,
296    mut main_world: ResMut<MainWorld>,
297    mut needs_extracting: Local<HashSet<AssetId<A::SourceAsset>>>,
298) {
299    extracted_assets.extracted.clear();
300    extracted_assets.removed.clear();
301    extracted_assets.modified.clear();
302    extracted_assets.added.clear();
303    needs_extracting.clear();
304
305    let reextract_ids = to_reextract
306        .as_mut()
307        .map(|r| core::mem::take(&mut r.ids))
308        .filter(|ids| !ids.is_empty());
309
310    main_world.resource_scope(
311        |world, mut cached_state: Mut<CachedExtractRenderAssetSystemState<A>>| {
312            let (mut events, mut assets, maybe_render_assets) = cached_state.state.get_mut(world).unwrap();
313
314            if let Some(reextract_ids) = reextract_ids {
315                needs_extracting.extend(reextract_ids);
316            }
317
318            for event in events.read() {
319                #[expect(
320                    clippy::match_same_arms,
321                    reason = "LoadedWithDependencies is marked as a TODO, so it's likely this will no longer lint soon."
322                )]
323                match event {
324                    AssetEvent::Added { id } => {
325                        needs_extracting.insert(*id);
326                    }
327                    AssetEvent::Modified { id } => {
328                        needs_extracting.insert(*id);
329                        extracted_assets.modified.insert(*id);
330                    }
331                    AssetEvent::Removed { .. } => {
332                        // We don't care that the asset was removed from Assets<T> in the main world.
333                        // An asset is only removed from RenderAssets<T> when its last handle is dropped (AssetEvent::Unused).
334                    }
335                    AssetEvent::Unused { id } => {
336                        needs_extracting.remove(id);
337                        extracted_assets.modified.remove(id);
338                        extracted_assets.removed.insert(*id);
339                    }
340                    AssetEvent::LoadedWithDependencies { .. } => {
341                        // TODO: handle this
342                    }
343                }
344            }
345
346            for id in needs_extracting.drain() {
347                if let Some(asset) = assets.get(id) {
348                    let asset_usage = A::asset_usage(asset);
349                    if asset_usage.contains(RenderAssetUsages::RENDER_WORLD) {
350                        if asset_usage == RenderAssetUsages::RENDER_WORLD {
351                            if let Some(asset) = assets.get_mut_untracked(id) {
352                                let previous_asset = maybe_render_assets.as_ref().and_then(|render_assets| render_assets.get(id));
353                                match A::take_gpu_data(asset, previous_asset) {
354                                    Ok(gpu_data_asset) => {
355                                        extracted_assets.extracted.push((id, gpu_data_asset));
356                                        extracted_assets.added.insert(id);
357                                    }
358                                    Err(e) => {
359                                        error!("{} with RenderAssetUsages == RENDER_WORLD cannot be extracted: {e}", core::any::type_name::<A>());
360                                    }
361                                };
362                            }
363                        } else {
364                            extracted_assets.extracted.push((id, asset.clone()));
365                            extracted_assets.added.insert(id);
366                        }
367                    }
368                }
369            }
370
371            cached_state.state.apply(world);
372        },
373    );
374}
375
376// TODO: consider storing inside system?
377/// All assets that should be prepared next frame.
378#[derive(Resource)]
379pub struct PrepareNextFrameAssets<A: RenderAsset> {
380    assets: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
381}
382
383impl<A: RenderAsset> Default for PrepareNextFrameAssets<A> {
384    fn default() -> Self {
385        Self {
386            assets: Default::default(),
387        }
388    }
389}
390
391/// This system prepares all assets of the corresponding [`RenderAsset::SourceAsset`] type
392/// which where extracted this frame for the GPU.
393pub fn prepare_assets<A: RenderAsset>(
394    mut extracted_assets: ResMut<ExtractedAssets<A>>,
395    mut render_assets: ResMut<RenderAssets<A>>,
396    mut prepare_next_frame: ResMut<PrepareNextFrameAssets<A>>,
397    param: StaticSystemParam<<A as RenderAsset>::Param>,
398    bpf: Res<RenderAssetBytesPerFrameLimiter>,
399) {
400    let mut wrote_asset_count = 0;
401
402    let mut param = param.into_inner();
403    let queued_assets = core::mem::take(&mut prepare_next_frame.assets);
404    for (id, extracted_asset) in queued_assets {
405        if extracted_assets.removed.contains(&id) || extracted_assets.added.contains(&id) {
406            // skip previous frame's assets that have been removed or updated
407            continue;
408        }
409
410        let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
411            // we could check if available bytes > byte_len here, but we want to make some
412            // forward progress even if the asset is larger than the max bytes per frame.
413            // this way we always write at least one (sized) asset per frame.
414            // in future we could also consider partial asset uploads.
415            if bpf.exhausted() {
416                prepare_next_frame.assets.push((id, extracted_asset));
417                continue;
418            }
419            size
420        } else {
421            0
422        };
423
424        let previous_asset = render_assets.get(id);
425        match A::prepare_asset(extracted_asset, id, &mut param, previous_asset) {
426            Ok(prepared_asset) => {
427                render_assets.insert(id, prepared_asset);
428                bpf.write_bytes(write_bytes);
429                wrote_asset_count += 1;
430            }
431            Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
432                prepare_next_frame.assets.push((id, extracted_asset));
433            }
434            Err(PrepareAssetError::AsBindGroupError(e)) => {
435                error!(
436                    "{} Bind group construction failed: {e}",
437                    core::any::type_name::<A>()
438                );
439            }
440        }
441    }
442
443    for removed in extracted_assets.removed.drain() {
444        render_assets.remove(removed);
445        A::unload_asset(removed, &mut param);
446    }
447
448    for (id, extracted_asset) in extracted_assets.extracted.drain(..) {
449        // we remove previous here to ensure that if we are updating the asset then
450        // any users will not see the old asset after a new asset is extracted,
451        // even if the new asset is not yet ready or we are out of bytes to write.
452        let previous_asset = render_assets.remove(id);
453
454        let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
455            if bpf.exhausted() {
456                prepare_next_frame.assets.push((id, extracted_asset));
457                continue;
458            }
459            size
460        } else {
461            0
462        };
463
464        match A::prepare_asset(extracted_asset, id, &mut param, previous_asset.as_ref()) {
465            Ok(prepared_asset) => {
466                render_assets.insert(id, prepared_asset);
467                bpf.write_bytes(write_bytes);
468                wrote_asset_count += 1;
469            }
470            Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
471                prepare_next_frame.assets.push((id, extracted_asset));
472            }
473            Err(PrepareAssetError::AsBindGroupError(e)) => {
474                error!(
475                    "{} Bind group construction failed: {e}",
476                    core::any::type_name::<A>()
477                );
478            }
479        }
480    }
481
482    if bpf.exhausted() && !prepare_next_frame.assets.is_empty() {
483        debug!(
484            "{} write budget exhausted with {} assets remaining (wrote {})",
485            core::any::type_name::<A>(),
486            prepare_next_frame.assets.len(),
487            wrote_asset_count
488        );
489    }
490}
491
492pub fn reset_render_asset_bytes_per_frame(
493    mut bpf_limiter: ResMut<RenderAssetBytesPerFrameLimiter>,
494) {
495    bpf_limiter.reset();
496}
497
498pub fn extract_render_asset_bytes_per_frame(
499    bpf: Extract<Res<RenderAssetBytesPerFrame>>,
500    mut bpf_limiter: ResMut<RenderAssetBytesPerFrameLimiter>,
501) {
502    bpf_limiter.max_bytes = bpf.max_bytes;
503}
504
505/// A resource that defines the amount of data allowed to be transferred from CPU to GPU
506/// each frame, preventing choppy frames at the cost of waiting longer for GPU assets
507/// to become available.
508#[derive(Resource, Default)]
509pub struct RenderAssetBytesPerFrame {
510    pub max_bytes: Option<usize>,
511}
512
513impl RenderAssetBytesPerFrame {
514    /// `max_bytes`: the number of bytes to write per frame.
515    ///
516    /// This is a soft limit: only full assets are written currently, uploading stops
517    /// after the first asset that exceeds the limit.
518    ///
519    /// To participate, assets should implement [`RenderAsset::byte_len`]. If the default
520    /// is not overridden, the assets are assumed to be small enough to upload without restriction.
521    pub fn new(max_bytes: usize) -> Self {
522        Self {
523            max_bytes: Some(max_bytes),
524        }
525    }
526}
527
528/// A render-world resource that facilitates limiting the data transferred from CPU to GPU
529/// each frame, preventing choppy frames at the cost of waiting longer for GPU assets
530/// to become available.
531#[derive(Resource, Default)]
532pub struct RenderAssetBytesPerFrameLimiter {
533    /// Populated by [`RenderAssetBytesPerFrame`] during extraction.
534    pub max_bytes: Option<usize>,
535    /// Bytes written this frame.
536    pub bytes_written: AtomicUsize,
537}
538
539impl RenderAssetBytesPerFrameLimiter {
540    /// Reset the available bytes. Called once per frame during extraction by [`crate::RenderPlugin`].
541    pub fn reset(&mut self) {
542        if self.max_bytes.is_none() {
543            return;
544        }
545        self.bytes_written.store(0, Ordering::Relaxed);
546    }
547
548    /// Check how many bytes are available for writing.
549    pub fn available_bytes(&self, required_bytes: usize) -> usize {
550        if let Some(max_bytes) = self.max_bytes {
551            let total_bytes = self
552                .bytes_written
553                .fetch_add(required_bytes, Ordering::Relaxed);
554
555            // The bytes available is the inverse of the amount we overshot max_bytes
556            if total_bytes >= max_bytes {
557                required_bytes.saturating_sub(total_bytes - max_bytes)
558            } else {
559                required_bytes
560            }
561        } else {
562            required_bytes
563        }
564    }
565
566    /// Decreases the available bytes for the current frame.
567    pub(crate) fn write_bytes(&self, bytes: usize) {
568        if self.max_bytes.is_some() && bytes > 0 {
569            self.bytes_written.fetch_add(bytes, Ordering::Relaxed);
570        }
571    }
572
573    /// Returns `true` if there are no remaining bytes available for writing this frame.
574    pub(crate) fn exhausted(&self) -> bool {
575        if let Some(max_bytes) = self.max_bytes {
576            let bytes_written = self.bytes_written.load(Ordering::Relaxed);
577            bytes_written >= max_bytes
578        } else {
579            false
580        }
581    }
582}