bevy_asset/
assets.rs

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/// A generational runtime-only identifier for a specific [`Asset`] stored in [`Assets`]. This is optimized for efficient runtime
21/// usage and is not suitable for identifying assets across app runs.
22#[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    /// Convert the [`AssetIndex`] into an opaque blob of bits to transport it in circumstances where carrying a strongly typed index isn't possible.
32    ///
33    /// The result of this function should not be relied upon for anything except putting it back into [`AssetIndex::from_bits`] to recover the index.
34    pub fn to_bits(self) -> u64 {
35        let Self { generation, index } = self;
36        ((generation as u64) << 32) | index as u64
37    }
38    /// Convert an opaque `u64` acquired from [`AssetIndex::to_bits`] back into an [`AssetIndex`]. This should not be used with any inputs other than those
39    /// derived from [`AssetIndex::to_bits`], as there are no guarantees for what will happen with such inputs.
40    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
47/// Allocates generational [`AssetIndex`] values and facilitates their reuse.
48pub(crate) struct AssetIndexAllocator {
49    /// A monotonically increasing index.
50    next_index: AtomicU32,
51    recycled_queue_sender: Sender<AssetIndex>,
52    /// This receives every recycled [`AssetIndex`]. It serves as a buffer/queue to store indices ready for reuse.
53    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    /// Reserves a new [`AssetIndex`], either by reusing a recycled index (with an incremented generation), or by creating a new index
74    /// by incrementing the index counter for a given asset type `A`.
75    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    /// Queues the given `index` for reuse. This should only be done if the `index` is no longer being used.
91    pub fn recycle(&self, index: AssetIndex) {
92        self.recycled_queue_sender.send(index).unwrap();
93    }
94}
95
96/// A "loaded asset" containing the untyped handle for an asset stored in a given [`AssetPath`].
97///
98/// [`AssetPath`]: crate::AssetPath
99#[derive(Asset, TypePath)]
100pub struct LoadedUntypedAsset {
101    #[dependency]
102    pub handle: UntypedHandle,
103}
104
105// PERF: do we actually need this to be an enum? Can we just use an "invalid" generation instead
106#[derive(Default)]
107enum Entry<A: Asset> {
108    /// None is an indicator that this entry does not have live handles.
109    #[default]
110    None,
111    /// Some is an indicator that there is a live handle active for the entry at this [`AssetIndex`]
112    Some { value: Option<A>, generation: u32 },
113}
114
115/// Stores [`Asset`] values in a Vec-like storage identified by [`AssetIndex`].
116struct 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    // Returns the number of assets stored.
134    pub(crate) fn len(&self) -> usize {
135        self.len as usize
136    }
137
138    // Returns `true` if there are no assets stored.
139    pub(crate) fn is_empty(&self) -> bool {
140        self.len == 0
141    }
142
143    /// Insert the value at the given index. Returns true if a value already exists (and was replaced)
144    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    /// Removes the asset stored at the given `index` and returns it as [`Some`] (if the asset exists).
171    /// This will recycle the id and allow new entries to be inserted.
172    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    /// Removes the asset stored at the given `index` and returns it as [`Some`] (if the asset exists).
180    /// This will _not_ recycle the id. New values with the current ID can still be inserted. The ID will
181    /// not be reused until [`DenseAssetStorage::remove_dropped`] is called.
182    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        // NOTE: this assumes the allocator index is monotonically increasing.
239        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/// Stores [`Asset`] values identified by their [`AssetId`].
281///
282/// Assets identified by [`AssetId::Index`] will be stored in a "dense" vec-like storage. This is more efficient, but it means that
283/// the assets can only be identified at runtime. This is the default behavior.
284///
285/// Assets identified by [`AssetId::Uuid`] will be stored in a hashmap. This is less efficient, but it means that the assets can be referenced
286/// at compile time.
287///
288/// This tracks (and queues) [`AssetEvent`] events whenever changes to the collection occur.
289#[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    /// Assets managed by the `Assets` struct with live strong `Handle`s
296    /// originating from `get_strong_handle`.
297    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    /// Retrieves an [`AssetHandleProvider`] capable of reserving new [`Handle`] values for assets that will be stored in this
317    /// collection.
318    pub fn get_handle_provider(&self) -> AssetHandleProvider {
319        self.handle_provider.clone()
320    }
321
322    /// Reserves a new [`Handle`] for an asset that will be stored in this collection.
323    pub fn reserve_handle(&self) -> Handle<A> {
324        self.handle_provider.reserve_handle().typed::<A>()
325    }
326
327    /// Inserts the given `asset`, identified by the given `id`. If an asset already exists for `id`, it will be replaced.
328    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    /// Retrieves an [`Asset`] stored for the given `id` if it exists. If it does not exist, it will be inserted using `insert_fn`.
340    // PERF: Optimize this or remove it
341    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    /// Returns `true` if the `id` exists in this collection. Otherwise it returns `false`.
354    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    /// Adds the given `asset` and allocates a new strong [`Handle`] for it.
389    #[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    /// Upgrade an `AssetId` into a strong `Handle` that will prevent asset drop.
400    ///
401    /// Returns `None` if the provided `id` is not part of this `Assets` collection.
402    /// For example, it may have been dropped earlier.
403    #[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    /// Retrieves a reference to the [`Asset`] with the given `id`, if it exists.
419    /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
420    #[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    /// Retrieves a mutable reference to the [`Asset`] with the given `id`, if it exists.
429    /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
430    #[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    /// Removes (and returns) the [`Asset`] with the given `id`, if it exists.
444    /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
445    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    /// Removes (and returns) the [`Asset`] with the given `id`, if it exists. This skips emitting [`AssetEvent::Removed`].
455    /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
456    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    /// Removes the [`Asset`] with the given `id`.
466    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    /// Returns `true` if there are no assets in this collection.
484    pub fn is_empty(&self) -> bool {
485        self.dense_storage.is_empty() && self.hash_map.is_empty()
486    }
487
488    /// Returns the number of assets currently stored in the collection.
489    pub fn len(&self) -> usize {
490        self.dense_storage.len() + self.hash_map.len()
491    }
492
493    /// Returns an iterator over the [`AssetId`] of every [`Asset`] stored in this collection.
494    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    /// Returns an iterator over the [`AssetId`] and [`Asset`] ref of every asset in this collection.
501    // PERF: this could be accelerated if we implement a skip list. Consider the cost/benefits
502    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    /// Returns an iterator over the [`AssetId`] and mutable [`Asset`] ref of every asset in this collection.
528    // PERF: this could be accelerated if we implement a skip list. Consider the cost/benefits
529    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    /// A system that synchronizes the state of assets in this collection with the [`AssetServer`]. This manages
538    /// [`Handle`] drop events.
539    pub fn track_assets(mut assets: ResMut<Self>, asset_server: Res<AssetServer>) {
540        let assets = &mut *assets;
541        // note that we must hold this lock for the entire duration of this function to ensure
542        // that `asset_server.load` calls that occur during it block, which ensures that
543        // re-loads are kicked off appropriately. This function must be "transactional" relative
544        // to other asset info operations
545        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                // the process_handle_drop call checks whether new handles have been created since the drop event was fired, before removing the asset
553                if !infos.process_handle_drop(untyped_id) {
554                    // a new handle has been created, or the asset doesn't exist
555                    continue;
556                }
557            }
558
559            assets.queued_events.push(AssetEvent::Unused { id });
560            assets.remove_dropped(id);
561        }
562    }
563
564    /// A system that applies accumulated asset change events to the [`Events`] resource.
565    ///
566    /// [`Events`]: bevy_ecs::event::Events
567    pub fn asset_events(mut assets: ResMut<Self>, mut events: EventWriter<AssetEvent<A>>) {
568        events.send_batch(assets.queued_events.drain(..));
569    }
570
571    /// A run condition for [`asset_events`]. The system will not run if there are no events to
572    /// flush.
573    ///
574    /// [`asset_events`]: Self::asset_events
575    pub(crate) fn asset_events_condition(assets: Res<Self>) -> bool {
576        !assets.queued_events.is_empty()
577    }
578}
579
580/// A mutable iterator over [`Assets`].
581pub 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}