Send asset events for UUID asset IDs (duplicating their entity's asset events).

Author: andriyDev

crates/bevy_asset/src/assets.rs

@@ -15,7 +15,7 @@ use bevy_ecs::{
     world::{DeferredWorld, Mut, Ref, World},
 };
 use bevy_reflect::TypePath;
-use core::{marker::PhantomData, ops::Deref};
+use core::{any::TypeId, marker::PhantomData, ops::Deref};
 use derive_more::{Deref, DerefMut};
 use thiserror::Error;
 
@@ -220,6 +220,14 @@ fn write_added_asset_event<A: Asset>(
     if let Some(mut changes) = world.get_resource_mut::<AssetChanges<A>>() {
         changes.insert(entity, tick);
     }
+    if let Some(uuid_map) = world.get_resource::<AssetUuidMap>()
+        && let Some(uuids) = uuid_map.entity_to_uuids(entity, TypeId::of::<A>())
+    {
+        for uuid in uuids {
+            let id = AssetId::<A>::Uuid { uuid };
+            world.write_message(AssetEvent::<A>::Added { id });
+        }
+    }
 }
 
 /// Writes [`AssetEvent::Removed`] for this asset.
@@ -237,6 +245,14 @@ fn write_removed_asset_event<A: Asset>(
     if let Some(mut changes) = world.get_resource_mut::<AssetChanges<A>>() {
         changes.remove(&entity);
     }
+    if let Some(uuid_map) = world.get_resource::<AssetUuidMap>()
+        && let Some(uuids) = uuid_map.entity_to_uuids(entity, TypeId::of::<A>())
+    {
+        for uuid in uuids {
+            let id = AssetId::<A>::Uuid { uuid };
+            world.write_message(AssetEvent::<A>::Removed { id });
+        }
+    }
 }
 
 /// Writes [`AssetEvent::Modified`] messages for any assets that have changed.
@@ -246,6 +262,7 @@ pub(crate) fn write_modified_asset_events<A: Asset>(
     ticks: SystemChangeTick,
     mut messages: MessageWriter<AssetEvent<A>>,
     mut changes: Option<ResMut<AssetChanges<A>>>,
+    asset_uuid_map: Option<Res<AssetUuidMap>>,
 ) {
     for (entity, data_ref) in assets.iter() {
         let entity = AssetEntity::new_unchecked(entity);
@@ -265,6 +282,27 @@ pub(crate) fn write_modified_asset_events<A: Asset>(
             },
         });
     }
+
+    let Some(uuid_map) = asset_uuid_map else {
+        return;
+    };
+
+    for (_, inner) in uuid_map.read().iter() {
+        for (&uuid, handle) in inner.uuid_to_handle.iter() {
+            let Ok((_, data_ref)) = assets.get(handle.entity().raw_entity()) else {
+                continue;
+            };
+
+            if data_ref.last_changed() == data_ref.added() {
+                // The change corresponds to new asset data, which would have been handled by
+                // `AssetData`s hooks.
+                continue;
+            }
+            messages.write(AssetEvent::Modified {
+                id: AssetId::Uuid { uuid },
+            });
+        }
+    }
 }
 
 /// Despawns assets whose handles have been dropped and so are "unused".

crates/bevy_asset/src/direct_access_ext.rs

@@ -13,7 +13,7 @@ use uuid::Uuid;
 
 use crate::{
     insert_asset, meta::Settings, setup_asset, Asset, AssetData, AssetEntity,
-    AssetEntityDoesNotExistError, AssetHandleProvider, AssetId, AssetMut, AssetPath,
+    AssetEntityDoesNotExistError, AssetEvent, AssetHandleProvider, AssetId, AssetMut, AssetPath,
     AssetSelfHandle, AssetServer, AssetUuidMap, EntityHandle, Handle, ResolveUuidError,
     UntypedEntityHandle, UntypedHandle,
 };
@@ -89,6 +89,13 @@ impl DirectAssetAccessExt for World {
         let entity_handle = UntypedEntityHandle::from(EntityHandle::try_from(handle).unwrap());
         self.resource_mut::<AssetUuidMap>()
             .set_uuid(uuid, entity_handle);
+        // Send an asset event that this UUID has been added.
+        // TODO: This isn't quite sufficient, since users can call set_uuid themselves. We should be
+        // sending a message in that case as well. This also doesn't work correctly if the UUID
+        // was already set.
+        self.write_message(AssetEvent::<A>::Added {
+            id: AssetId::Uuid { uuid },
+        });
         Handle::Uuid(uuid, PhantomData)
     }
 
