1use crate::asset_changed::AssetChanges;
2use crate::{Asset, AssetEvent, AssetHandleProvider, AssetId, AssetServer, Handle, UntypedHandle};
3use alloc::{sync::Arc, vec::Vec};
4use bevy_ecs::{
5 message::MessageWriter,
6 resource::Resource,
7 system::{Res, ResMut, SystemChangeTick},
8};
9use bevy_platform::collections::HashMap;
10use bevy_reflect::{Reflect, TypePath};
11use core::{any::TypeId, iter::Enumerate, marker::PhantomData, sync::atomic::AtomicU32};
12use crossbeam_channel::{Receiver, Sender};
13use serde::{Deserialize, Serialize};
14use thiserror::Error;
15use uuid::Uuid;
16
17#[derive(
20 Debug, Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd, Reflect, Serialize, Deserialize,
21)]
22pub struct AssetIndex {
23 pub(crate) generation: u32,
24 pub(crate) index: u32,
25}
26
27impl AssetIndex {
28 pub fn to_bits(self) -> u64 {
32 let Self { generation, index } = self;
33 ((generation as u64) << 32) | index as u64
34 }
35 pub fn from_bits(bits: u64) -> Self {
38 let index = ((bits << 32) >> 32) as u32;
39 let generation = (bits >> 32) as u32;
40 Self { generation, index }
41 }
42}
43
44pub(crate) struct AssetIndexAllocator {
46 next_index: AtomicU32,
48 recycled_queue_sender: Sender<AssetIndex>,
49 recycled_queue_receiver: Receiver<AssetIndex>,
51 recycled_sender: Sender<AssetIndex>,
52 recycled_receiver: Receiver<AssetIndex>,
53}
54
55impl Default for AssetIndexAllocator {
56 fn default() -> Self {
57 let (recycled_queue_sender, recycled_queue_receiver) = crossbeam_channel::unbounded();
58 let (recycled_sender, recycled_receiver) = crossbeam_channel::unbounded();
59 Self {
60 recycled_queue_sender,
61 recycled_queue_receiver,
62 recycled_sender,
63 recycled_receiver,
64 next_index: Default::default(),
65 }
66 }
67}
68
69impl AssetIndexAllocator {
70 pub fn reserve(&self) -> AssetIndex {
73 if let Ok(mut recycled) = self.recycled_queue_receiver.try_recv() {
74 recycled.generation += 1;
75 self.recycled_sender.send(recycled).unwrap();
76 recycled
77 } else {
78 AssetIndex {
79 index: self
80 .next_index
81 .fetch_add(1, core::sync::atomic::Ordering::Relaxed),
82 generation: 0,
83 }
84 }
85 }
86
87 pub fn recycle(&self, index: AssetIndex) {
89 self.recycled_queue_sender.send(index).unwrap();
90 }
91}
92
93#[derive(Asset, TypePath)]
97pub struct LoadedUntypedAsset {
98 #[dependency]
100 pub handle: UntypedHandle,
101}
102
103#[derive(Default)]
105enum Entry<A: Asset> {
106 #[default]
108 None,
109 Some { value: Option<A>, generation: u32 },
111}
112
113struct DenseAssetStorage<A: Asset> {
115 storage: Vec<Entry<A>>,
116 len: u32,
117 allocator: Arc<AssetIndexAllocator>,
118}
119
120impl<A: Asset> Default for DenseAssetStorage<A> {
121 fn default() -> Self {
122 Self {
123 len: 0,
124 storage: Default::default(),
125 allocator: Default::default(),
126 }
127 }
128}
129
130impl<A: Asset> DenseAssetStorage<A> {
131 pub(crate) fn len(&self) -> usize {
133 self.len as usize
134 }
135
136 pub(crate) fn is_empty(&self) -> bool {
138 self.len == 0
139 }
140
141 pub(crate) fn insert(
143 &mut self,
144 index: AssetIndex,
145 asset: A,
146 ) -> Result<bool, InvalidGenerationError> {
147 self.flush();
148 let entry = &mut self.storage[index.index as usize];
149 if let Entry::Some { value, generation } = entry {
150 if *generation == index.generation {
151 let exists = value.is_some();
152 if !exists {
153 self.len += 1;
154 }
155 *value = Some(asset);
156 Ok(exists)
157 } else {
158 Err(InvalidGenerationError::Occupied {
159 index,
160 current_generation: *generation,
161 })
162 }
163 } else {
164 Err(InvalidGenerationError::Removed { index })
165 }
166 }
167
168 pub(crate) fn remove_dropped(&mut self, index: AssetIndex) -> Option<A> {
171 self.remove_internal(index, |dense_storage| {
172 dense_storage.storage[index.index as usize] = Entry::None;
173 dense_storage.allocator.recycle(index);
174 })
175 }
176
177 pub(crate) fn remove_still_alive(&mut self, index: AssetIndex) -> Option<A> {
181 self.remove_internal(index, |_| {})
182 }
183
184 fn remove_internal(
185 &mut self,
186 index: AssetIndex,
187 removed_action: impl FnOnce(&mut Self),
188 ) -> Option<A> {
189 self.flush();
190 let value = match &mut self.storage[index.index as usize] {
191 Entry::None => return None,
192 Entry::Some { value, generation } => {
193 if *generation == index.generation {
194 value.take().inspect(|_| self.len -= 1)
195 } else {
196 return None;
197 }
198 }
199 };
200 removed_action(self);
201 value
202 }
203
204 pub(crate) fn get(&self, index: AssetIndex) -> Option<&A> {
205 let entry = self.storage.get(index.index as usize)?;
206 match entry {
207 Entry::None => None,
208 Entry::Some { value, generation } => {
209 if *generation == index.generation {
210 value.as_ref()
211 } else {
212 None
213 }
214 }
215 }
216 }
217
218 pub(crate) fn get_mut(&mut self, index: AssetIndex) -> Option<&mut A> {
219 let entry = self.storage.get_mut(index.index as usize)?;
220 match entry {
221 Entry::None => None,
222 Entry::Some { value, generation } => {
223 if *generation == index.generation {
224 value.as_mut()
225 } else {
226 None
227 }
228 }
229 }
230 }
231
232 pub(crate) fn flush(&mut self) {
233 let new_len = self
235 .allocator
236 .next_index
237 .load(core::sync::atomic::Ordering::Relaxed);
238 self.storage.resize_with(new_len as usize, || Entry::Some {
239 value: None,
240 generation: 0,
241 });
242 while let Ok(recycled) = self.allocator.recycled_receiver.try_recv() {
243 let entry = &mut self.storage[recycled.index as usize];
244 *entry = Entry::Some {
245 value: None,
246 generation: recycled.generation,
247 };
248 }
249 }
250
251 pub(crate) fn get_index_allocator(&self) -> Arc<AssetIndexAllocator> {
252 self.allocator.clone()
253 }
254
255 pub(crate) fn ids(&self) -> impl Iterator<Item = AssetId<A>> + '_ {
256 self.storage
257 .iter()
258 .enumerate()
259 .filter_map(|(i, v)| match v {
260 Entry::None => None,
261 Entry::Some { value, generation } => {
262 if value.is_some() {
263 Some(AssetId::from(AssetIndex {
264 index: i as u32,
265 generation: *generation,
266 }))
267 } else {
268 None
269 }
270 }
271 })
272 }
273}
274
275#[derive(Resource)]
287pub struct Assets<A: Asset> {
288 dense_storage: DenseAssetStorage<A>,
289 hash_map: HashMap<Uuid, A>,
290 handle_provider: AssetHandleProvider,
291 queued_events: Vec<AssetEvent<A>>,
292 duplicate_handles: HashMap<AssetId<A>, u16>,
295}
296
297impl<A: Asset> Default for Assets<A> {
298 fn default() -> Self {
299 let dense_storage = DenseAssetStorage::default();
300 let handle_provider =
301 AssetHandleProvider::new(TypeId::of::<A>(), dense_storage.get_index_allocator());
302 Self {
303 dense_storage,
304 handle_provider,
305 hash_map: Default::default(),
306 queued_events: Default::default(),
307 duplicate_handles: Default::default(),
308 }
309 }
310}
311
312impl<A: Asset> Assets<A> {
313 pub fn get_handle_provider(&self) -> AssetHandleProvider {
316 self.handle_provider.clone()
317 }
318
319 pub fn reserve_handle(&self) -> Handle<A> {
321 self.handle_provider.reserve_handle().typed::<A>()
322 }
323
324 pub fn insert(
329 &mut self,
330 id: impl Into<AssetId<A>>,
331 asset: A,
332 ) -> Result<(), InvalidGenerationError> {
333 match id.into() {
334 AssetId::Index { index, .. } => self.insert_with_index(index, asset).map(|_| ()),
335 AssetId::Uuid { uuid } => {
336 self.insert_with_uuid(uuid, asset);
337 Ok(())
338 }
339 }
340 }
341
342 pub fn get_or_insert_with(
348 &mut self,
349 id: impl Into<AssetId<A>>,
350 insert_fn: impl FnOnce() -> A,
351 ) -> Result<&mut A, InvalidGenerationError> {
352 let id: AssetId<A> = id.into();
353 if self.get(id).is_none() {
354 self.insert(id, insert_fn())?;
355 }
356 Ok(self
359 .get_mut(id)
360 .expect("the Asset was none even though we checked or inserted"))
361 }
362
363 pub fn contains(&self, id: impl Into<AssetId<A>>) -> bool {
365 match id.into() {
366 AssetId::Index { index, .. } => self.dense_storage.get(index).is_some(),
367 AssetId::Uuid { uuid } => self.hash_map.contains_key(&uuid),
368 }
369 }
370
371 pub(crate) fn insert_with_uuid(&mut self, uuid: Uuid, asset: A) -> Option<A> {
372 let result = self.hash_map.insert(uuid, asset);
373 if result.is_some() {
374 self.queued_events
375 .push(AssetEvent::Modified { id: uuid.into() });
376 } else {
377 self.queued_events
378 .push(AssetEvent::Added { id: uuid.into() });
379 }
380 result
381 }
382 pub(crate) fn insert_with_index(
383 &mut self,
384 index: AssetIndex,
385 asset: A,
386 ) -> Result<bool, InvalidGenerationError> {
387 let replaced = self.dense_storage.insert(index, asset)?;
388 if replaced {
389 self.queued_events
390 .push(AssetEvent::Modified { id: index.into() });
391 } else {
392 self.queued_events
393 .push(AssetEvent::Added { id: index.into() });
394 }
395 Ok(replaced)
396 }
397
398 #[inline]
400 pub fn add(&mut self, asset: impl Into<A>) -> Handle<A> {
401 let index = self.dense_storage.allocator.reserve();
402 self.insert_with_index(index, asset.into()).unwrap();
403 Handle::Strong(
404 self.handle_provider
405 .get_handle(index.into(), false, None, None),
406 )
407 }
408
409 #[inline]
414 pub fn get_strong_handle(&mut self, id: AssetId<A>) -> Option<Handle<A>> {
415 if !self.contains(id) {
416 return None;
417 }
418 *self.duplicate_handles.entry(id).or_insert(0) += 1;
419 let index = match id {
420 AssetId::Index { index, .. } => index.into(),
421 AssetId::Uuid { uuid } => uuid.into(),
422 };
423 Some(Handle::Strong(
424 self.handle_provider.get_handle(index, false, None, None),
425 ))
426 }
427
428 #[inline]
431 pub fn get(&self, id: impl Into<AssetId<A>>) -> Option<&A> {
432 match id.into() {
433 AssetId::Index { index, .. } => self.dense_storage.get(index),
434 AssetId::Uuid { uuid } => self.hash_map.get(&uuid),
435 }
436 }
437
438 #[inline]
441 pub fn get_mut(&mut self, id: impl Into<AssetId<A>>) -> Option<&mut A> {
442 let id: AssetId<A> = id.into();
443 let result = match id {
444 AssetId::Index { index, .. } => self.dense_storage.get_mut(index),
445 AssetId::Uuid { uuid } => self.hash_map.get_mut(&uuid),
446 };
447 if result.is_some() {
448 self.queued_events.push(AssetEvent::Modified { id });
449 }
450 result
451 }
452
453 #[inline]
457 pub fn get_mut_untracked(&mut self, id: impl Into<AssetId<A>>) -> Option<&mut A> {
458 let id: AssetId<A> = id.into();
459 match id {
460 AssetId::Index { index, .. } => self.dense_storage.get_mut(index),
461 AssetId::Uuid { uuid } => self.hash_map.get_mut(&uuid),
462 }
463 }
464
465 pub fn remove(&mut self, id: impl Into<AssetId<A>>) -> Option<A> {
468 let id: AssetId<A> = id.into();
469 let result = self.remove_untracked(id);
470 if result.is_some() {
471 self.queued_events.push(AssetEvent::Removed { id });
472 }
473 result
474 }
475
476 pub fn remove_untracked(&mut self, id: impl Into<AssetId<A>>) -> Option<A> {
481 let id: AssetId<A> = id.into();
482 self.duplicate_handles.remove(&id);
483 match id {
484 AssetId::Index { index, .. } => self.dense_storage.remove_still_alive(index),
485 AssetId::Uuid { uuid } => self.hash_map.remove(&uuid),
486 }
487 }
488
489 pub(crate) fn remove_dropped(&mut self, id: AssetId<A>) {
491 match self.duplicate_handles.get_mut(&id) {
492 None => {}
493 Some(0) => {
494 self.duplicate_handles.remove(&id);
495 }
496 Some(value) => {
497 *value -= 1;
498 return;
499 }
500 }
501
502 let existed = match id {
503 AssetId::Index { index, .. } => self.dense_storage.remove_dropped(index).is_some(),
504 AssetId::Uuid { uuid } => self.hash_map.remove(&uuid).is_some(),
505 };
506
507 self.queued_events.push(AssetEvent::Unused { id });
508 if existed {
509 self.queued_events.push(AssetEvent::Removed { id });
510 }
511 }
512
513 pub fn is_empty(&self) -> bool {
515 self.dense_storage.is_empty() && self.hash_map.is_empty()
516 }
517
518 pub fn len(&self) -> usize {
520 self.dense_storage.len() + self.hash_map.len()
521 }
522
523 pub fn ids(&self) -> impl Iterator<Item = AssetId<A>> + '_ {
525 self.dense_storage
526 .ids()
527 .chain(self.hash_map.keys().map(|uuid| AssetId::from(*uuid)))
528 }
529
530 pub fn iter(&self) -> impl Iterator<Item = (AssetId<A>, &A)> {
533 self.dense_storage
534 .storage
535 .iter()
536 .enumerate()
537 .filter_map(|(i, v)| match v {
538 Entry::None => None,
539 Entry::Some { value, generation } => value.as_ref().map(|v| {
540 let id = AssetId::Index {
541 index: AssetIndex {
542 generation: *generation,
543 index: i as u32,
544 },
545 marker: PhantomData,
546 };
547 (id, v)
548 }),
549 })
550 .chain(
551 self.hash_map
552 .iter()
553 .map(|(i, v)| (AssetId::Uuid { uuid: *i }, v)),
554 )
555 }
556
557 pub fn iter_mut(&mut self) -> AssetsMutIterator<'_, A> {
560 AssetsMutIterator {
561 dense_storage: self.dense_storage.storage.iter_mut().enumerate(),
562 hash_map: self.hash_map.iter_mut(),
563 queued_events: &mut self.queued_events,
564 }
565 }
566
567 pub fn track_assets(mut assets: ResMut<Self>, asset_server: Res<AssetServer>) {
570 let assets = &mut *assets;
571 let mut infos = asset_server.data.infos.write();
576 while let Ok(drop_event) = assets.handle_provider.drop_receiver.try_recv() {
577 let id = drop_event.id.typed();
578
579 if drop_event.asset_server_managed {
580 let untyped_id = id.untyped();
581
582 if !infos.process_handle_drop(untyped_id) {
584 continue;
586 }
587 }
588
589 assets.remove_dropped(id);
590 }
591 }
592
593 pub(crate) fn asset_events(
597 mut assets: ResMut<Self>,
598 mut messages: MessageWriter<AssetEvent<A>>,
599 asset_changes: Option<ResMut<AssetChanges<A>>>,
600 ticks: SystemChangeTick,
601 ) {
602 use AssetEvent::{Added, LoadedWithDependencies, Modified, Removed};
603
604 if let Some(mut asset_changes) = asset_changes {
605 for new_event in &assets.queued_events {
606 match new_event {
607 Removed { id } | AssetEvent::Unused { id } => asset_changes.remove(id),
608 Added { id } | Modified { id } | LoadedWithDependencies { id } => {
609 asset_changes.insert(*id, ticks.this_run());
610 }
611 };
612 }
613 }
614 messages.write_batch(assets.queued_events.drain(..));
615 }
616
617 pub(crate) fn asset_events_condition(assets: Res<Self>) -> bool {
622 !assets.queued_events.is_empty()
623 }
624}
625
626pub struct AssetsMutIterator<'a, A: Asset> {
628 queued_events: &'a mut Vec<AssetEvent<A>>,
629 dense_storage: Enumerate<core::slice::IterMut<'a, Entry<A>>>,
630 hash_map: bevy_platform::collections::hash_map::IterMut<'a, Uuid, A>,
631}
632
633impl<'a, A: Asset> Iterator for AssetsMutIterator<'a, A> {
634 type Item = (AssetId<A>, &'a mut A);
635
636 fn next(&mut self) -> Option<Self::Item> {
637 for (i, entry) in &mut self.dense_storage {
638 match entry {
639 Entry::None => {
640 continue;
641 }
642 Entry::Some { value, generation } => {
643 let id = AssetId::Index {
644 index: AssetIndex {
645 generation: *generation,
646 index: i as u32,
647 },
648 marker: PhantomData,
649 };
650 self.queued_events.push(AssetEvent::Modified { id });
651 if let Some(value) = value {
652 return Some((id, value));
653 }
654 }
655 }
656 }
657 if let Some((key, value)) = self.hash_map.next() {
658 let id = AssetId::Uuid { uuid: *key };
659 self.queued_events.push(AssetEvent::Modified { id });
660 Some((id, value))
661 } else {
662 None
663 }
664 }
665}
666
667#[derive(Error, Debug, PartialEq, Eq)]
669pub enum InvalidGenerationError {
670 #[error("AssetIndex {index:?} has an invalid generation. The current generation is: '{current_generation}'.")]
671 Occupied {
672 index: AssetIndex,
673 current_generation: u32,
674 },
675 #[error("AssetIndex {index:?} has been removed")]
676 Removed { index: AssetIndex },
677}
678
679#[cfg(test)]
680mod test {
681 use crate::AssetIndex;
682
683 #[test]
684 fn asset_index_round_trip() {
685 let asset_index = AssetIndex {
686 generation: 42,
687 index: 1337,
688 };
689 let roundtripped = AssetIndex::from_bits(asset_index.to_bits());
690 assert_eq!(asset_index, roundtripped);
691 }
692}