Add a way to get a strong handle from an AssetId (bevyengine#12088)

# Objective

- Sometimes, it is useful to get a `Handle<T>` from an `AssetId<T>`. For
example, when iterating `Assets` to find a specific asset. So much so
that it's possible to do so with `AssetServer::get_id_handle`
- However, `AssetServer::get_id_handle` doesn't work with assets
directly added to `Assets` using `Assets::add`.
- Fixes bevyengine#12087

## Solution

- Add `Assets::get_strong_handle` to convert an `AssetId` into an
`Handle`
- Document `AssetServer::get_id_handle` to explain its limitation and
point to `get_strong_handle`.
- Add a test for `get_strong_handle`
- Add a `duplicate_handles` field to `Assets` to avoid dropping assets
with a live handle generated by `get_strong_handle` (more reference
counting yay…)
- Fix typos in `Assets` docs

---

## Changelog

- Add `Assets::get_strong_handle` to convert an `AssetId` into an
`Handle`

crates/bevy_asset/src/assets.rs

@@ -276,6 +276,9 @@ pub struct Assets<A: Asset> {
     hash_map: HashMap<Uuid, A>,
     handle_provider: AssetHandleProvider,
     queued_events: Vec<AssetEvent<A>>,
+    /// Assets managed by the `Assets` struct with live strong `Handle`s
+    /// originating from `get_strong_handle`.
+    duplicate_handles: HashMap<AssetId<A>, u16>,
 }
 
 impl<A: Asset> Default for Assets<A> {
@@ -288,6 +291,7 @@ impl<A: Asset> Default for Assets<A> {
             handle_provider,
             hash_map: Default::default(),
             queued_events: Default::default(),
+            duplicate_handles: Default::default(),
         }
     }
 }
