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#[derive(SystemSet, Clone, PartialEq, Eq, Debug, Hash)]
29pub struct AssetExtractionSystems;
30
31#[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
40pub trait RenderAsset: Send + Sync + 'static + Sized {
48 type SourceAsset: Asset + Clone;
50
51 type Param: SystemParam;
55
56 #[inline]
58 fn asset_usage(_source_asset: &Self::SourceAsset) -> RenderAssetUsages {
59 RenderAssetUsages::default()
60 }
61
62 #[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 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 fn unload_asset(
90 _source_asset: AssetId<Self::SourceAsset>,
91 _param: &mut SystemParamItem<Self::Param>,
92 ) {
93 }
94
95 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
108pub 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
156pub 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#[derive(Resource)]
188pub struct ExtractedAssets<A: RenderAsset> {
189 pub extracted: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
193
194 pub removed: HashSet<AssetId<A::SourceAsset>>,
198
199 pub modified: HashSet<AssetId<A::SourceAsset>>,
201
202 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#[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#[derive(Resource)]
274pub(crate) struct RenderAssetsToReExtract<A: RenderAsset> {
275 ids: Vec<AssetId<A::SourceAsset>>,
276}
277
278fn 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
291pub(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 }
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 }
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#[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
391pub 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 continue;
408 }
409
410 let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
411 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 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#[derive(Resource, Default)]
509pub struct RenderAssetBytesPerFrame {
510 pub max_bytes: Option<usize>,
511}
512
513impl RenderAssetBytesPerFrame {
514 pub fn new(max_bytes: usize) -> Self {
522 Self {
523 max_bytes: Some(max_bytes),
524 }
525 }
526}
527
528#[derive(Resource, Default)]
532pub struct RenderAssetBytesPerFrameLimiter {
533 pub max_bytes: Option<usize>,
535 pub bytes_written: AtomicUsize,
537}
538
539impl RenderAssetBytesPerFrameLimiter {
540 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 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 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 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 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}