1use crate::{
2 render_resource::AsBindGroupError, Extract, ExtractSchedule, MainWorld, Render, RenderApp,
3 RenderSystems, Res,
4};
5use bevy_app::{App, Plugin, SubApp};
6use bevy_asset::{Asset, AssetEvent, AssetId, Assets, RenderAssetUsages};
7use bevy_ecs::{
8 prelude::{Commands, IntoScheduleConfigs, MessageReader, ResMut, Resource},
9 schedule::{ScheduleConfigs, SystemSet},
10 system::{ScheduleSystem, StaticSystemParam, SystemParam, SystemParamItem, SystemState},
11 world::{FromWorld, Mut},
12};
13use bevy_platform::collections::{HashMap, HashSet};
14use core::marker::PhantomData;
15use core::sync::atomic::{AtomicUsize, Ordering};
16use thiserror::Error;
17use tracing::{debug, 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#[deprecated(since = "0.17.0", note = "Renamed to `AssetExtractionSystems`.")]
33pub type ExtractAssetsSet = AssetExtractionSystems;
34
35pub trait RenderAsset: Send + Sync + 'static + Sized {
43 type SourceAsset: Asset + Clone;
45
46 type Param: SystemParam;
50
51 #[inline]
53 fn asset_usage(_source_asset: &Self::SourceAsset) -> RenderAssetUsages {
54 RenderAssetUsages::default()
55 }
56
57 #[inline]
60 #[expect(
61 unused_variables,
62 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."
63 )]
64 fn byte_len(source_asset: &Self::SourceAsset) -> Option<usize> {
65 None
66 }
67
68 fn prepare_asset(
72 source_asset: Self::SourceAsset,
73 asset_id: AssetId<Self::SourceAsset>,
74 param: &mut SystemParamItem<Self::Param>,
75 previous_asset: Option<&Self>,
76 ) -> Result<Self, PrepareAssetError<Self::SourceAsset>>;
77
78 fn unload_asset(
85 _source_asset: AssetId<Self::SourceAsset>,
86 _param: &mut SystemParamItem<Self::Param>,
87 ) {
88 }
89}
90
91pub struct RenderAssetPlugin<A: RenderAsset, AFTER: RenderAssetDependency + 'static = ()> {
102 phantom: PhantomData<fn() -> (A, AFTER)>,
103}
104
105impl<A: RenderAsset, AFTER: RenderAssetDependency + 'static> Default
106 for RenderAssetPlugin<A, AFTER>
107{
108 fn default() -> Self {
109 Self {
110 phantom: Default::default(),
111 }
112 }
113}
114
115impl<A: RenderAsset, AFTER: RenderAssetDependency + 'static> Plugin
116 for RenderAssetPlugin<A, AFTER>
117{
118 fn build(&self, app: &mut App) {
119 app.init_resource::<CachedExtractRenderAssetSystemState<A>>();
120 if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
121 render_app
122 .init_resource::<ExtractedAssets<A>>()
123 .init_resource::<RenderAssets<A>>()
124 .init_resource::<PrepareNextFrameAssets<A>>()
125 .add_systems(
126 ExtractSchedule,
127 extract_render_asset::<A>.in_set(AssetExtractionSystems),
128 );
129 AFTER::register_system(
130 render_app,
131 prepare_assets::<A>.in_set(RenderSystems::PrepareAssets),
132 );
133 }
134 }
135}
136
137pub trait RenderAssetDependency {
139 fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>);
140}
141
142impl RenderAssetDependency for () {
143 fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
144 render_app.add_systems(Render, system);
145 }
146}
147
148impl<A: RenderAsset> RenderAssetDependency for A {
149 fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
150 render_app.add_systems(Render, system.after(prepare_assets::<A>));
151 }
152}
153
154#[derive(Resource)]
156pub struct ExtractedAssets<A: RenderAsset> {
157 pub extracted: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
161
162 pub removed: HashSet<AssetId<A::SourceAsset>>,
166
167 pub modified: HashSet<AssetId<A::SourceAsset>>,
169
170 pub added: HashSet<AssetId<A::SourceAsset>>,
172}
173
174impl<A: RenderAsset> Default for ExtractedAssets<A> {
175 fn default() -> Self {
176 Self {
177 extracted: Default::default(),
178 removed: Default::default(),
179 modified: Default::default(),
180 added: Default::default(),
181 }
182 }
183}
184
185#[derive(Resource)]
188pub struct RenderAssets<A: RenderAsset>(HashMap<AssetId<A::SourceAsset>, A>);
189
190impl<A: RenderAsset> Default for RenderAssets<A> {
191 fn default() -> Self {
192 Self(Default::default())
193 }
194}
195
196impl<A: RenderAsset> RenderAssets<A> {
197 pub fn get(&self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<&A> {
198 self.0.get(&id.into())
199 }
200
201 pub fn get_mut(&mut self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<&mut A> {
202 self.0.get_mut(&id.into())
203 }
204
205 pub fn insert(&mut self, id: impl Into<AssetId<A::SourceAsset>>, value: A) -> Option<A> {
206 self.0.insert(id.into(), value)
207 }
208
209 pub fn remove(&mut self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<A> {
210 self.0.remove(&id.into())
211 }
212
213 pub fn iter(&self) -> impl Iterator<Item = (AssetId<A::SourceAsset>, &A)> {
214 self.0.iter().map(|(k, v)| (*k, v))
215 }
216
217 pub fn iter_mut(&mut self) -> impl Iterator<Item = (AssetId<A::SourceAsset>, &mut A)> {
218 self.0.iter_mut().map(|(k, v)| (*k, v))
219 }
220}
221
222#[derive(Resource)]
223struct CachedExtractRenderAssetSystemState<A: RenderAsset> {
224 state: SystemState<(
225 MessageReader<'static, 'static, AssetEvent<A::SourceAsset>>,
226 ResMut<'static, Assets<A::SourceAsset>>,
227 )>,
228}
229
230impl<A: RenderAsset> FromWorld for CachedExtractRenderAssetSystemState<A> {
231 fn from_world(world: &mut bevy_ecs::world::World) -> Self {
232 Self {
233 state: SystemState::new(world),
234 }
235 }
236}
237
238pub(crate) fn extract_render_asset<A: RenderAsset>(
241 mut commands: Commands,
242 mut main_world: ResMut<MainWorld>,
243) {
244 main_world.resource_scope(
245 |world, mut cached_state: Mut<CachedExtractRenderAssetSystemState<A>>| {
246 let (mut events, mut assets) = cached_state.state.get_mut(world);
247
248 let mut needs_extracting = <HashSet<_>>::default();
249 let mut removed = <HashSet<_>>::default();
250 let mut modified = <HashSet<_>>::default();
251
252 for event in events.read() {
253 #[expect(
254 clippy::match_same_arms,
255 reason = "LoadedWithDependencies is marked as a TODO, so it's likely this will no longer lint soon."
256 )]
257 match event {
258 AssetEvent::Added { id } => {
259 needs_extracting.insert(*id);
260 }
261 AssetEvent::Modified { id } => {
262 needs_extracting.insert(*id);
263 modified.insert(*id);
264 }
265 AssetEvent::Removed { .. } => {
266 }
269 AssetEvent::Unused { id } => {
270 needs_extracting.remove(id);
271 modified.remove(id);
272 removed.insert(*id);
273 }
274 AssetEvent::LoadedWithDependencies { .. } => {
275 }
277 }
278 }
279
280 let mut extracted_assets = Vec::new();
281 let mut added = <HashSet<_>>::default();
282 for id in needs_extracting.drain() {
283 if let Some(asset) = assets.get(id) {
284 let asset_usage = A::asset_usage(asset);
285 if asset_usage.contains(RenderAssetUsages::RENDER_WORLD) {
286 if asset_usage == RenderAssetUsages::RENDER_WORLD {
287 if let Some(asset) = assets.remove(id) {
288 extracted_assets.push((id, asset));
289 added.insert(id);
290 }
291 } else {
292 extracted_assets.push((id, asset.clone()));
293 added.insert(id);
294 }
295 }
296 }
297 }
298
299 commands.insert_resource(ExtractedAssets::<A> {
300 extracted: extracted_assets,
301 removed,
302 modified,
303 added,
304 });
305 cached_state.state.apply(world);
306 },
307 );
308}
309
310#[derive(Resource)]
313pub struct PrepareNextFrameAssets<A: RenderAsset> {
314 assets: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
315}
316
317impl<A: RenderAsset> Default for PrepareNextFrameAssets<A> {
318 fn default() -> Self {
319 Self {
320 assets: Default::default(),
321 }
322 }
323}
324
325pub fn prepare_assets<A: RenderAsset>(
328 mut extracted_assets: ResMut<ExtractedAssets<A>>,
329 mut render_assets: ResMut<RenderAssets<A>>,
330 mut prepare_next_frame: ResMut<PrepareNextFrameAssets<A>>,
331 param: StaticSystemParam<<A as RenderAsset>::Param>,
332 bpf: Res<RenderAssetBytesPerFrameLimiter>,
333) {
334 let mut wrote_asset_count = 0;
335
336 let mut param = param.into_inner();
337 let queued_assets = core::mem::take(&mut prepare_next_frame.assets);
338 for (id, extracted_asset) in queued_assets {
339 if extracted_assets.removed.contains(&id) || extracted_assets.added.contains(&id) {
340 continue;
342 }
343
344 let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
345 if bpf.exhausted() {
350 prepare_next_frame.assets.push((id, extracted_asset));
351 continue;
352 }
353 size
354 } else {
355 0
356 };
357
358 let previous_asset = render_assets.get(id);
359 match A::prepare_asset(extracted_asset, id, &mut param, previous_asset) {
360 Ok(prepared_asset) => {
361 render_assets.insert(id, prepared_asset);
362 bpf.write_bytes(write_bytes);
363 wrote_asset_count += 1;
364 }
365 Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
366 prepare_next_frame.assets.push((id, extracted_asset));
367 }
368 Err(PrepareAssetError::AsBindGroupError(e)) => {
369 error!(
370 "{} Bind group construction failed: {e}",
371 core::any::type_name::<A>()
372 );
373 }
374 }
375 }
376
377 for removed in extracted_assets.removed.drain() {
378 render_assets.remove(removed);
379 A::unload_asset(removed, &mut param);
380 }
381
382 for (id, extracted_asset) in extracted_assets.extracted.drain(..) {
383 let previous_asset = render_assets.remove(id);
387
388 let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
389 if bpf.exhausted() {
390 prepare_next_frame.assets.push((id, extracted_asset));
391 continue;
392 }
393 size
394 } else {
395 0
396 };
397
398 match A::prepare_asset(extracted_asset, id, &mut param, previous_asset.as_ref()) {
399 Ok(prepared_asset) => {
400 render_assets.insert(id, prepared_asset);
401 bpf.write_bytes(write_bytes);
402 wrote_asset_count += 1;
403 }
404 Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
405 prepare_next_frame.assets.push((id, extracted_asset));
406 }
407 Err(PrepareAssetError::AsBindGroupError(e)) => {
408 error!(
409 "{} Bind group construction failed: {e}",
410 core::any::type_name::<A>()
411 );
412 }
413 }
414 }
415
416 if bpf.exhausted() && !prepare_next_frame.assets.is_empty() {
417 debug!(
418 "{} write budget exhausted with {} assets remaining (wrote {})",
419 core::any::type_name::<A>(),
420 prepare_next_frame.assets.len(),
421 wrote_asset_count
422 );
423 }
424}
425
426pub fn reset_render_asset_bytes_per_frame(
427 mut bpf_limiter: ResMut<RenderAssetBytesPerFrameLimiter>,
428) {
429 bpf_limiter.reset();
430}
431
432pub fn extract_render_asset_bytes_per_frame(
433 bpf: Extract<Res<RenderAssetBytesPerFrame>>,
434 mut bpf_limiter: ResMut<RenderAssetBytesPerFrameLimiter>,
435) {
436 bpf_limiter.max_bytes = bpf.max_bytes;
437}
438
439#[derive(Resource, Default)]
443pub struct RenderAssetBytesPerFrame {
444 pub max_bytes: Option<usize>,
445}
446
447impl RenderAssetBytesPerFrame {
448 pub fn new(max_bytes: usize) -> Self {
456 Self {
457 max_bytes: Some(max_bytes),
458 }
459 }
460}
461
462#[derive(Resource, Default)]
466pub struct RenderAssetBytesPerFrameLimiter {
467 pub max_bytes: Option<usize>,
469 pub bytes_written: AtomicUsize,
471}
472
473impl RenderAssetBytesPerFrameLimiter {
474 pub fn reset(&mut self) {
476 if self.max_bytes.is_none() {
477 return;
478 }
479 self.bytes_written.store(0, Ordering::Relaxed);
480 }
481
482 pub fn available_bytes(&self, required_bytes: usize) -> usize {
484 if let Some(max_bytes) = self.max_bytes {
485 let total_bytes = self
486 .bytes_written
487 .fetch_add(required_bytes, Ordering::Relaxed);
488
489 if total_bytes >= max_bytes {
491 required_bytes.saturating_sub(total_bytes - max_bytes)
492 } else {
493 required_bytes
494 }
495 } else {
496 required_bytes
497 }
498 }
499
500 pub(crate) fn write_bytes(&self, bytes: usize) {
502 if self.max_bytes.is_some() && bytes > 0 {
503 self.bytes_written.fetch_add(bytes, Ordering::Relaxed);
504 }
505 }
506
507 pub(crate) fn exhausted(&self) -> bool {
509 if let Some(max_bytes) = self.max_bytes {
510 let bytes_written = self.bytes_written.load(Ordering::Relaxed);
511 bytes_written >= max_bytes
512 } else {
513 false
514 }
515 }
516}