bevy_asset/
assets.rs

1use crate::asset_changed::AssetChanges;
2use crate::{Asset, AssetEvent, AssetHandleProvider, AssetId, AssetServer, Handle, UntypedHandle};
3use alloc::{sync::Arc, vec::Vec};
4use bevy_ecs::{
5    prelude::EventWriter,
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/// A generational runtime-only identifier for a specific [`Asset`] stored in [`Assets`]. This is optimized for efficient runtime
18/// usage and is not suitable for identifying assets across app runs.
19#[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    /// Convert the [`AssetIndex`] into an opaque blob of bits to transport it in circumstances where carrying a strongly typed index isn't possible.
29    ///
30    /// The result of this function should not be relied upon for anything except putting it back into [`AssetIndex::from_bits`] to recover the index.
31    pub fn to_bits(self) -> u64 {
32        let Self { generation, index } = self;
33        ((generation as u64) << 32) | index as u64
34    }
35    /// Convert an opaque `u64` acquired from [`AssetIndex::to_bits`] back into an [`AssetIndex`]. This should not be used with any inputs other than those
36    /// derived from [`AssetIndex::to_bits`], as there are no guarantees for what will happen with such inputs.
37    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
44/// Allocates generational [`AssetIndex`] values and facilitates their reuse.
45pub(crate) struct AssetIndexAllocator {
46    /// A monotonically increasing index.
47    next_index: AtomicU32,
48    recycled_queue_sender: Sender<AssetIndex>,
49    /// This receives every recycled [`AssetIndex`]. It serves as a buffer/queue to store indices ready for reuse.
50    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    /// Reserves a new [`AssetIndex`], either by reusing a recycled index (with an incremented generation), or by creating a new index
71    /// by incrementing the index counter for a given asset type `A`.
72    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    /// Queues the given `index` for reuse. This should only be done if the `index` is no longer being used.
88    pub fn recycle(&self, index: AssetIndex) {
89        self.recycled_queue_sender.send(index).unwrap();
90    }
91}
92
93/// A "loaded asset" containing the untyped handle for an asset stored in a given [`AssetPath`].
94///
95/// [`AssetPath`]: crate::AssetPath
96#[derive(Asset, TypePath)]
97pub struct LoadedUntypedAsset {
98    /// The handle to the loaded asset.
99    #[dependency]
100    pub handle: UntypedHandle,
101}
102
103// PERF: do we actually need this to be an enum? Can we just use an "invalid" generation instead
104#[derive(Default)]
105enum Entry<A: Asset> {
106    /// None is an indicator that this entry does not have live handles.
107    #[default]
108    None,
109    /// Some is an indicator that there is a live handle active for the entry at this [`AssetIndex`]
110    Some { value: Option<A>, generation: u32 },
111}
112
113/// Stores [`Asset`] values in a Vec-like storage identified by [`AssetIndex`].
114struct 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    // Returns the number of assets stored.
132    pub(crate) fn len(&self) -> usize {
133        self.len as usize
134    }
135
136    // Returns `true` if there are no assets stored.
137    pub(crate) fn is_empty(&self) -> bool {
138        self.len == 0
139    }
140
141    /// Insert the value at the given index. Returns true if a value already exists (and was replaced)
142    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 {
159                    index,
160                    current_generation: *generation,
161                })
162            }
163        } else {
164            unreachable!("entries should always be valid after a flush");
165        }
166    }
167
168    /// Removes the asset stored at the given `index` and returns it as [`Some`] (if the asset exists).
169    /// This will recycle the id and allow new entries to be inserted.
170    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    /// Removes the asset stored at the given `index` and returns it as [`Some`] (if the asset exists).
178    /// This will _not_ recycle the id. New values with the current ID can still be inserted. The ID will
179    /// not be reused until [`DenseAssetStorage::remove_dropped`] is called.
180    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        // NOTE: this assumes the allocator index is monotonically increasing.
234        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/// Stores [`Asset`] values identified by their [`AssetId`].
276///
277/// Assets identified by [`AssetId::Index`] will be stored in a "dense" vec-like storage. This is more efficient, but it means that
278/// the assets can only be identified at runtime. This is the default behavior.
279///
280/// Assets identified by [`AssetId::Uuid`] will be stored in a hashmap. This is less efficient, but it means that the assets can be referenced
281/// at compile time.
282///
283/// This tracks (and queues) [`AssetEvent`] events whenever changes to the collection occur.
284/// To check whether the asset used by a given component has changed (due to a change in the handle or the underlying asset)
285/// use the [`AssetChanged`](crate::asset_changed::AssetChanged) query filter.
286#[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    /// Assets managed by the `Assets` struct with live strong `Handle`s
293    /// originating from `get_strong_handle`.
294    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    /// Retrieves an [`AssetHandleProvider`] capable of reserving new [`Handle`] values for assets that will be stored in this
314    /// collection.
315    pub fn get_handle_provider(&self) -> AssetHandleProvider {
316        self.handle_provider.clone()
317    }
318
319    /// Reserves a new [`Handle`] for an asset that will be stored in this collection.
320    pub fn reserve_handle(&self) -> Handle<A> {
321        self.handle_provider.reserve_handle().typed::<A>()
322    }
323
324    /// Inserts the given `asset`, identified by the given `id`. If an asset already exists for `id`, it will be replaced.
325    pub fn insert(&mut self, id: impl Into<AssetId<A>>, asset: A) {
326        match id.into() {
327            AssetId::Index { index, .. } => {
328                self.insert_with_index(index, asset).unwrap();
329            }
330            AssetId::Uuid { uuid } => {
331                self.insert_with_uuid(uuid, asset);
332            }
333        }
334    }
335
336    /// Retrieves an [`Asset`] stored for the given `id` if it exists. If it does not exist, it will be inserted using `insert_fn`.
337    // PERF: Optimize this or remove it
338    pub fn get_or_insert_with(
339        &mut self,
340        id: impl Into<AssetId<A>>,
341        insert_fn: impl FnOnce() -> A,
342    ) -> &mut A {
343        let id: AssetId<A> = id.into();
344        if self.get(id).is_none() {
345            self.insert(id, insert_fn());
346        }
347        self.get_mut(id).unwrap()
348    }
349
350    /// Returns `true` if the `id` exists in this collection. Otherwise it returns `false`.
351    pub fn contains(&self, id: impl Into<AssetId<A>>) -> bool {
352        match id.into() {
353            AssetId::Index { index, .. } => self.dense_storage.get(index).is_some(),
354            AssetId::Uuid { uuid } => self.hash_map.contains_key(&uuid),
355        }
356    }
357
358    pub(crate) fn insert_with_uuid(&mut self, uuid: Uuid, asset: A) -> Option<A> {
359        let result = self.hash_map.insert(uuid, asset);
360        if result.is_some() {
361            self.queued_events
362                .push(AssetEvent::Modified { id: uuid.into() });
363        } else {
364            self.queued_events
365                .push(AssetEvent::Added { id: uuid.into() });
366        }
367        result
368    }
369    pub(crate) fn insert_with_index(
370        &mut self,
371        index: AssetIndex,
372        asset: A,
373    ) -> Result<bool, InvalidGenerationError> {
374        let replaced = self.dense_storage.insert(index, asset)?;
375        if replaced {
376            self.queued_events
377                .push(AssetEvent::Modified { id: index.into() });
378        } else {
379            self.queued_events
380                .push(AssetEvent::Added { id: index.into() });
381        }
382        Ok(replaced)
383    }
384
385    /// Adds the given `asset` and allocates a new strong [`Handle`] for it.
386    #[inline]
387    pub fn add(&mut self, asset: impl Into<A>) -> Handle<A> {
388        let index = self.dense_storage.allocator.reserve();
389        self.insert_with_index(index, asset.into()).unwrap();
390        Handle::Strong(
391            self.handle_provider
392                .get_handle(index.into(), false, None, None),
393        )
394    }
395
396    /// Upgrade an `AssetId` into a strong `Handle` that will prevent asset drop.
397    ///
398    /// Returns `None` if the provided `id` is not part of this `Assets` collection.
399    /// For example, it may have been dropped earlier.
400    #[inline]
401    pub fn get_strong_handle(&mut self, id: AssetId<A>) -> Option<Handle<A>> {
402        if !self.contains(id) {
403            return None;
404        }
405        *self.duplicate_handles.entry(id).or_insert(0) += 1;
406        let index = match id {
407            AssetId::Index { index, .. } => index.into(),
408            AssetId::Uuid { uuid } => uuid.into(),
409        };
410        Some(Handle::Strong(
411            self.handle_provider.get_handle(index, false, None, None),
412        ))
413    }
414
415    /// Retrieves a reference to the [`Asset`] with the given `id`, if it exists.
416    /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
417    #[inline]
418    pub fn get(&self, id: impl Into<AssetId<A>>) -> Option<&A> {
419        match id.into() {
420            AssetId::Index { index, .. } => self.dense_storage.get(index),
421            AssetId::Uuid { uuid } => self.hash_map.get(&uuid),
422        }
423    }
424
425    /// Retrieves a mutable reference to the [`Asset`] with the given `id`, if it exists.
426    /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
427    #[inline]
428    pub fn get_mut(&mut self, id: impl Into<AssetId<A>>) -> Option<&mut A> {
429        let id: AssetId<A> = id.into();
430        let result = match id {
431            AssetId::Index { index, .. } => self.dense_storage.get_mut(index),
432            AssetId::Uuid { uuid } => self.hash_map.get_mut(&uuid),
433        };
434        if result.is_some() {
435            self.queued_events.push(AssetEvent::Modified { id });
436        }
437        result
438    }
439
440    /// Removes (and returns) the [`Asset`] with the given `id`, if it exists.
441    /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
442    pub fn remove(&mut self, id: impl Into<AssetId<A>>) -> Option<A> {
443        let id: AssetId<A> = id.into();
444        let result = self.remove_untracked(id);
445        if result.is_some() {
446            self.queued_events.push(AssetEvent::Removed { id });
447        }
448        result
449    }
450
451    /// Removes (and returns) the [`Asset`] with the given `id`, if it exists. This skips emitting [`AssetEvent::Removed`].
452    /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
453    pub fn remove_untracked(&mut self, id: impl Into<AssetId<A>>) -> Option<A> {
454        let id: AssetId<A> = id.into();
455        self.duplicate_handles.remove(&id);
456        match id {
457            AssetId::Index { index, .. } => self.dense_storage.remove_still_alive(index),
458            AssetId::Uuid { uuid } => self.hash_map.remove(&uuid),
459        }
460    }
461
462    /// Removes the [`Asset`] with the given `id`.
463    pub(crate) fn remove_dropped(&mut self, id: AssetId<A>) {
464        match self.duplicate_handles.get_mut(&id) {
465            None => {}
466            Some(0) => {
467                self.duplicate_handles.remove(&id);
468            }
469            Some(value) => {
470                *value -= 1;
471                return;
472            }
473        }
474
475        let existed = match id {
476            AssetId::Index { index, .. } => self.dense_storage.remove_dropped(index).is_some(),
477            AssetId::Uuid { uuid } => self.hash_map.remove(&uuid).is_some(),
478        };
479
480        self.queued_events.push(AssetEvent::Unused { id });
481        if existed {
482            self.queued_events.push(AssetEvent::Removed { id });
483        }
484    }
485
486    /// Returns `true` if there are no assets in this collection.
487    pub fn is_empty(&self) -> bool {
488        self.dense_storage.is_empty() && self.hash_map.is_empty()
489    }
490
491    /// Returns the number of assets currently stored in the collection.
492    pub fn len(&self) -> usize {
493        self.dense_storage.len() + self.hash_map.len()
494    }
495
496    /// Returns an iterator over the [`AssetId`] of every [`Asset`] stored in this collection.
497    pub fn ids(&self) -> impl Iterator<Item = AssetId<A>> + '_ {
498        self.dense_storage
499            .ids()
500            .chain(self.hash_map.keys().map(|uuid| AssetId::from(*uuid)))
501    }
502
503    /// Returns an iterator over the [`AssetId`] and [`Asset`] ref of every asset in this collection.
504    // PERF: this could be accelerated if we implement a skip list. Consider the cost/benefits
505    pub fn iter(&self) -> impl Iterator<Item = (AssetId<A>, &A)> {
506        self.dense_storage
507            .storage
508            .iter()
509            .enumerate()
510            .filter_map(|(i, v)| match v {
511                Entry::None => None,
512                Entry::Some { value, generation } => value.as_ref().map(|v| {
513                    let id = AssetId::Index {
514                        index: AssetIndex {
515                            generation: *generation,
516                            index: i as u32,
517                        },
518                        marker: PhantomData,
519                    };
520                    (id, v)
521                }),
522            })
523            .chain(
524                self.hash_map
525                    .iter()
526                    .map(|(i, v)| (AssetId::Uuid { uuid: *i }, v)),
527            )
528    }
529
530    /// Returns an iterator over the [`AssetId`] and mutable [`Asset`] ref of every asset in this collection.
531    // PERF: this could be accelerated if we implement a skip list. Consider the cost/benefits
532    pub fn iter_mut(&mut self) -> AssetsMutIterator<'_, A> {
533        AssetsMutIterator {
534            dense_storage: self.dense_storage.storage.iter_mut().enumerate(),
535            hash_map: self.hash_map.iter_mut(),
536            queued_events: &mut self.queued_events,
537        }
538    }
539
540    /// A system that synchronizes the state of assets in this collection with the [`AssetServer`]. This manages
541    /// [`Handle`] drop events.
542    pub fn track_assets(mut assets: ResMut<Self>, asset_server: Res<AssetServer>) {
543        let assets = &mut *assets;
544        // note that we must hold this lock for the entire duration of this function to ensure
545        // that `asset_server.load` calls that occur during it block, which ensures that
546        // re-loads are kicked off appropriately. This function must be "transactional" relative
547        // to other asset info operations
548        let mut infos = asset_server.data.infos.write();
549        while let Ok(drop_event) = assets.handle_provider.drop_receiver.try_recv() {
550            let id = drop_event.id.typed();
551
552            if drop_event.asset_server_managed {
553                let untyped_id = id.untyped();
554
555                // the process_handle_drop call checks whether new handles have been created since the drop event was fired, before removing the asset
556                if !infos.process_handle_drop(untyped_id) {
557                    // a new handle has been created, or the asset doesn't exist
558                    continue;
559                }
560            }
561
562            assets.remove_dropped(id);
563        }
564    }
565
566    /// A system that applies accumulated asset change events to the [`Events`] resource.
567    ///
568    /// [`Events`]: bevy_ecs::event::Events
569    pub(crate) fn asset_events(
570        mut assets: ResMut<Self>,
571        mut events: EventWriter<AssetEvent<A>>,
572        asset_changes: Option<ResMut<AssetChanges<A>>>,
573        ticks: SystemChangeTick,
574    ) {
575        use AssetEvent::{Added, LoadedWithDependencies, Modified, Removed};
576
577        if let Some(mut asset_changes) = asset_changes {
578            for new_event in &assets.queued_events {
579                match new_event {
580                    Removed { id } | AssetEvent::Unused { id } => asset_changes.remove(id),
581                    Added { id } | Modified { id } | LoadedWithDependencies { id } => {
582                        asset_changes.insert(*id, ticks.this_run());
583                    }
584                };
585            }
586        }
587        events.write_batch(assets.queued_events.drain(..));
588    }
589
590    /// A run condition for [`asset_events`]. The system will not run if there are no events to
591    /// flush.
592    ///
593    /// [`asset_events`]: Self::asset_events
594    pub(crate) fn asset_events_condition(assets: Res<Self>) -> bool {
595        !assets.queued_events.is_empty()
596    }
597}
598
599/// A mutable iterator over [`Assets`].
600pub struct AssetsMutIterator<'a, A: Asset> {
601    queued_events: &'a mut Vec<AssetEvent<A>>,
602    dense_storage: Enumerate<core::slice::IterMut<'a, Entry<A>>>,
603    hash_map: bevy_platform::collections::hash_map::IterMut<'a, Uuid, A>,
604}
605
606impl<'a, A: Asset> Iterator for AssetsMutIterator<'a, A> {
607    type Item = (AssetId<A>, &'a mut A);
608
609    fn next(&mut self) -> Option<Self::Item> {
610        for (i, entry) in &mut self.dense_storage {
611            match entry {
612                Entry::None => {
613                    continue;
614                }
615                Entry::Some { value, generation } => {
616                    let id = AssetId::Index {
617                        index: AssetIndex {
618                            generation: *generation,
619                            index: i as u32,
620                        },
621                        marker: PhantomData,
622                    };
623                    self.queued_events.push(AssetEvent::Modified { id });
624                    if let Some(value) = value {
625                        return Some((id, value));
626                    }
627                }
628            }
629        }
630        if let Some((key, value)) = self.hash_map.next() {
631            let id = AssetId::Uuid { uuid: *key };
632            self.queued_events.push(AssetEvent::Modified { id });
633            Some((id, value))
634        } else {
635            None
636        }
637    }
638}
639
640/// An error returned when an [`AssetIndex`] has an invalid generation.
641#[derive(Error, Debug)]
642#[error("AssetIndex {index:?} has an invalid generation. The current generation is: '{current_generation}'.")]
643pub struct InvalidGenerationError {
644    index: AssetIndex,
645    current_generation: u32,
646}
647
648#[cfg(test)]
649mod test {
650    use crate::AssetIndex;
651
652    #[test]
653    fn asset_index_round_trip() {
654        let asset_index = AssetIndex {
655            generation: 42,
656            index: 1337,
657        };
658        let roundtripped = AssetIndex::from_bits(asset_index.to_bits());
659        assert_eq!(asset_index, roundtripped);
660    }
661}