1use crate::{self as bevy_asset};
2use crate::{Asset, AssetEvent, AssetHandleProvider, AssetId, AssetServer, Handle, UntypedHandle};
3use bevy_ecs::{
4 prelude::EventWriter,
5 system::{Res, ResMut, Resource},
6};
7use bevy_reflect::{Reflect, TypePath};
8use bevy_utils::HashMap;
9use crossbeam_channel::{Receiver, Sender};
10use serde::{Deserialize, Serialize};
11use std::{
12 any::TypeId,
13 iter::Enumerate,
14 marker::PhantomData,
15 sync::{atomic::AtomicU32, Arc},
16};
17use thiserror::Error;
18use uuid::Uuid;
19
20#[derive(
23 Debug, Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd, Reflect, Serialize, Deserialize,
24)]
25pub struct AssetIndex {
26 pub(crate) generation: u32,
27 pub(crate) index: u32,
28}
29
30impl AssetIndex {
31 pub fn to_bits(self) -> u64 {
35 let Self { generation, index } = self;
36 ((generation as u64) << 32) | index as u64
37 }
38 pub fn from_bits(bits: u64) -> Self {
41 let index = ((bits << 32) >> 32) as u32;
42 let generation = (bits >> 32) as u32;
43 Self { generation, index }
44 }
45}
46
47pub(crate) struct AssetIndexAllocator {
49 next_index: AtomicU32,
51 recycled_queue_sender: Sender<AssetIndex>,
52 recycled_queue_receiver: Receiver<AssetIndex>,
54 recycled_sender: Sender<AssetIndex>,
55 recycled_receiver: Receiver<AssetIndex>,
56}
57
58impl Default for AssetIndexAllocator {
59 fn default() -> Self {
60 let (recycled_queue_sender, recycled_queue_receiver) = crossbeam_channel::unbounded();
61 let (recycled_sender, recycled_receiver) = crossbeam_channel::unbounded();
62 Self {
63 recycled_queue_sender,
64 recycled_queue_receiver,
65 recycled_sender,
66 recycled_receiver,
67 next_index: Default::default(),
68 }
69 }
70}
71
72impl AssetIndexAllocator {
73 pub fn reserve(&self) -> AssetIndex {
76 if let Ok(mut recycled) = self.recycled_queue_receiver.try_recv() {
77 recycled.generation += 1;
78 self.recycled_sender.send(recycled).unwrap();
79 recycled
80 } else {
81 AssetIndex {
82 index: self
83 .next_index
84 .fetch_add(1, std::sync::atomic::Ordering::Relaxed),
85 generation: 0,
86 }
87 }
88 }
89
90 pub fn recycle(&self, index: AssetIndex) {
92 self.recycled_queue_sender.send(index).unwrap();
93 }
94}
95
96#[derive(Asset, TypePath)]
100pub struct LoadedUntypedAsset {
101 #[dependency]
102 pub handle: UntypedHandle,
103}
104
105#[derive(Default)]
107enum Entry<A: Asset> {
108 #[default]
110 None,
111 Some { value: Option<A>, generation: u32 },
113}
114
115struct DenseAssetStorage<A: Asset> {
117 storage: Vec<Entry<A>>,
118 len: u32,
119 allocator: Arc<AssetIndexAllocator>,
120}
121
122impl<A: Asset> Default for DenseAssetStorage<A> {
123 fn default() -> Self {
124 Self {
125 len: 0,
126 storage: Default::default(),
127 allocator: Default::default(),
128 }
129 }
130}
131
132impl<A: Asset> DenseAssetStorage<A> {
133 pub(crate) fn len(&self) -> usize {
135 self.len as usize
136 }
137
138 pub(crate) fn is_empty(&self) -> bool {
140 self.len == 0
141 }
142
143 pub(crate) fn insert(
145 &mut self,
146 index: AssetIndex,
147 asset: A,
148 ) -> Result<bool, InvalidGenerationError> {
149 self.flush();
150 let entry = &mut self.storage[index.index as usize];
151 if let Entry::Some { value, generation } = entry {
152 if *generation == index.generation {
153 let exists = value.is_some();
154 if !exists {
155 self.len += 1;
156 }
157 *value = Some(asset);
158 Ok(exists)
159 } else {
160 Err(InvalidGenerationError {
161 index,
162 current_generation: *generation,
163 })
164 }
165 } else {
166 unreachable!("entries should always be valid after a flush");
167 }
168 }
169
170 pub(crate) fn remove_dropped(&mut self, index: AssetIndex) -> Option<A> {
173 self.remove_internal(index, |dense_storage| {
174 dense_storage.storage[index.index as usize] = Entry::None;
175 dense_storage.allocator.recycle(index);
176 })
177 }
178
179 pub(crate) fn remove_still_alive(&mut self, index: AssetIndex) -> Option<A> {
183 self.remove_internal(index, |_| {})
184 }
185
186 fn remove_internal(
187 &mut self,
188 index: AssetIndex,
189 removed_action: impl FnOnce(&mut Self),
190 ) -> Option<A> {
191 self.flush();
192 let value = match &mut self.storage[index.index as usize] {
193 Entry::None => return None,
194 Entry::Some { value, generation } => {
195 if *generation == index.generation {
196 value.take().map(|value| {
197 self.len -= 1;
198 value
199 })
200 } else {
201 return None;
202 }
203 }
204 };
205 removed_action(self);
206 value
207 }
208
209 pub(crate) fn get(&self, index: AssetIndex) -> Option<&A> {
210 let entry = self.storage.get(index.index as usize)?;
211 match entry {
212 Entry::None => None,
213 Entry::Some { value, generation } => {
214 if *generation == index.generation {
215 value.as_ref()
216 } else {
217 None
218 }
219 }
220 }
221 }
222
223 pub(crate) fn get_mut(&mut self, index: AssetIndex) -> Option<&mut A> {
224 let entry = self.storage.get_mut(index.index as usize)?;
225 match entry {
226 Entry::None => None,
227 Entry::Some { value, generation } => {
228 if *generation == index.generation {
229 value.as_mut()
230 } else {
231 None
232 }
233 }
234 }
235 }
236
237 pub(crate) fn flush(&mut self) {
238 let new_len = self
240 .allocator
241 .next_index
242 .load(std::sync::atomic::Ordering::Relaxed);
243 self.storage.resize_with(new_len as usize, || Entry::Some {
244 value: None,
245 generation: 0,
246 });
247 while let Ok(recycled) = self.allocator.recycled_receiver.try_recv() {
248 let entry = &mut self.storage[recycled.index as usize];
249 *entry = Entry::Some {
250 value: None,
251 generation: recycled.generation,
252 };
253 }
254 }
255
256 pub(crate) fn get_index_allocator(&self) -> Arc<AssetIndexAllocator> {
257 self.allocator.clone()
258 }
259
260 pub(crate) fn ids(&self) -> impl Iterator<Item = AssetId<A>> + '_ {
261 self.storage
262 .iter()
263 .enumerate()
264 .filter_map(|(i, v)| match v {
265 Entry::None => None,
266 Entry::Some { value, generation } => {
267 if value.is_some() {
268 Some(AssetId::from(AssetIndex {
269 index: i as u32,
270 generation: *generation,
271 }))
272 } else {
273 None
274 }
275 }
276 })
277 }
278}
279
280#[derive(Resource)]
290pub struct Assets<A: Asset> {
291 dense_storage: DenseAssetStorage<A>,
292 hash_map: HashMap<Uuid, A>,
293 handle_provider: AssetHandleProvider,
294 queued_events: Vec<AssetEvent<A>>,
295 duplicate_handles: HashMap<AssetId<A>, u16>,
298}
299
300impl<A: Asset> Default for Assets<A> {
301 fn default() -> Self {
302 let dense_storage = DenseAssetStorage::default();
303 let handle_provider =
304 AssetHandleProvider::new(TypeId::of::<A>(), dense_storage.get_index_allocator());
305 Self {
306 dense_storage,
307 handle_provider,
308 hash_map: Default::default(),
309 queued_events: Default::default(),
310 duplicate_handles: Default::default(),
311 }
312 }
313}
314
315impl<A: Asset> Assets<A> {
316 pub fn get_handle_provider(&self) -> AssetHandleProvider {
319 self.handle_provider.clone()
320 }
321
322 pub fn reserve_handle(&self) -> Handle<A> {
324 self.handle_provider.reserve_handle().typed::<A>()
325 }
326
327 pub fn insert(&mut self, id: impl Into<AssetId<A>>, asset: A) {
329 match id.into() {
330 AssetId::Index { index, .. } => {
331 self.insert_with_index(index, asset).unwrap();
332 }
333 AssetId::Uuid { uuid } => {
334 self.insert_with_uuid(uuid, asset);
335 }
336 }
337 }
338
339 pub fn get_or_insert_with(
342 &mut self,
343 id: impl Into<AssetId<A>>,
344 insert_fn: impl FnOnce() -> A,
345 ) -> &mut A {
346 let id: AssetId<A> = id.into();
347 if self.get(id).is_none() {
348 self.insert(id, insert_fn());
349 }
350 self.get_mut(id).unwrap()
351 }
352
353 pub fn contains(&self, id: impl Into<AssetId<A>>) -> bool {
355 match id.into() {
356 AssetId::Index { index, .. } => self.dense_storage.get(index).is_some(),
357 AssetId::Uuid { uuid } => self.hash_map.contains_key(&uuid),
358 }
359 }
360
361 pub(crate) fn insert_with_uuid(&mut self, uuid: Uuid, asset: A) -> Option<A> {
362 let result = self.hash_map.insert(uuid, asset);
363 if result.is_some() {
364 self.queued_events
365 .push(AssetEvent::Modified { id: uuid.into() });
366 } else {
367 self.queued_events
368 .push(AssetEvent::Added { id: uuid.into() });
369 }
370 result
371 }
372 pub(crate) fn insert_with_index(
373 &mut self,
374 index: AssetIndex,
375 asset: A,
376 ) -> Result<bool, InvalidGenerationError> {
377 let replaced = self.dense_storage.insert(index, asset)?;
378 if replaced {
379 self.queued_events
380 .push(AssetEvent::Modified { id: index.into() });
381 } else {
382 self.queued_events
383 .push(AssetEvent::Added { id: index.into() });
384 }
385 Ok(replaced)
386 }
387
388 #[inline]
390 pub fn add(&mut self, asset: impl Into<A>) -> Handle<A> {
391 let index = self.dense_storage.allocator.reserve();
392 self.insert_with_index(index, asset.into()).unwrap();
393 Handle::Strong(
394 self.handle_provider
395 .get_handle(index.into(), false, None, None),
396 )
397 }
398
399 #[inline]
404 pub fn get_strong_handle(&mut self, id: AssetId<A>) -> Option<Handle<A>> {
405 if !self.contains(id) {
406 return None;
407 }
408 *self.duplicate_handles.entry(id).or_insert(0) += 1;
409 let index = match id {
410 AssetId::Index { index, .. } => index.into(),
411 AssetId::Uuid { uuid } => uuid.into(),
412 };
413 Some(Handle::Strong(
414 self.handle_provider.get_handle(index, false, None, None),
415 ))
416 }
417
418 #[inline]
421 pub fn get(&self, id: impl Into<AssetId<A>>) -> Option<&A> {
422 match id.into() {
423 AssetId::Index { index, .. } => self.dense_storage.get(index),
424 AssetId::Uuid { uuid } => self.hash_map.get(&uuid),
425 }
426 }
427
428 #[inline]
431 pub fn get_mut(&mut self, id: impl Into<AssetId<A>>) -> Option<&mut A> {
432 let id: AssetId<A> = id.into();
433 let result = match id {
434 AssetId::Index { index, .. } => self.dense_storage.get_mut(index),
435 AssetId::Uuid { uuid } => self.hash_map.get_mut(&uuid),
436 };
437 if result.is_some() {
438 self.queued_events.push(AssetEvent::Modified { id });
439 }
440 result
441 }
442
443 pub fn remove(&mut self, id: impl Into<AssetId<A>>) -> Option<A> {
446 let id: AssetId<A> = id.into();
447 let result = self.remove_untracked(id);
448 if result.is_some() {
449 self.queued_events.push(AssetEvent::Removed { id });
450 }
451 result
452 }
453
454 pub fn remove_untracked(&mut self, id: impl Into<AssetId<A>>) -> Option<A> {
457 let id: AssetId<A> = id.into();
458 self.duplicate_handles.remove(&id);
459 match id {
460 AssetId::Index { index, .. } => self.dense_storage.remove_still_alive(index),
461 AssetId::Uuid { uuid } => self.hash_map.remove(&uuid),
462 }
463 }
464
465 pub(crate) fn remove_dropped(&mut self, id: AssetId<A>) {
467 match self.duplicate_handles.get_mut(&id) {
468 None | Some(0) => {}
469 Some(value) => {
470 *value -= 1;
471 return;
472 }
473 }
474 let existed = match id {
475 AssetId::Index { index, .. } => self.dense_storage.remove_dropped(index).is_some(),
476 AssetId::Uuid { uuid } => self.hash_map.remove(&uuid).is_some(),
477 };
478 if existed {
479 self.queued_events.push(AssetEvent::Removed { id });
480 }
481 }
482
483 pub fn is_empty(&self) -> bool {
485 self.dense_storage.is_empty() && self.hash_map.is_empty()
486 }
487
488 pub fn len(&self) -> usize {
490 self.dense_storage.len() + self.hash_map.len()
491 }
492
493 pub fn ids(&self) -> impl Iterator<Item = AssetId<A>> + '_ {
495 self.dense_storage
496 .ids()
497 .chain(self.hash_map.keys().map(|uuid| AssetId::from(*uuid)))
498 }
499
500 pub fn iter(&self) -> impl Iterator<Item = (AssetId<A>, &A)> {
503 self.dense_storage
504 .storage
505 .iter()
506 .enumerate()
507 .filter_map(|(i, v)| match v {
508 Entry::None => None,
509 Entry::Some { value, generation } => value.as_ref().map(|v| {
510 let id = AssetId::Index {
511 index: AssetIndex {
512 generation: *generation,
513 index: i as u32,
514 },
515 marker: PhantomData,
516 };
517 (id, v)
518 }),
519 })
520 .chain(
521 self.hash_map
522 .iter()
523 .map(|(i, v)| (AssetId::Uuid { uuid: *i }, v)),
524 )
525 }
526
527 pub fn iter_mut(&mut self) -> AssetsMutIterator<'_, A> {
530 AssetsMutIterator {
531 dense_storage: self.dense_storage.storage.iter_mut().enumerate(),
532 hash_map: self.hash_map.iter_mut(),
533 queued_events: &mut self.queued_events,
534 }
535 }
536
537 pub fn track_assets(mut assets: ResMut<Self>, asset_server: Res<AssetServer>) {
540 let assets = &mut *assets;
541 let mut infos = asset_server.data.infos.write();
546 while let Ok(drop_event) = assets.handle_provider.drop_receiver.try_recv() {
547 let id = drop_event.id.typed();
548
549 if drop_event.asset_server_managed {
550 let untyped_id = id.untyped();
551
552 if !infos.process_handle_drop(untyped_id) {
554 continue;
556 }
557 }
558
559 assets.queued_events.push(AssetEvent::Unused { id });
560 assets.remove_dropped(id);
561 }
562 }
563
564 pub fn asset_events(mut assets: ResMut<Self>, mut events: EventWriter<AssetEvent<A>>) {
568 events.send_batch(assets.queued_events.drain(..));
569 }
570
571 pub(crate) fn asset_events_condition(assets: Res<Self>) -> bool {
576 !assets.queued_events.is_empty()
577 }
578}
579
580pub struct AssetsMutIterator<'a, A: Asset> {
582 queued_events: &'a mut Vec<AssetEvent<A>>,
583 dense_storage: Enumerate<std::slice::IterMut<'a, Entry<A>>>,
584 hash_map: bevy_utils::hashbrown::hash_map::IterMut<'a, Uuid, A>,
585}
586
587impl<'a, A: Asset> Iterator for AssetsMutIterator<'a, A> {
588 type Item = (AssetId<A>, &'a mut A);
589
590 fn next(&mut self) -> Option<Self::Item> {
591 for (i, entry) in &mut self.dense_storage {
592 match entry {
593 Entry::None => {
594 continue;
595 }
596 Entry::Some { value, generation } => {
597 let id = AssetId::Index {
598 index: AssetIndex {
599 generation: *generation,
600 index: i as u32,
601 },
602 marker: PhantomData,
603 };
604 self.queued_events.push(AssetEvent::Modified { id });
605 if let Some(value) = value {
606 return Some((id, value));
607 }
608 }
609 }
610 }
611 if let Some((key, value)) = self.hash_map.next() {
612 let id = AssetId::Uuid { uuid: *key };
613 self.queued_events.push(AssetEvent::Modified { id });
614 Some((id, value))
615 } else {
616 None
617 }
618 }
619}
620
621#[derive(Error, Debug)]
622#[error("AssetIndex {index:?} has an invalid generation. The current generation is: '{current_generation}'.")]
623pub struct InvalidGenerationError {
624 index: AssetIndex,
625 current_generation: u32,
626}
627
628#[cfg(test)]
629mod test {
630 use crate::AssetIndex;
631
632 #[test]
633 fn asset_index_round_trip() {
634 let asset_index = AssetIndex {
635 generation: 42,
636 index: 1337,
637 };
638 let roundtripped = AssetIndex::from_bits(asset_index.to_bits());
639 assert_eq!(asset_index, roundtripped);
640 }
641}