@@ -350,6 +357,13 @@ impl InternalAssetCommands<'_, '_, '_> {
             world
                 .resource_mut::<AssetUuidMap>()
                 .set_uuid(uuid, entity_handle);
+            // Send an asset event that this UUID has been added.
+            // TODO: This isn't quite sufficient, since users can call set_uuid themselves. We should be
+            // sending a message in that case as well. This also doesn't work correctly if the UUID
+            // was already set.
+            world.write_message(AssetEvent::<A>::Added {
+                id: AssetId::Uuid { uuid },
+            });
         });
         Handle::Uuid(uuid, PhantomData)
     }

crates/bevy_asset/src/uuid_map.rs

@@ -1,8 +1,10 @@
+use core::any::TypeId;
+
 use alloc::sync::Arc;
 
 use bevy_ecs::resource::Resource;
 use bevy_platform::{
-    collections::HashMap,
+    collections::{hash_map::Entry, HashMap, HashSet},
     sync::{PoisonError, RwLock, RwLockReadGuard},
 };
 use bevy_utils::TypeIdMap;
@@ -15,21 +17,43 @@ use crate::{
 
 /// Maps asset UUIDs to the asset handle assigned to it.
 #[derive(Resource, Clone, Default)]
-pub struct AssetUuidMap(Arc<RwLock<TypeIdMap<HashMap<Uuid, UntypedEntityHandle>>>>);
+pub struct AssetUuidMap(Arc<RwLock<TypeIdMap<AssetUuidMapInner>>>);
+
+#[derive(Default)]
+pub(crate) struct AssetUuidMapInner {
+    pub(crate) uuid_to_handle: HashMap<Uuid, UntypedEntityHandle>,
+    entity_to_uuids: HashMap<AssetEntity, HashSet<Uuid>>,
+}
 
 impl AssetUuidMap {
     /// Sets the handle that a UUID refers to.
     pub fn set_uuid(&mut self, uuid: Uuid, handle: UntypedEntityHandle) {
-        self.0
-            .write()
-            .unwrap_or_else(PoisonError::into_inner)
-            .entry(handle.0.type_id)
+        let mut type_id_map = self.0.write().unwrap_or_else(PoisonError::into_inner);
+        let inner = type_id_map.entry(handle.0.type_id).or_default();
+        let new_entity = handle.entity();
+        match inner.uuid_to_handle.entry(uuid) {
+            Entry::Vacant(entry) => {
+                entry.insert(handle);
+            }
+            Entry::Occupied(mut entry) => {
+                let old_entity = entry.get().entity();
+                inner
+                    .entity_to_uuids
+                    .get_mut(&old_entity)
+                    .unwrap()
+                    .remove(&uuid);
+                entry.insert(handle);
+            }
+        }
+        inner
+            .entity_to_uuids
+            .entry(new_entity)
             .or_default()
-            .insert(uuid, handle);
+            .insert(uuid);
     }
 
     /// Convenience function for accessing the internal uuid map.
-    fn read(&self) -> RwLockReadGuard<'_, TypeIdMap<HashMap<Uuid, UntypedEntityHandle>>> {
+    pub(crate) fn read(&self) -> RwLockReadGuard<'_, TypeIdMap<AssetUuidMapInner>> {
         self.0.read().unwrap_or_else(PoisonError::into_inner)
     }
 
@@ -46,7 +70,7 @@ impl AssetUuidMap {
             UntypedHandle::Uuid { type_id, uuid } => self
                 .read()
                 .get(&type_id)
-                .and_then(|map| map.get(&uuid))
+                .and_then(|inner| inner.uuid_to_handle.get(&uuid))
                 .cloned()
                 .ok_or(ResolveUuidError(uuid)),
         }
@@ -79,11 +103,26 @@ impl AssetUuidMap {
             UntypedAssetId::Uuid { type_id, uuid } => self
                 .read()
                 .get(&type_id)
-                .and_then(|map| map.get(&uuid))
+                .and_then(|inner| inner.uuid_to_handle.get(&uuid))
                 .map(|value| value.0.entity)
                 .ok_or(ResolveUuidError(uuid)),
         }
     }
+
+    /// Returns a reverse mapping from an entity to all UUIDs that reference it.
+    pub(crate) fn entity_to_uuids(
+        &self,
+        entity: AssetEntity,
+        type_id: TypeId,
+    ) -> Option<HashSet<Uuid>> {
+        Some(
+            self.read()
+                .get(&type_id)?
+                .entity_to_uuids
+                .get(&entity)?
+                .clone(),
+        )
+    }
 }
 
 /// An error while resolve a [`Uuid`] in the [`AssetUuidMap`].