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
173#[derive(Resource)]
175pub struct ExtractedAssets<A: RenderAsset> {
176 pub extracted: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
180
181 pub removed: HashSet<AssetId<A::SourceAsset>>,
185
186 pub modified: HashSet<AssetId<A::SourceAsset>>,
188
189 pub added: HashSet<AssetId<A::SourceAsset>>,
191}
192
193impl<A: RenderAsset> Default for ExtractedAssets<A> {
194 fn default() -> Self {
195 Self {
196 extracted: Default::default(),
197 removed: Default::default(),
198 modified: Default::default(),
199 added: Default::default(),
200 }
201 }
202}
203
204#[derive(Resource)]
207pub struct RenderAssets<A: RenderAsset>(HashMap<AssetId<A::SourceAsset>, A>);
208
209impl<A: RenderAsset> Default for RenderAssets<A> {
210 fn default() -> Self {
211 Self(Default::default())
212 }
213}
214
215impl<A: RenderAsset> RenderAssets<A> {
216 pub fn get(&self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<&A> {
217 self.0.get(&id.into())
218 }
219
220 pub fn get_mut(&mut self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<&mut A> {
221 self.0.get_mut(&id.into())
222 }
223
224 pub fn insert(&mut self, id: impl Into<AssetId<A::SourceAsset>>, value: A) -> Option<A> {
225 self.0.insert(id.into(), value)
226 }
227
228 pub fn remove(&mut self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<A> {
229 self.0.remove(&id.into())
230 }
231
232 pub fn iter(&self) -> impl Iterator<Item = (AssetId<A::SourceAsset>, &A)> {
233 self.0.iter().map(|(k, v)| (*k, v))
234 }
235
236 pub fn iter_mut(&mut self) -> impl Iterator<Item = (AssetId<A::SourceAsset>, &mut A)> {
237 self.0.iter_mut().map(|(k, v)| (*k, v))
238 }
239}
240
241#[derive(Resource)]
242struct CachedExtractRenderAssetSystemState<A: RenderAsset> {
243 state: SystemState<(
244 MessageReader<'static, 'static, AssetEvent<A::SourceAsset>>,
245 ResMut<'static, Assets<A::SourceAsset>>,
246 Option<Res<'static, RenderAssets<A>>>,
247 )>,
248}
249
250impl<A: RenderAsset> FromWorld for CachedExtractRenderAssetSystemState<A> {
251 fn from_world(world: &mut bevy_ecs::world::World) -> Self {
252 Self {
253 state: SystemState::new(world),
254 }
255 }
256}
257
258#[derive(Resource)]
261pub(crate) struct RenderAssetsToReExtract<A: RenderAsset> {
262 ids: Vec<AssetId<A::SourceAsset>>,
263}
264
265fn collect_render_assets_to_reextract<A: RenderAsset>(
267 mut commands: Commands,
268 mut render_assets: ResMut<RenderAssets<A>>,
269 mut prepare_next_frame: ResMut<PrepareNextFrameAssets<A>>,
270) {
271 let ids: Vec<_> = render_assets.0.drain().map(|(id, _)| id).collect();
272 prepare_next_frame.assets.clear();
273 if !ids.is_empty() {
274 commands.insert_resource(RenderAssetsToReExtract::<A> { ids });
275 }
276}
277
278pub(crate) fn extract_render_asset<A: RenderAsset>(
281 mut to_reextract: Option<ResMut<RenderAssetsToReExtract<A>>>,
282 mut extracted_assets: ResMut<ExtractedAssets<A>>,
283 mut main_world: ResMut<MainWorld>,
284 mut needs_extracting: Local<HashSet<AssetId<A::SourceAsset>>>,
285) {
286 extracted_assets.extracted.clear();
287 extracted_assets.removed.clear();
288 extracted_assets.modified.clear();
289 extracted_assets.added.clear();
290 needs_extracting.clear();
291
292 let reextract_ids = to_reextract
293 .as_mut()
294 .map(|r| core::mem::take(&mut r.ids))
295 .filter(|ids| !ids.is_empty());
296
297 main_world.resource_scope(
298 |world, mut cached_state: Mut<CachedExtractRenderAssetSystemState<A>>| {
299 let (mut events, mut assets, maybe_render_assets) = cached_state.state.get_mut(world).unwrap();
300
301 if let Some(reextract_ids) = reextract_ids {
302 needs_extracting.extend(reextract_ids);
303 }
304
305 for event in events.read() {
306 #[expect(
307 clippy::match_same_arms,
308 reason = "LoadedWithDependencies is marked as a TODO, so it's likely this will no longer lint soon."
309 )]
310 match event {
311 AssetEvent::Added { id } => {
312 needs_extracting.insert(*id);
313 }
314 AssetEvent::Modified { id } => {
315 needs_extracting.insert(*id);
316 extracted_assets.modified.insert(*id);
317 }
318 AssetEvent::Removed { .. } => {
319 }
322 AssetEvent::Unused { id } => {
323 needs_extracting.remove(id);
324 extracted_assets.modified.remove(id);
325 extracted_assets.removed.insert(*id);
326 }
327 AssetEvent::LoadedWithDependencies { .. } => {
328 }
330 }
331 }
332
333 for id in needs_extracting.drain() {
334 if let Some(asset) = assets.get(id) {
335 let asset_usage = A::asset_usage(asset);
336 if asset_usage.contains(RenderAssetUsages::RENDER_WORLD) {
337 if asset_usage == RenderAssetUsages::RENDER_WORLD {
338 if let Some(asset) = assets.get_mut_untracked(id) {
339 let previous_asset = maybe_render_assets.as_ref().and_then(|render_assets| render_assets.get(id));
340 match A::take_gpu_data(asset, previous_asset) {
341 Ok(gpu_data_asset) => {
342 extracted_assets.extracted.push((id, gpu_data_asset));
343 extracted_assets.added.insert(id);
344 }
345 Err(e) => {
346 error!("{} with RenderAssetUsages == RENDER_WORLD cannot be extracted: {e}", core::any::type_name::<A>());
347 }
348 };
349 }
350 } else {
351 extracted_assets.extracted.push((id, asset.clone()));
352 extracted_assets.added.insert(id);
353 }
354 }
355 }
356 }
357
358 cached_state.state.apply(world);
359 },
360 );
361}
362
363#[derive(Resource)]
366pub struct PrepareNextFrameAssets<A: RenderAsset> {
367 assets: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
368}
369
370impl<A: RenderAsset> Default for PrepareNextFrameAssets<A> {
371 fn default() -> Self {
372 Self {
373 assets: Default::default(),
374 }
375 }
376}
377
378pub fn prepare_assets<A: RenderAsset>(
381 mut extracted_assets: ResMut<ExtractedAssets<A>>,
382 mut render_assets: ResMut<RenderAssets<A>>,
383 mut prepare_next_frame: ResMut<PrepareNextFrameAssets<A>>,
384 param: StaticSystemParam<<A as RenderAsset>::Param>,
385 bpf: Res<RenderAssetBytesPerFrameLimiter>,
386) {
387 let mut wrote_asset_count = 0;
388
389 let mut param = param.into_inner();
390 let queued_assets = core::mem::take(&mut prepare_next_frame.assets);
391 for (id, extracted_asset) in queued_assets {
392 if extracted_assets.removed.contains(&id) || extracted_assets.added.contains(&id) {
393 continue;
395 }
396
397 let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
398 if bpf.exhausted() {
403 prepare_next_frame.assets.push((id, extracted_asset));
404 continue;
405 }
406 size
407 } else {
408 0
409 };
410
411 let previous_asset = render_assets.get(id);
412 match A::prepare_asset(extracted_asset, id, &mut param, previous_asset) {
413 Ok(prepared_asset) => {
414 render_assets.insert(id, prepared_asset);
415 bpf.write_bytes(write_bytes);
416 wrote_asset_count += 1;
417 }
418 Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
419 prepare_next_frame.assets.push((id, extracted_asset));
420 }
421 Err(PrepareAssetError::AsBindGroupError(e)) => {
422 error!(
423 "{} Bind group construction failed: {e}",
424 core::any::type_name::<A>()
425 );
426 }
427 }
428 }
429
430 for removed in extracted_assets.removed.drain() {
431 render_assets.remove(removed);
432 A::unload_asset(removed, &mut param);
433 }
434
435 for (id, extracted_asset) in extracted_assets.extracted.drain(..) {
436 let previous_asset = render_assets.remove(id);
440
441 let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
442 if bpf.exhausted() {
443 prepare_next_frame.assets.push((id, extracted_asset));
444 continue;
445 }
446 size
447 } else {
448 0
449 };
450
451 match A::prepare_asset(extracted_asset, id, &mut param, previous_asset.as_ref()) {
452 Ok(prepared_asset) => {
453 render_assets.insert(id, prepared_asset);
454 bpf.write_bytes(write_bytes);
455 wrote_asset_count += 1;
456 }
457 Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
458 prepare_next_frame.assets.push((id, extracted_asset));
459 }
460 Err(PrepareAssetError::AsBindGroupError(e)) => {
461 error!(
462 "{} Bind group construction failed: {e}",
463 core::any::type_name::<A>()
464 );
465 }
466 }
467 }
468
469 if bpf.exhausted() && !prepare_next_frame.assets.is_empty() {
470 debug!(
471 "{} write budget exhausted with {} assets remaining (wrote {})",
472 core::any::type_name::<A>(),
473 prepare_next_frame.assets.len(),
474 wrote_asset_count
475 );
476 }
477}
478
479pub fn reset_render_asset_bytes_per_frame(
480 mut bpf_limiter: ResMut<RenderAssetBytesPerFrameLimiter>,
481) {
482 bpf_limiter.reset();
483}
484
485pub fn extract_render_asset_bytes_per_frame(
486 bpf: Extract<Res<RenderAssetBytesPerFrame>>,
487 mut bpf_limiter: ResMut<RenderAssetBytesPerFrameLimiter>,
488) {
489 bpf_limiter.max_bytes = bpf.max_bytes;
490}
491
492#[derive(Resource, Default)]
496pub struct RenderAssetBytesPerFrame {
497 pub max_bytes: Option<usize>,
498}
499
500impl RenderAssetBytesPerFrame {
501 pub fn new(max_bytes: usize) -> Self {
509 Self {
510 max_bytes: Some(max_bytes),
511 }
512 }
513}
514
515#[derive(Resource, Default)]
519pub struct RenderAssetBytesPerFrameLimiter {
520 pub max_bytes: Option<usize>,
522 pub bytes_written: AtomicUsize,
524}
525
526impl RenderAssetBytesPerFrameLimiter {
527 pub fn reset(&mut self) {
529 if self.max_bytes.is_none() {
530 return;
531 }
532 self.bytes_written.store(0, Ordering::Relaxed);
533 }
534
535 pub fn available_bytes(&self, required_bytes: usize) -> usize {
537 if let Some(max_bytes) = self.max_bytes {
538 let total_bytes = self
539 .bytes_written
540 .fetch_add(required_bytes, Ordering::Relaxed);
541
542 if total_bytes >= max_bytes {
544 required_bytes.saturating_sub(total_bytes - max_bytes)
545 } else {
546 required_bytes
547 }
548 } else {
549 required_bytes
550 }
551 }
552
553 pub(crate) fn write_bytes(&self, bytes: usize) {
555 if self.max_bytes.is_some() && bytes > 0 {
556 self.bytes_written.fetch_add(bytes, Ordering::Relaxed);
557 }
558 }
559
560 pub(crate) fn exhausted(&self) -> bool {
562 if let Some(max_bytes) = self.max_bytes {
563 let bytes_written = self.bytes_written.load(Ordering::Relaxed);
564 bytes_written >= max_bytes
565 } else {
566 false
567 }
568 }
569}