1use crate::{
2 render_resource::AsBindGroupError, Extract, ExtractSchedule, MainWorld, Render, RenderApp,
3 RenderSet, Res,
4};
5use bevy_app::{App, Plugin, SubApp};
6pub use bevy_asset::RenderAssetUsages;
7use bevy_asset::{Asset, AssetEvent, AssetId, Assets};
8use bevy_ecs::{
9 prelude::{Commands, EventReader, IntoScheduleConfigs, 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 core::marker::PhantomData;
16use core::sync::atomic::{AtomicUsize, Ordering};
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#[derive(SystemSet, Clone, PartialEq, Eq, Debug, Hash)]
30pub struct ExtractAssetsSet;
31
32pub trait RenderAsset: Send + Sync + 'static + Sized {
40 type SourceAsset: Asset + Clone;
42
43 type Param: SystemParam;
47
48 #[inline]
50 fn asset_usage(_source_asset: &Self::SourceAsset) -> RenderAssetUsages {
51 RenderAssetUsages::default()
52 }
53
54 #[inline]
57 #[expect(
58 unused_variables,
59 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."
60 )]
61 fn byte_len(source_asset: &Self::SourceAsset) -> Option<usize> {
62 None
63 }
64
65 fn prepare_asset(
69 source_asset: Self::SourceAsset,
70 asset_id: AssetId<Self::SourceAsset>,
71 param: &mut SystemParamItem<Self::Param>,
72 ) -> Result<Self, PrepareAssetError<Self::SourceAsset>>;
73
74 fn unload_asset(
81 _source_asset: AssetId<Self::SourceAsset>,
82 _param: &mut SystemParamItem<Self::Param>,
83 ) {
84 }
85}
86
87pub struct RenderAssetPlugin<A: RenderAsset, AFTER: RenderAssetDependency + 'static = ()> {
98 phantom: PhantomData<fn() -> (A, AFTER)>,
99}
100
101impl<A: RenderAsset, AFTER: RenderAssetDependency + 'static> Default
102 for RenderAssetPlugin<A, AFTER>
103{
104 fn default() -> Self {
105 Self {
106 phantom: Default::default(),
107 }
108 }
109}
110
111impl<A: RenderAsset, AFTER: RenderAssetDependency + 'static> Plugin
112 for RenderAssetPlugin<A, AFTER>
113{
114 fn build(&self, app: &mut App) {
115 app.init_resource::<CachedExtractRenderAssetSystemState<A>>();
116 if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
117 render_app
118 .init_resource::<ExtractedAssets<A>>()
119 .init_resource::<RenderAssets<A>>()
120 .init_resource::<PrepareNextFrameAssets<A>>()
121 .add_systems(
122 ExtractSchedule,
123 extract_render_asset::<A>.in_set(ExtractAssetsSet),
124 );
125 AFTER::register_system(
126 render_app,
127 prepare_assets::<A>.in_set(RenderSet::PrepareAssets),
128 );
129 }
130 }
131}
132
133pub trait RenderAssetDependency {
135 fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>);
136}
137
138impl RenderAssetDependency for () {
139 fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
140 render_app.add_systems(Render, system);
141 }
142}
143
144impl<A: RenderAsset> RenderAssetDependency for A {
145 fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
146 render_app.add_systems(Render, system.after(prepare_assets::<A>));
147 }
148}
149
150#[derive(Resource)]
152pub struct ExtractedAssets<A: RenderAsset> {
153 pub extracted: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
157
158 pub removed: HashSet<AssetId<A::SourceAsset>>,
162
163 pub modified: HashSet<AssetId<A::SourceAsset>>,
165
166 pub added: HashSet<AssetId<A::SourceAsset>>,
168}
169
170impl<A: RenderAsset> Default for ExtractedAssets<A> {
171 fn default() -> Self {
172 Self {
173 extracted: Default::default(),
174 removed: Default::default(),
175 modified: Default::default(),
176 added: Default::default(),
177 }
178 }
179}
180
181#[derive(Resource)]
184pub struct RenderAssets<A: RenderAsset>(HashMap<AssetId<A::SourceAsset>, A>);
185
186impl<A: RenderAsset> Default for RenderAssets<A> {
187 fn default() -> Self {
188 Self(Default::default())
189 }
190}
191
192impl<A: RenderAsset> RenderAssets<A> {
193 pub fn get(&self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<&A> {
194 self.0.get(&id.into())
195 }
196
197 pub fn get_mut(&mut self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<&mut A> {
198 self.0.get_mut(&id.into())
199 }
200
201 pub fn insert(&mut self, id: impl Into<AssetId<A::SourceAsset>>, value: A) -> Option<A> {
202 self.0.insert(id.into(), value)
203 }
204
205 pub fn remove(&mut self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<A> {
206 self.0.remove(&id.into())
207 }
208
209 pub fn iter(&self) -> impl Iterator<Item = (AssetId<A::SourceAsset>, &A)> {
210 self.0.iter().map(|(k, v)| (*k, v))
211 }
212
213 pub fn iter_mut(&mut self) -> impl Iterator<Item = (AssetId<A::SourceAsset>, &mut A)> {
214 self.0.iter_mut().map(|(k, v)| (*k, v))
215 }
216}
217
218#[derive(Resource)]
219struct CachedExtractRenderAssetSystemState<A: RenderAsset> {
220 state: SystemState<(
221 EventReader<'static, 'static, AssetEvent<A::SourceAsset>>,
222 ResMut<'static, Assets<A::SourceAsset>>,
223 )>,
224}
225
226impl<A: RenderAsset> FromWorld for CachedExtractRenderAssetSystemState<A> {
227 fn from_world(world: &mut bevy_ecs::world::World) -> Self {
228 Self {
229 state: SystemState::new(world),
230 }
231 }
232}
233
234pub(crate) fn extract_render_asset<A: RenderAsset>(
237 mut commands: Commands,
238 mut main_world: ResMut<MainWorld>,
239) {
240 main_world.resource_scope(
241 |world, mut cached_state: Mut<CachedExtractRenderAssetSystemState<A>>| {
242 let (mut events, mut assets) = cached_state.state.get_mut(world);
243
244 let mut needs_extracting = <HashSet<_>>::default();
245 let mut removed = <HashSet<_>>::default();
246 let mut modified = <HashSet<_>>::default();
247
248 for event in events.read() {
249 #[expect(
250 clippy::match_same_arms,
251 reason = "LoadedWithDependencies is marked as a TODO, so it's likely this will no longer lint soon."
252 )]
253 match event {
254 AssetEvent::Added { id } => {
255 needs_extracting.insert(*id);
256 }
257 AssetEvent::Modified { id } => {
258 needs_extracting.insert(*id);
259 modified.insert(*id);
260 }
261 AssetEvent::Removed { .. } => {
262 }
265 AssetEvent::Unused { id } => {
266 needs_extracting.remove(id);
267 modified.remove(id);
268 removed.insert(*id);
269 }
270 AssetEvent::LoadedWithDependencies { .. } => {
271 }
273 }
274 }
275
276 let mut extracted_assets = Vec::new();
277 let mut added = <HashSet<_>>::default();
278 for id in needs_extracting.drain() {
279 if let Some(asset) = assets.get(id) {
280 let asset_usage = A::asset_usage(asset);
281 if asset_usage.contains(RenderAssetUsages::RENDER_WORLD) {
282 if asset_usage == RenderAssetUsages::RENDER_WORLD {
283 if let Some(asset) = assets.remove(id) {
284 extracted_assets.push((id, asset));
285 added.insert(id);
286 }
287 } else {
288 extracted_assets.push((id, asset.clone()));
289 added.insert(id);
290 }
291 }
292 }
293 }
294
295 commands.insert_resource(ExtractedAssets::<A> {
296 extracted: extracted_assets,
297 removed,
298 modified,
299 added,
300 });
301 cached_state.state.apply(world);
302 },
303 );
304}
305
306#[derive(Resource)]
309pub struct PrepareNextFrameAssets<A: RenderAsset> {
310 assets: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
311}
312
313impl<A: RenderAsset> Default for PrepareNextFrameAssets<A> {
314 fn default() -> Self {
315 Self {
316 assets: Default::default(),
317 }
318 }
319}
320
321pub fn prepare_assets<A: RenderAsset>(
324 mut extracted_assets: ResMut<ExtractedAssets<A>>,
325 mut render_assets: ResMut<RenderAssets<A>>,
326 mut prepare_next_frame: ResMut<PrepareNextFrameAssets<A>>,
327 param: StaticSystemParam<<A as RenderAsset>::Param>,
328 bpf: Res<RenderAssetBytesPerFrameLimiter>,
329) {
330 let mut wrote_asset_count = 0;
331
332 let mut param = param.into_inner();
333 let queued_assets = core::mem::take(&mut prepare_next_frame.assets);
334 for (id, extracted_asset) in queued_assets {
335 if extracted_assets.removed.contains(&id) || extracted_assets.added.contains(&id) {
336 continue;
338 }
339
340 let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
341 if bpf.exhausted() {
346 prepare_next_frame.assets.push((id, extracted_asset));
347 continue;
348 }
349 size
350 } else {
351 0
352 };
353
354 match A::prepare_asset(extracted_asset, id, &mut param) {
355 Ok(prepared_asset) => {
356 render_assets.insert(id, prepared_asset);
357 bpf.write_bytes(write_bytes);
358 wrote_asset_count += 1;
359 }
360 Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
361 prepare_next_frame.assets.push((id, extracted_asset));
362 }
363 Err(PrepareAssetError::AsBindGroupError(e)) => {
364 error!(
365 "{} Bind group construction failed: {e}",
366 core::any::type_name::<A>()
367 );
368 }
369 }
370 }
371
372 for removed in extracted_assets.removed.drain() {
373 render_assets.remove(removed);
374 A::unload_asset(removed, &mut param);
375 }
376
377 for (id, extracted_asset) in extracted_assets.extracted.drain(..) {
378 render_assets.remove(id);
382
383 let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
384 if bpf.exhausted() {
385 prepare_next_frame.assets.push((id, extracted_asset));
386 continue;
387 }
388 size
389 } else {
390 0
391 };
392
393 match A::prepare_asset(extracted_asset, id, &mut param) {
394 Ok(prepared_asset) => {
395 render_assets.insert(id, prepared_asset);
396 bpf.write_bytes(write_bytes);
397 wrote_asset_count += 1;
398 }
399 Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
400 prepare_next_frame.assets.push((id, extracted_asset));
401 }
402 Err(PrepareAssetError::AsBindGroupError(e)) => {
403 error!(
404 "{} Bind group construction failed: {e}",
405 core::any::type_name::<A>()
406 );
407 }
408 }
409 }
410
411 if bpf.exhausted() && !prepare_next_frame.assets.is_empty() {
412 debug!(
413 "{} write budget exhausted with {} assets remaining (wrote {})",
414 core::any::type_name::<A>(),
415 prepare_next_frame.assets.len(),
416 wrote_asset_count
417 );
418 }
419}
420
421pub fn reset_render_asset_bytes_per_frame(
422 mut bpf_limiter: ResMut<RenderAssetBytesPerFrameLimiter>,
423) {
424 bpf_limiter.reset();
425}
426
427pub fn extract_render_asset_bytes_per_frame(
428 bpf: Extract<Res<RenderAssetBytesPerFrame>>,
429 mut bpf_limiter: ResMut<RenderAssetBytesPerFrameLimiter>,
430) {
431 bpf_limiter.max_bytes = bpf.max_bytes;
432}
433
434#[derive(Resource, Default)]
438pub struct RenderAssetBytesPerFrame {
439 pub max_bytes: Option<usize>,
440}
441
442impl RenderAssetBytesPerFrame {
443 pub fn new(max_bytes: usize) -> Self {
451 Self {
452 max_bytes: Some(max_bytes),
453 }
454 }
455}
456
457#[derive(Resource, Default)]
461pub struct RenderAssetBytesPerFrameLimiter {
462 pub max_bytes: Option<usize>,
464 pub bytes_written: AtomicUsize,
466}
467
468impl RenderAssetBytesPerFrameLimiter {
469 pub fn reset(&mut self) {
471 if self.max_bytes.is_none() {
472 return;
473 }
474 self.bytes_written.store(0, Ordering::Relaxed);
475 }
476
477 pub fn available_bytes(&self, required_bytes: usize) -> usize {
479 if let Some(max_bytes) = self.max_bytes {
480 let total_bytes = self
481 .bytes_written
482 .fetch_add(required_bytes, Ordering::Relaxed);
483
484 if total_bytes >= max_bytes {
486 required_bytes.saturating_sub(total_bytes - max_bytes)
487 } else {
488 required_bytes
489 }
490 } else {
491 required_bytes
492 }
493 }
494
495 fn write_bytes(&self, bytes: usize) {
497 if self.max_bytes.is_some() && bytes > 0 {
498 self.bytes_written.fetch_add(bytes, Ordering::Relaxed);
499 }
500 }
501
502 fn exhausted(&self) -> bool {
504 if let Some(max_bytes) = self.max_bytes {
505 let bytes_written = self.bytes_written.load(Ordering::Relaxed);
506 bytes_written >= max_bytes
507 } else {
508 false
509 }
510 }
511}