@@ -306,8 +310,7 @@ impl<A: Asset> Assets<A> {
 
     /// Inserts the given `asset`, identified by the given `id`. If an asset already exists for `id`, it will be replaced.
     pub fn insert(&mut self, id: impl Into<AssetId<A>>, asset: A) {
-        let id: AssetId<A> = id.into();
-        match id {
+        match id.into() {
             AssetId::Index { index, .. } => {
                 self.insert_with_index(index, asset).unwrap();
             }
@@ -332,9 +335,11 @@ impl<A: Asset> Assets<A> {
     }
 
     /// Returns `true` if the `id` exists in this collection. Otherwise it returns `false`.
-    // PERF: Optimize this or remove it
     pub fn contains(&self, id: impl Into<AssetId<A>>) -> bool {
-        self.get(id).is_some()
+        match id.into() {
+            AssetId::Index { index, .. } => self.dense_storage.get(index).is_some(),
+            AssetId::Uuid { uuid } => self.hash_map.contains_key(&uuid),
+        }
     }
 
     pub(crate) fn insert_with_uuid(&mut self, uuid: Uuid, asset: A) -> Option<A> {
@@ -375,18 +380,36 @@ impl<A: Asset> Assets<A> {
         )
     }
 
-    /// Retrieves a reference to the [`Asset`] with the given `id`, if its exists.
+    /// Upgrade an `AssetId` into a strong `Handle` that will prevent asset drop.
+    ///
+    /// Returns `None` if the provided `id` is not part of this `Assets` collection.
+    /// For example, it may have been dropped earlier.
+    #[inline]
+    pub fn get_strong_handle(&mut self, id: AssetId<A>) -> Option<Handle<A>> {
+        if !self.contains(id) {
+            return None;
+        }
+        *self.duplicate_handles.entry(id).or_insert(0) += 1;
+        let index = match id {
+            AssetId::Index { index, .. } => index.into(),
+            AssetId::Uuid { uuid } => uuid.into(),
+        };
+        Some(Handle::Strong(
+            self.handle_provider.get_handle(index, false, None, None),
+        ))
+    }
+
+    /// Retrieves a reference to the [`Asset`] with the given `id`, if it exists.
     /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
     #[inline]
     pub fn get(&self, id: impl Into<AssetId<A>>) -> Option<&A> {
-        let id: AssetId<A> = id.into();
-        match id {
+        match id.into() {
             AssetId::Index { index, .. } => self.dense_storage.get(index),
             AssetId::Uuid { uuid } => self.hash_map.get(&uuid),
         }
     }
 
-    /// Retrieves a mutable reference to the [`Asset`] with the given `id`, if its exists.
+    /// Retrieves a mutable reference to the [`Asset`] with the given `id`, if it exists.
     /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
     #[inline]
     pub fn get_mut(&mut self, id: impl Into<AssetId<A>>) -> Option<&mut A> {
@@ -401,7 +424,7 @@ impl<A: Asset> Assets<A> {
         result
     }
 
-    /// Removes (and returns) the [`Asset`] with the given `id`, if its exists.
+    /// Removes (and returns) the [`Asset`] with the given `id`, if it exists.
     /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
     pub fn remove(&mut self, id: impl Into<AssetId<A>>) -> Option<A> {
         let id: AssetId<A> = id.into();
@@ -412,28 +435,33 @@ impl<A: Asset> Assets<A> {
         result
     }
 
-    /// Removes (and returns) the [`Asset`] with the given `id`, if its exists. This skips emitting [`AssetEvent::Removed`].
+    /// Removes (and returns) the [`Asset`] with the given `id`, if it exists. This skips emitting [`AssetEvent::Removed`].
     /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
     pub fn remove_untracked(&mut self, id: impl Into<AssetId<A>>) -> Option<A> {
         let id: AssetId<A> = id.into();
+        self.duplicate_handles.remove(&id);
         match id {
             AssetId::Index { index, .. } => self.dense_storage.remove_still_alive(index),
             AssetId::Uuid { uuid } => self.hash_map.remove(&uuid),
         }
     }
 
-    /// Removes (and returns) the [`Asset`] with the given `id`, if its exists.
-    /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
-    pub(crate) fn remove_dropped(&mut self, id: impl Into<AssetId<A>>) -> Option<A> {
-        let id: AssetId<A> = id.into();
-        let result = match id {
-            AssetId::Index { index, .. } => self.dense_storage.remove_dropped(index),
-            AssetId::Uuid { uuid } => self.hash_map.remove(&uuid),
+    /// Removes the [`Asset`] with the given `id`.
+    pub(crate) fn remove_dropped(&mut self, id: AssetId<A>) {
+        match self.duplicate_handles.get_mut(&id) {
+            None | Some(0) => {}
+            Some(value) => {
+                *value -= 1;
+                return;
+            }
+        }
+        let existed = match id {
+            AssetId::Index { index, .. } => self.dense_storage.remove_dropped(index).is_some(),
+            AssetId::Uuid { uuid } => self.hash_map.remove(&uuid).is_some(),
         };
-        if result.is_some() {
+        if existed {
             self.queued_events.push(AssetEvent::Removed { id });
         }
-        result
     }
 
     /// Returns `true` if there are no assets in this collection.
@@ -505,20 +533,19 @@ impl<A: Asset> Assets<A> {
 
             assets.queued_events.push(AssetEvent::Unused { id });
 
+            let mut remove_asset = true;
+
             if drop_event.asset_server_managed {
                 let untyped_id = drop_event.id.untyped(TypeId::of::<A>());
                 if let Some(info) = infos.get(untyped_id) {
-                    if info.load_state == LoadState::Loading
-                        || info.load_state == LoadState::NotLoaded
-                    {
+                    if let LoadState::Loading | LoadState::NotLoaded = info.load_state {
                         not_ready.push(drop_event);
                         continue;
                     }
                 }
-                if infos.process_handle_drop(untyped_id) {
-                    assets.remove_dropped(id);
-                }
-            } else {
+                remove_asset = infos.process_handle_drop(untyped_id);
+            }
+            if remove_asset {
                 assets.remove_dropped(id);
             }
         }

crates/bevy_asset/src/handle.rs

@@ -23,6 +23,7 @@ pub struct AssetHandleProvider {
     pub(crate) type_id: TypeId,
 }
 
+#[derive(Debug)]
 pub(crate) struct DropEvent {
     pub(crate) id: InternalAssetId,
     pub(crate) asset_server_managed: bool,
@@ -507,13 +508,13 @@ impl<A: Asset> TryFrom<UntypedHandle> for Handle<A> {
     }
 }
 
-/// Errors preventing the conversion of to/from an [`UntypedHandle`] and an [`Handle`].
+/// Errors preventing the conversion of to/from an [`UntypedHandle`] and a [`Handle`].
 #[derive(Error, Debug, PartialEq, Clone)]
 #[non_exhaustive]
 pub enum UntypedAssetConversionError {
-    /// Caused when trying to convert an [`UntypedHandle`] into an [`Handle`] of the wrong type.
+    /// Caused when trying to convert an [`UntypedHandle`] into a [`Handle`] of the wrong type.
     #[error(
-        "This UntypedHandle is for {found:?} and cannot be converted into an Handle<{expected:?}>"
+        "This UntypedHandle is for {found:?} and cannot be converted into a Handle<{expected:?}>"
     )]
     TypeIdMismatch { expected: TypeId, found: TypeId },
 }

crates/bevy_asset/src/lib.rs

@@ -466,7 +466,7 @@ mod tests {
     };
     use thiserror::Error;
 
-    #[derive(Asset, TypePath, Debug)]
+    #[derive(Asset, TypePath, Debug, Default)]
     pub struct CoolText {
         pub text: String,
         pub embedded: String,
@@ -1114,24 +1114,58 @@ mod tests {
         });
     }
 
+    const SIMPLE_TEXT: &str = r#"
+(
+    text: "dep",
+    dependencies: [],
+    embedded_dependencies: [],
+    sub_texts: [],
+)"#;
+    #[test]
+    fn keep_gotten_strong_handles() {
+        let dir = Dir::default();
+        dir.insert_asset_text(Path::new("dep.cool.ron"), SIMPLE_TEXT);
+
+        let (mut app, _) = test_app(dir);
+        app.init_asset::<CoolText>()
+            .init_asset::<SubText>()
+            .init_resource::<StoredEvents>()
+            .register_asset_loader(CoolTextLoader)
+            .add_systems(Update, store_asset_events);
+
+        let id = {
+            let handle = {
+                let mut texts = app.world.resource_mut::<Assets<CoolText>>();
+                let handle = texts.add(CoolText::default());
+                texts.get_strong_handle(handle.id()).unwrap()
+            };
+
+            app.update();
+
+            {
+                let text = app.world.resource::<Assets<CoolText>>().get(&handle);
+                assert!(text.is_some());
+            }
+            handle.id()
+        };
+        // handle is dropped
+        app.update();
+        assert!(
+            app.world.resource::<Assets<CoolText>>().get(id).is_none(),
+            "asset has no handles, so it should have been dropped last update"
+        );
+    }
+
     #[test]
     fn manual_asset_management() {
         // The particular usage of GatedReader in this test will cause deadlocking if running single-threaded
         #[cfg(not(feature = "multi-threaded"))]
         panic!("This test requires the \"multi-threaded\" feature, otherwise it will deadlock.\ncargo test --package bevy_asset --features multi-threaded");
 
         let dir = Dir::default();
-
         let dep_path = "dep.cool.ron";
-        let dep_ron = r#"
-(
-    text: "dep",
-    dependencies: [],
-    embedded_dependencies: [],
-    sub_texts: [],
-)"#;
 
-        dir.insert_asset_text(Path::new(dep_path), dep_ron);
+        dir.insert_asset_text(Path::new(dep_path), SIMPLE_TEXT);
 
         let (mut app, gate_opener) = test_app(dir);
         app.init_asset::<CoolText>()

crates/bevy_asset/src/server/info.rs

@@ -358,7 +358,7 @@ impl AssetInfos {
         }
     }
 
-    // Returns `true` if the asset should be removed from the collection
+    /// Returns `true` if the asset should be removed from the collection.
     pub(crate) fn process_handle_drop(&mut self, id: UntypedAssetId) -> bool {
         Self::process_handle_drop_internal(
             &mut self.infos,

crates/bevy_asset/src/server/mod.rs

@@ -797,16 +797,25 @@ impl AssetServer {
             .map(|h| h.typed_debug_checked())
     }
 
+    /// Get a `Handle` from an `AssetId`.
+    ///
+    /// This only returns `Some` if `id` is derived from a `Handle` that was
+    /// loaded through an `AssetServer`, otherwise it returns `None`.
+    ///
+    /// Consider using [`Assets::get_strong_handle`] in the case the `Handle`
+    /// comes from [`Assets::add`].
     pub fn get_id_handle<A: Asset>(&self, id: AssetId<A>) -> Option<Handle<A>> {
         self.get_id_handle_untyped(id.untyped()).map(|h| h.typed())
     }
 
+    /// Get an `UntypedHandle` from an `UntypedAssetId`.
+    /// See [`AssetServer::get_id_handle`] for details.
     pub fn get_id_handle_untyped(&self, id: UntypedAssetId) -> Option<UntypedHandle> {
         self.data.infos.read().get_id_handle(id)
     }
 
     /// Returns `true` if the given `id` corresponds to an asset that is managed by this [`AssetServer`].
-    /// Otherwise, returns false.
+    /// Otherwise, returns `false`.
     pub fn is_managed(&self, id: impl Into<UntypedAssetId>) -> bool {
         self.data.infos.read().contains_key(id.into())
     }