Change RenderAsset::unload_asset to be called whenever RenderAsset would be dropped

crates/bevy_pbr/src/render/mesh.rs

@@ -43,7 +43,7 @@ use bevy_render::batching::gpu_preprocessing::PreviousInstanceInputUniformBuffer
 use bevy_render::impl_atomic_pod;
 use bevy_render::mesh::allocator::{MeshSlabId, MeshSlabs};
 use bevy_render::mesh::morph::{
-    MorphTargetImage, MorphTargetsResource, RenderMorphTargetAllocator,
+    MorphTargetImages, MorphTargetsResource, RenderMorphTargetAllocator,
 };
 use bevy_render::{
     batching::{
@@ -3964,7 +3964,9 @@ fn prepare_mesh_bind_groups_for_phase(
     if weights_uniform.current_buffer.buffer().is_some() {
         match (render_morph_target_allocator, &mut groups.morph_targets) {
             (
-                RenderMorphTargetAllocator::Image { mesh_id_to_image },
+                RenderMorphTargetAllocator::Image {
+                    morph_target_images,
+                },
                 &mut MeshMorphTargetBindGroups::Uniform(ref mut morph_targets),
             ) => {
                 prepare_mesh_morph_target_bind_groups_for_phase_using_uniforms(
@@ -3975,7 +3977,7 @@ fn prepare_mesh_bind_groups_for_phase(
                     pipeline_cache,
                     skins_uniform,
                     weights_uniform,
-                    mesh_id_to_image,
+                    morph_target_images,
                     morph_targets,
                 );
             }
@@ -4034,7 +4036,7 @@ fn prepare_mesh_morph_target_bind_groups_for_phase_using_uniforms(
     pipeline_cache: &PipelineCache,
     skins_uniform: &SkinUniforms,
     weights_uniform: &MorphUniforms,
-    mesh_id_to_image: &HashMap<AssetId<Mesh>, MorphTargetImage>,
+    morph_target_images: &MorphTargetImages,
     morph_targets: &mut HashMap<AssetId<Mesh>, MeshBindGroupPair>,
 ) {
     let (skin, prev_skin) = (&skins_uniform.current_buffer, &skins_uniform.prev_buffer);
@@ -4049,12 +4051,11 @@ fn prepare_mesh_morph_target_bind_groups_for_phase_using_uniforms(
         .and_then(|descriptors_buffer| descriptors_buffer.buffer());
 
     for (id, gpu_mesh) in meshes.iter() {
-        if !gpu_mesh.has_morph_targets() {
-            continue;
-        }
-        let Some(morph_targets_image) = mesh_id_to_image.get(&id) else {
+        let Some(handle) = &gpu_mesh.morph_target_handle else {
             continue;
         };
+
+        let morph_targets_image = morph_target_images.get(handle);
         let targets = MorphTargetsResource::Texture(&morph_targets_image.texture_view);
         let bind_group_pair = if is_skinned(&gpu_mesh.layout) {
             MeshBindGroupPair {

crates/bevy_render/Cargo.toml

@@ -128,6 +128,7 @@ indexmap = { version = "2" }
 bitflags = "2"
 itertools = "0.14"
 weak-table = "0.3"
+slotmap = "1.1.1"
 
 [dev-dependencies]
 proptest = "1"

crates/bevy_render/src/mesh/mod.rs

@@ -26,7 +26,7 @@ use glam::Vec3;
 use wgpu::IndexFormat;
 
 #[cfg(feature = "morph")]
-use crate::mesh::morph::RenderMorphTargetAllocator;
+use crate::mesh::morph::{MorphTargetImageHandle, RenderMorphTargetAllocator};
 
 /// Makes sure that [`Mesh`]es are extracted and prepared for the GPU.
 /// Does *not* add the [`Mesh`] as an asset. Use [`MeshPlugin`] for that.
@@ -57,7 +57,7 @@ impl Plugin for MeshRenderAssetPlugin {
 }
 
 /// The render world representation of a [`Mesh`].
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct RenderMesh {
     /// The number of vertices in the mesh.
     pub vertex_count: u32,
@@ -77,6 +77,9 @@ pub struct RenderMesh {
     /// Combined with [`RenderMesh::buffer_info`], this specifies the complete
     /// layout of the buffers associated with this mesh.
     pub layout: MeshVertexBufferLayoutRef,
+
+    #[cfg(feature = "morph")]
+    pub morph_target_handle: Option<MorphTargetImageHandle>,
 }
 
 impl RenderMesh {
@@ -200,15 +203,14 @@ impl RenderAsset for RenderMesh {
 
         // Place the morph displacements in an image if necessary.
         #[cfg(feature = "morph")]
-        if let Some(morph_targets) = mesh.morph_targets() {
+        let morph_target_handle = mesh.morph_targets().and_then(|morph_targets| {
             _render_morph_targets_allocator.allocate(
                 _render_device,
                 _render_queue,
-                _mesh_id,
                 morph_targets,
                 mesh.count_vertices(),
-            );
-        }
+            )
+        });
 
         Ok(RenderMesh {
             vertex_count: mesh.count_vertices() as u32,
@@ -219,15 +221,20 @@ impl RenderAsset for RenderMesh {
             buffer_info,
             key_bits,
             layout: mesh_vertex_buffer_layout,
+            #[cfg(feature = "morph")]
+            morph_target_handle,
         })
     }
 
     fn unload_asset(
+        self,
         _mesh_id: AssetId<Self::SourceAsset>,
         (_, _, _, _render_morph_targets_allocator): &mut SystemParamItem<Self::Param>,
     ) {
         // Free the morph target images if necessary.
         #[cfg(feature = "morph")]
-        _render_morph_targets_allocator.free(_mesh_id);
+        if let Some(handle) = self.morph_target_handle {
+            _render_morph_targets_allocator.free(handle);
+        }
     }
 }

crates/bevy_render/src/mesh/morph.rs

@@ -1,14 +1,10 @@
-use bevy_asset::AssetId;
 use bevy_ecs::{
     resource::Resource,
     world::{FromWorld, World},
 };
 use bevy_log::error;
-use bevy_mesh::{
-    morph::{MorphAttributes, MorphBuildError, MAX_MORPH_WEIGHTS, MAX_TEXTURE_WIDTH},
-    Mesh,
-};
-use bevy_platform::collections::HashMap;
+use bevy_mesh::morph::{MorphAttributes, MorphBuildError, MAX_MORPH_WEIGHTS, MAX_TEXTURE_WIDTH};
+use slotmap::SlotMap;
 use wgpu::{
     Extent3d, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
     TextureViewDescriptor,
@@ -111,19 +107,45 @@ impl MorphTargetImage {
     }
 }
 
+slotmap::new_key_type! {
+    pub struct MorphTargetImageKey;
+}
+
+#[derive(Debug)]
+pub struct MorphTargetImageHandle(MorphTargetImageKey);
+
+#[derive(Debug, Clone, Default)]
+pub struct MorphTargetImages(SlotMap<MorphTargetImageKey, MorphTargetImage>);
+
+impl MorphTargetImages {
+    pub fn insert(&mut self, morph_target_image: MorphTargetImage) -> MorphTargetImageHandle {
+        MorphTargetImageHandle(self.0.insert(morph_target_image))
+    }
+
+    pub fn remove(&mut self, handle: MorphTargetImageHandle) {
+        self.0
+            .remove(handle.0)
+            .expect("MorphTargetImageHandle cannot double-free");
+    }
+
+    pub fn get(&self, handle: &MorphTargetImageHandle) -> &MorphTargetImage {
+        self.0
+            .get(handle.0)
+            .expect("MorphTargetImageHandle cannot be use-after-free")
+    }
+}
+
 /// Stores the images for all morph target displacement data, if the current
 /// platform doesn't support storage buffers.
 ///
 /// If the current platform does support storage buffers, the mesh allocator
 /// stores displacement data instead.
-#[derive(Resource)]
+#[derive(Debug, Resource)]
 pub enum RenderMorphTargetAllocator {
     /// The variant used when the current platform doesn't support storage
     /// buffers.
     Image {
-        /// Maps the ID of each mesh to the image containing its morph target
-        /// displacements.
-        mesh_id_to_image: HashMap<AssetId<Mesh>, MorphTargetImage>,
+        morph_target_images: MorphTargetImages,
     },
     /// The variant used when the current platform does support storage buffers.
     ///
@@ -137,7 +159,7 @@ impl FromWorld for RenderMorphTargetAllocator {
         let render_device = world.resource::<RenderDevice>();
         if bevy_render::storage_buffers_are_unsupported(&render_device.limits()) {
             RenderMorphTargetAllocator::Image {
-                mesh_id_to_image: HashMap::default(),
+                morph_target_images: MorphTargetImages::default(),
             }
         } else {
             RenderMorphTargetAllocator::Storage
@@ -175,70 +197,57 @@ impl RenderMorphTargetAllocator {
         &mut self,
         render_device: &RenderDevice,
         render_queue: &RenderQueue,
-        mesh_id: AssetId<Mesh>,
         targets: &[MorphAttributes],
         vertex_count: usize,
-    ) {
+    ) -> Option<MorphTargetImageHandle> {
         match *self {
             RenderMorphTargetAllocator::Image {
-                ref mut mesh_id_to_image,
-            } => {
-                if let Ok(morph_target_image) =
-                    MorphTargetImage::new(render_device, render_queue, targets, vertex_count)
-                {
-                    mesh_id_to_image.insert(mesh_id, morph_target_image);
+                ref mut morph_target_images,
+            } => match MorphTargetImage::new(render_device, render_queue, targets, vertex_count) {
+                Ok(morph_target_image) => {
+                    let handle = morph_target_images.insert(morph_target_image);
+                    Some(handle)
                 }
-            }
-
+                Err(e) => {
+                    error!("Failed to build morph target image for mesh {e:?}");
+                    None
+                }
+            },
             RenderMorphTargetAllocator::Storage => {
                 // Do nothing. Morph target displacements are managed by the
                 // mesh allocator in this case.
+                None
             }
         }
     }
 
     /// Frees the storage associated with morph target displacements for the
     /// mesh with the given ID.
-    ///
-    /// If the current platform doesn't support storage buffers, this drops the
-    /// reference to the [`MorphTargetImage`] that stores the data for the
-    /// mesh's morph target displacements. If the current platform does support
-    /// storage buffers, this method does nothing, as morph target displacements
-    /// are managed by the mesh allocator in this case.
-    ///
-    /// If passed a mesh without morph targets, this method does nothing.
-    pub fn free(&mut self, mesh_id: AssetId<Mesh>) {
+    pub fn free(&mut self, handle: MorphTargetImageHandle) {
         match *self {
             RenderMorphTargetAllocator::Image {
-                ref mut mesh_id_to_image,
-            } => {
-                if mesh_id_to_image.remove(&mesh_id).is_none() {
-                    error!(
-                        "Attempted to free a morph target allocation that wasn't allocated: {:?}",
-                        mesh_id
-                    );
-                }
-            }
-            RenderMorphTargetAllocator::Storage => {
-                // Do nothing. Morph target displacements are managed by the
-                // mesh allocator in this case.
-            }
+                ref mut morph_target_images,
+            } => morph_target_images.remove(handle),
+            RenderMorphTargetAllocator::Storage => error!(
+                "Attempted to free a morph target allocation {:?} when using storage allocator {:?}",
+                handle,
+                self
+            ),
         }
     }
 
     /// Returns the [`MorphTargetImage`] containing the packed morph target
-    /// displacements for the mesh with the given ID, if that image is present.
-    ///
-    /// A [`MorphTargetImage`] is only available if storage buffers aren't
-    /// supported on the given platform. If storage buffers are supported, this
-    /// method returns `None`, as the mesh allocator stores the morph target
-    /// displacements in that case.
-    pub fn get_image(&self, mesh_id: AssetId<Mesh>) -> Option<MorphTargetImage> {
+    /// displacements for the mesh with the given ID.
+    pub fn get_image(&self, handle: &MorphTargetImageHandle) -> &MorphTargetImage {
         match *self {
             RenderMorphTargetAllocator::Image {
-                ref mesh_id_to_image,
-            } => mesh_id_to_image.get(&mesh_id).cloned(),
-            RenderMorphTargetAllocator::Storage => None,
+                ref morph_target_images,
+            } => morph_target_images.get(handle),
+            RenderMorphTargetAllocator::Storage => panic!(
+                "Attempted to get a morph target image with allocation {:?} when using storage allocator {:?}",
+                handle,
+                self
+            ),
         }
     }
 }