1use crate::{
2 render_resource::AsBindGroupError, ExtractSchedule, MainWorld, Render, RenderApp,
3 RenderSystems, Res,
4};
5use bevy_app::{App, Plugin, SubApp};
6use bevy_asset::RenderAssetUsages;
7use bevy_asset::{Asset, AssetEvent, AssetId, Assets, UntypedAssetId};
8use bevy_ecs::{
9 prelude::{Commands, IntoScheduleConfigs, MessageReader, 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 bevy_render::render_asset::RenderAssetBytesPerFrameLimiter;
16use core::marker::PhantomData;
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 AssetExtractionSystems;
31
32#[deprecated(since = "0.17.0", note = "Renamed to `AssetExtractionSystems`.")]
34pub type ExtractAssetsSet = AssetExtractionSystems;
35
36pub trait ErasedRenderAsset: Send + Sync + 'static {
44 type SourceAsset: Asset + Clone;
46 type ErasedAsset: Send + Sync + 'static + Sized;
48
49 type Param: SystemParam;
53
54 #[inline]
56 fn asset_usage(_source_asset: &Self::SourceAsset) -> RenderAssetUsages {
57 RenderAssetUsages::default()
58 }
59
60 #[inline]
63 #[expect(
64 unused_variables,
65 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."
66 )]
67 fn byte_len(erased_asset: &Self::SourceAsset) -> Option<usize> {
68 None
69 }
70
71 fn prepare_asset(
75 source_asset: Self::SourceAsset,
76 asset_id: AssetId<Self::SourceAsset>,
77 param: &mut SystemParamItem<Self::Param>,
78 ) -> Result<Self::ErasedAsset, PrepareAssetError<Self::SourceAsset>>;
79
80 fn unload_asset(
87 _source_asset: AssetId<Self::SourceAsset>,
88 _param: &mut SystemParamItem<Self::Param>,
89 ) {
90 }
91}
92
93pub struct ErasedRenderAssetPlugin<
104 A: ErasedRenderAsset,
105 AFTER: ErasedRenderAssetDependency + 'static = (),
106> {
107 phantom: PhantomData<fn() -> (A, AFTER)>,
108}
109
110impl<A: ErasedRenderAsset, AFTER: ErasedRenderAssetDependency + 'static> Default
111 for ErasedRenderAssetPlugin<A, AFTER>
112{
113 fn default() -> Self {
114 Self {
115 phantom: Default::default(),
116 }
117 }
118}
119
120impl<A: ErasedRenderAsset, AFTER: ErasedRenderAssetDependency + 'static> Plugin
121 for ErasedRenderAssetPlugin<A, AFTER>
122{
123 fn build(&self, app: &mut App) {
124 app.init_resource::<CachedExtractErasedRenderAssetSystemState<A>>();
125 }
126
127 fn finish(&self, app: &mut App) {
128 if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
129 render_app
130 .init_resource::<ExtractedAssets<A>>()
131 .init_resource::<ErasedRenderAssets<A::ErasedAsset>>()
132 .init_resource::<PrepareNextFrameAssets<A>>()
133 .add_systems(
134 ExtractSchedule,
135 extract_erased_render_asset::<A>.in_set(AssetExtractionSystems),
136 );
137 AFTER::register_system(
138 render_app,
139 prepare_erased_assets::<A>.in_set(RenderSystems::PrepareAssets),
140 );
141 }
142 }
143}
144
145pub trait ErasedRenderAssetDependency {
147 fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>);
148}
149
150impl ErasedRenderAssetDependency for () {
151 fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
152 render_app.add_systems(Render, system);
153 }
154}
155
156impl<A: ErasedRenderAsset> ErasedRenderAssetDependency for A {
157 fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
158 render_app.add_systems(Render, system.after(prepare_erased_assets::<A>));
159 }
160}
161
162#[derive(Resource)]
164pub struct ExtractedAssets<A: ErasedRenderAsset> {
165 pub extracted: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
169
170 pub removed: HashSet<AssetId<A::SourceAsset>>,
174
175 pub modified: HashSet<AssetId<A::SourceAsset>>,
177
178 pub added: HashSet<AssetId<A::SourceAsset>>,
180}
181
182impl<A: ErasedRenderAsset> Default for ExtractedAssets<A> {
183 fn default() -> Self {
184 Self {
185 extracted: Default::default(),
186 removed: Default::default(),
187 modified: Default::default(),
188 added: Default::default(),
189 }
190 }
191}
192
193#[derive(Resource)]
196pub struct ErasedRenderAssets<ERA>(HashMap<UntypedAssetId, ERA>);
197
198impl<ERA> Default for ErasedRenderAssets<ERA> {
199 fn default() -> Self {
200 Self(Default::default())
201 }
202}
203
204impl<ERA> ErasedRenderAssets<ERA> {
205 pub fn get(&self, id: impl Into<UntypedAssetId>) -> Option<&ERA> {
206 self.0.get(&id.into())
207 }
208
209 pub fn get_mut(&mut self, id: impl Into<UntypedAssetId>) -> Option<&mut ERA> {
210 self.0.get_mut(&id.into())
211 }
212
213 pub fn insert(&mut self, id: impl Into<UntypedAssetId>, value: ERA) -> Option<ERA> {
214 self.0.insert(id.into(), value)
215 }
216
217 pub fn remove(&mut self, id: impl Into<UntypedAssetId>) -> Option<ERA> {
218 self.0.remove(&id.into())
219 }
220
221 pub fn iter(&self) -> impl Iterator<Item = (UntypedAssetId, &ERA)> {
222 self.0.iter().map(|(k, v)| (*k, v))
223 }
224
225 pub fn iter_mut(&mut self) -> impl Iterator<Item = (UntypedAssetId, &mut ERA)> {
226 self.0.iter_mut().map(|(k, v)| (*k, v))
227 }
228}
229
230#[derive(Resource)]
231struct CachedExtractErasedRenderAssetSystemState<A: ErasedRenderAsset> {
232 state: SystemState<(
233 MessageReader<'static, 'static, AssetEvent<A::SourceAsset>>,
234 ResMut<'static, Assets<A::SourceAsset>>,
235 )>,
236}
237
238impl<A: ErasedRenderAsset> FromWorld for CachedExtractErasedRenderAssetSystemState<A> {
239 fn from_world(world: &mut bevy_ecs::world::World) -> Self {
240 Self {
241 state: SystemState::new(world),
242 }
243 }
244}
245
246pub(crate) fn extract_erased_render_asset<A: ErasedRenderAsset>(
249 mut commands: Commands,
250 mut main_world: ResMut<MainWorld>,
251) {
252 main_world.resource_scope(
253 |world, mut cached_state: Mut<CachedExtractErasedRenderAssetSystemState<A>>| {
254 let (mut events, mut assets) = cached_state.state.get_mut(world);
255
256 let mut needs_extracting = <HashSet<_>>::default();
257 let mut removed = <HashSet<_>>::default();
258 let mut modified = <HashSet<_>>::default();
259
260 for event in events.read() {
261 #[expect(
262 clippy::match_same_arms,
263 reason = "LoadedWithDependencies is marked as a TODO, so it's likely this will no longer lint soon."
264 )]
265 match event {
266 AssetEvent::Added { id } => {
267 needs_extracting.insert(*id);
268 }
269 AssetEvent::Modified { id } => {
270 needs_extracting.insert(*id);
271 modified.insert(*id);
272 }
273 AssetEvent::Removed { .. } => {
274 }
277 AssetEvent::Unused { id } => {
278 needs_extracting.remove(id);
279 modified.remove(id);
280 removed.insert(*id);
281 }
282 AssetEvent::LoadedWithDependencies { .. } => {
283 }
285 }
286 }
287
288 let mut extracted_assets = Vec::new();
289 let mut added = <HashSet<_>>::default();
290 for id in needs_extracting.drain() {
291 if let Some(asset) = assets.get(id) {
292 let asset_usage = A::asset_usage(asset);
293 if asset_usage.contains(RenderAssetUsages::RENDER_WORLD) {
294 if asset_usage == RenderAssetUsages::RENDER_WORLD {
295 if let Some(asset) = assets.remove(id) {
296 extracted_assets.push((id, asset));
297 added.insert(id);
298 }
299 } else {
300 extracted_assets.push((id, asset.clone()));
301 added.insert(id);
302 }
303 }
304 }
305 }
306
307 commands.insert_resource(ExtractedAssets::<A> {
308 extracted: extracted_assets,
309 removed,
310 modified,
311 added,
312 });
313 cached_state.state.apply(world);
314 },
315 );
316}
317
318#[derive(Resource)]
321pub struct PrepareNextFrameAssets<A: ErasedRenderAsset> {
322 assets: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
323}
324
325impl<A: ErasedRenderAsset> Default for PrepareNextFrameAssets<A> {
326 fn default() -> Self {
327 Self {
328 assets: Default::default(),
329 }
330 }
331}
332
333pub fn prepare_erased_assets<A: ErasedRenderAsset>(
336 mut extracted_assets: ResMut<ExtractedAssets<A>>,
337 mut render_assets: ResMut<ErasedRenderAssets<A::ErasedAsset>>,
338 mut prepare_next_frame: ResMut<PrepareNextFrameAssets<A>>,
339 param: StaticSystemParam<<A as ErasedRenderAsset>::Param>,
340 bpf: Res<RenderAssetBytesPerFrameLimiter>,
341) {
342 let mut wrote_asset_count = 0;
343
344 let mut param = param.into_inner();
345 let queued_assets = core::mem::take(&mut prepare_next_frame.assets);
346 for (id, extracted_asset) in queued_assets {
347 if extracted_assets.removed.contains(&id) || extracted_assets.added.contains(&id) {
348 continue;
350 }
351
352 let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
353 if bpf.exhausted() {
358 prepare_next_frame.assets.push((id, extracted_asset));
359 continue;
360 }
361 size
362 } else {
363 0
364 };
365
366 match A::prepare_asset(extracted_asset, id, &mut param) {
367 Ok(prepared_asset) => {
368 render_assets.insert(id, prepared_asset);
369 bpf.write_bytes(write_bytes);
370 wrote_asset_count += 1;
371 }
372 Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
373 prepare_next_frame.assets.push((id, extracted_asset));
374 }
375 Err(PrepareAssetError::AsBindGroupError(e)) => {
376 error!(
377 "{} Bind group construction failed: {e}",
378 core::any::type_name::<A>()
379 );
380 }
381 }
382 }
383
384 for removed in extracted_assets.removed.drain() {
385 render_assets.remove(removed);
386 A::unload_asset(removed, &mut param);
387 }
388
389 for (id, extracted_asset) in extracted_assets.extracted.drain(..) {
390 render_assets.remove(id);
394
395 let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
396 if bpf.exhausted() {
397 prepare_next_frame.assets.push((id, extracted_asset));
398 continue;
399 }
400 size
401 } else {
402 0
403 };
404
405 match A::prepare_asset(extracted_asset, id, &mut param) {
406 Ok(prepared_asset) => {
407 render_assets.insert(id, prepared_asset);
408 bpf.write_bytes(write_bytes);
409 wrote_asset_count += 1;
410 }
411 Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
412 prepare_next_frame.assets.push((id, extracted_asset));
413 }
414 Err(PrepareAssetError::AsBindGroupError(e)) => {
415 error!(
416 "{} Bind group construction failed: {e}",
417 core::any::type_name::<A>()
418 );
419 }
420 }
421 }
422
423 if bpf.exhausted() && !prepare_next_frame.assets.is_empty() {
424 debug!(
425 "{} write budget exhausted with {} assets remaining (wrote {})",
426 core::any::type_name::<A>(),
427 prepare_next_frame.assets.len(),
428 wrote_asset_count
429 );
430 }
431}