From a74f31ca0e474cfe79bbc326531815d037275d57 Mon Sep 17 00:00:00 2001 From: Raphael Date: Fri, 8 May 2026 20:01:49 -0400 Subject: [PATCH 1/5] Store previous_asset in PrepareNextFrameAssets instead of getting removed value. --- crates/bevy_render/src/render_asset.rs | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/crates/bevy_render/src/render_asset.rs b/crates/bevy_render/src/render_asset.rs index a3ea3613d1c76..14588b3536678 100644 --- a/crates/bevy_render/src/render_asset.rs +++ b/crates/bevy_render/src/render_asset.rs @@ -364,7 +364,7 @@ pub(crate) fn extract_render_asset( /// All assets that should be prepared next frame. #[derive(Resource)] pub struct PrepareNextFrameAssets { - assets: Vec<(AssetId, A::SourceAsset)>, + assets: Vec<(AssetId, A::SourceAsset, Option)>, } impl Default for PrepareNextFrameAssets { @@ -388,7 +388,7 @@ pub fn prepare_assets( let mut param = param.into_inner(); let queued_assets = core::mem::take(&mut prepare_next_frame.assets); - for (id, extracted_asset) in queued_assets { + for (id, extracted_asset, previous_asset) in queued_assets { if extracted_assets.removed.contains(&id) || extracted_assets.added.contains(&id) { // skip previous frame's assets that have been removed or updated continue; @@ -400,7 +400,9 @@ pub fn prepare_assets( // this way we always write at least one (sized) asset per frame. // in future we could also consider partial asset uploads. if bpf.exhausted() { - prepare_next_frame.assets.push((id, extracted_asset)); + prepare_next_frame + .assets + .push((id, extracted_asset, previous_asset)); continue; } size @@ -408,15 +410,16 @@ pub fn prepare_assets( 0 }; - let previous_asset = render_assets.get(id); - match A::prepare_asset(extracted_asset, id, &mut param, previous_asset) { + match A::prepare_asset(extracted_asset, id, &mut param, previous_asset.as_ref()) { Ok(prepared_asset) => { render_assets.insert(id, prepared_asset); bpf.write_bytes(write_bytes); wrote_asset_count += 1; } Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => { - prepare_next_frame.assets.push((id, extracted_asset)); + prepare_next_frame + .assets + .push((id, extracted_asset, previous_asset)); } Err(PrepareAssetError::AsBindGroupError(e)) => { error!( @@ -440,7 +443,9 @@ pub fn prepare_assets( let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) { if bpf.exhausted() { - prepare_next_frame.assets.push((id, extracted_asset)); + prepare_next_frame + .assets + .push((id, extracted_asset, previous_asset)); continue; } size @@ -455,7 +460,9 @@ pub fn prepare_assets( wrote_asset_count += 1; } Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => { - prepare_next_frame.assets.push((id, extracted_asset)); + prepare_next_frame + .assets + .push((id, extracted_asset, previous_asset)); } Err(PrepareAssetError::AsBindGroupError(e)) => { error!( From 684cf644ccbe1f8c47d3273df6727b6df7c37cae Mon Sep 17 00:00:00 2001 From: Raphael Date: Sat, 9 May 2026 11:22:39 -0400 Subject: [PATCH 2/5] Change `RenderAsset::unload` to take ownership of self and called before drop. --- crates/bevy_render/src/render_asset.rs | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/crates/bevy_render/src/render_asset.rs b/crates/bevy_render/src/render_asset.rs index 14588b3536678..fd17fdc3cadd9 100644 --- a/crates/bevy_render/src/render_asset.rs +++ b/crates/bevy_render/src/render_asset.rs @@ -80,13 +80,14 @@ pub trait RenderAsset: Send + Sync + 'static + Sized { previous_asset: Option<&Self>, ) -> Result>; - /// Called whenever the [`RenderAsset::SourceAsset`] has been removed. + /// Called whenever the [`RenderAsset`] is dropped (the [`RenderAsset::SourceAsset`] is reextracted (updated) or removed). /// /// You can implement this method if you need to access ECS data (via /// `_param`) in order to perform cleanup tasks when the asset is removed. /// /// The default implementation does nothing. fn unload_asset( + self, _source_asset: AssetId, _param: &mut SystemParamItem, ) { @@ -420,6 +421,7 @@ pub fn prepare_assets( prepare_next_frame .assets .push((id, extracted_asset, previous_asset)); + continue; } Err(PrepareAssetError::AsBindGroupError(e)) => { error!( @@ -428,11 +430,16 @@ pub fn prepare_assets( ); } } + + if let Some(unload) = previous_asset { + unload.unload_asset(id, &mut param); + } } for removed in extracted_assets.removed.drain() { - render_assets.remove(removed); - A::unload_asset(removed, &mut param); + if let Some(unload) = render_assets.remove(removed) { + unload.unload_asset(removed, &mut param); + } } for (id, extracted_asset) in extracted_assets.extracted.drain(..) { @@ -463,6 +470,7 @@ pub fn prepare_assets( prepare_next_frame .assets .push((id, extracted_asset, previous_asset)); + continue; } Err(PrepareAssetError::AsBindGroupError(e)) => { error!( @@ -471,6 +479,10 @@ pub fn prepare_assets( ); } } + + if let Some(unload) = previous_asset { + unload.unload_asset(id, &mut param); + } } if bpf.exhausted() && !prepare_next_frame.assets.is_empty() { From e466b0061d955f1b5b54f44a79873c015e2c5e3c Mon Sep 17 00:00:00 2001 From: Raphael Date: Sat, 9 May 2026 11:24:07 -0400 Subject: [PATCH 3/5] Updated signiture of `RenderMesh::unload_asset` to take self as well. --- crates/bevy_render/src/mesh/mod.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/bevy_render/src/mesh/mod.rs b/crates/bevy_render/src/mesh/mod.rs index e60e10a41658a..537de865007bf 100644 --- a/crates/bevy_render/src/mesh/mod.rs +++ b/crates/bevy_render/src/mesh/mod.rs @@ -223,6 +223,7 @@ impl RenderAsset for RenderMesh { } fn unload_asset( + self, _mesh_id: AssetId, (_, _, _, _render_morph_targets_allocator): &mut SystemParamItem, ) { From c73074e2dbc7a4fcf3bff33f51675df5bdf00719 Mon Sep 17 00:00:00 2001 From: Raphael Date: Sat, 9 May 2026 16:56:42 -0400 Subject: [PATCH 4/5] Changed RenderMorphTargetAllocator to give out handles to match RenderAsset::unload_asset functionality --- crates/bevy_pbr/src/render/mesh.rs | 17 ++-- crates/bevy_render/Cargo.toml | 1 + crates/bevy_render/src/mesh/mod.rs | 20 +++-- crates/bevy_render/src/mesh/morph.rs | 115 +++++++++++++++------------ 4 files changed, 85 insertions(+), 68 deletions(-) diff --git a/crates/bevy_pbr/src/render/mesh.rs b/crates/bevy_pbr/src/render/mesh.rs index 253540c6dbcb9..585f6d4990740 100644 --- a/crates/bevy_pbr/src/render/mesh.rs +++ b/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, MorphTargetImage>, + morph_target_images: &MorphTargetImages, morph_targets: &mut HashMap, 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 { diff --git a/crates/bevy_render/Cargo.toml b/crates/bevy_render/Cargo.toml index c0d5b89d2bc9f..a566db1e8f2ae 100644 --- a/crates/bevy_render/Cargo.toml +++ b/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" diff --git a/crates/bevy_render/src/mesh/mod.rs b/crates/bevy_render/src/mesh/mod.rs index 537de865007bf..11125be3111d3 100644 --- a/crates/bevy_render/src/mesh/mod.rs +++ b/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, } 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,6 +221,8 @@ impl RenderAsset for RenderMesh { buffer_info, key_bits, layout: mesh_vertex_buffer_layout, + #[cfg(feature = "morph")] + morph_target_handle, }) } @@ -229,6 +233,8 @@ impl RenderAsset for RenderMesh { ) { // 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); + } } } diff --git a/crates/bevy_render/src/mesh/morph.rs b/crates/bevy_render/src/mesh/morph.rs index d3b05c07dac1f..f6e0abdc53d7e 100644 --- a/crates/bevy_render/src/mesh/morph.rs +++ b/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); + +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, 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::(); 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, targets: &[MorphAttributes], vertex_count: usize, - ) { + ) -> Option { 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) { + 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) -> Option { + /// 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, + } => return morph_target_images.get(handle), + RenderMorphTargetAllocator::Storage => panic!( + "Attempted to get a morph target image with allocation {:?} when using storage allocator {:?}", + handle, + self + ), } } } From f6bceed15953e9b6850c418eb057ddbdeac26060 Mon Sep 17 00:00:00 2001 From: Raphael Date: Mon, 11 May 2026 12:42:24 -0400 Subject: [PATCH 5/5] Remove unnececary return. --- crates/bevy_render/src/mesh/morph.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/bevy_render/src/mesh/morph.rs b/crates/bevy_render/src/mesh/morph.rs index f6e0abdc53d7e..e0276e9b55093 100644 --- a/crates/bevy_render/src/mesh/morph.rs +++ b/crates/bevy_render/src/mesh/morph.rs @@ -242,7 +242,7 @@ impl RenderMorphTargetAllocator { match *self { RenderMorphTargetAllocator::Image { ref morph_target_images, - } => return morph_target_images.get(handle), + } => morph_target_images.get(handle), RenderMorphTargetAllocator::Storage => panic!( "Attempted to get a morph target image with allocation {:?} when using storage allocator {:?}", handle,