Reuse heap allocations in prepare_mesh_bind_groups

crates/bevy_pbr/src/render/mesh.rs

@@ -3748,7 +3748,7 @@ pub struct MeshPhaseBindGroups {
     skinned: HashMap<MeshSlabId, MeshBindGroupPair>,
     /// Bind groups for meshes with morph targets.
     morph_targets: MeshMorphTargetBindGroups,
-    lightmaps: HashMap<LightmapSlabIndex, BindGroup>,
+    lightmaps: HashMap<(MeshSlabId, LightmapSlabIndex), BindGroup>,
 }
 
 /// Stores bind groups for each mesh with morph targets.
@@ -3888,7 +3888,7 @@ impl MeshPhaseBindGroups {
                 .get(&metadata_slab_id)
                 .map(|bind_group_pair| bind_group_pair.get(motion_vectors)),
             (false, MeshMorphBindGroupKey::NoMorphTargets, Some(lightmap_slab)) => {
-                self.lightmaps.get(&lightmap_slab)
+                self.lightmaps.get(&(metadata_slab_id, lightmap_slab))
             }
             (false, MeshMorphBindGroupKey::NoMorphTargets, None) => {
                 self.model_only.get(&metadata_slab_id)
@@ -3943,17 +3943,33 @@ pub fn prepare_mesh_bind_groups(
     weights_uniform: Res<MorphUniforms>,
     mesh_allocator: Res<MeshAllocator>,
     render_morph_target_allocator: Res<RenderMorphTargetAllocator>,
-    mut render_lightmaps: ResMut<RenderLightmaps>,
+    render_lightmaps: Res<RenderLightmaps>,
     metadata_fallback_buffer: Res<MeshMetadataFallbackBuffer>,
+    mut mesh_bind_groups: Option<ResMut<MeshBindGroups>>,
 ) {
     // CPU mesh preprocessing path.
     if let Some(cpu_batched_instance_buffer) = cpu_batched_instance_buffer
         && let Some(instance_data_binding) = cpu_batched_instance_buffer
             .into_inner()
             .instance_data_binding()
     {
+        // Reuse allocations
+        let mut cpu_preprocessing_mesh_bind_groups = match mesh_bind_groups.as_deref_mut() {
+            None | Some(MeshBindGroups::GpuPreprocessing(_)) => {
+                MeshPhaseBindGroups::new(&render_device)
+            }
+            Some(MeshBindGroups::CpuPreprocessing(cpu_preprocessing_mesh_bind_groups)) => {
+                core::mem::replace(
+                    cpu_preprocessing_mesh_bind_groups,
+                    MeshPhaseBindGroups::new(&render_device),
+                )
+            }
+        };
+        cpu_preprocessing_mesh_bind_groups.reset();
+
         // In this path, we only have a single set of bind groups for all phases.
-        let cpu_preprocessing_mesh_bind_groups = prepare_mesh_bind_groups_for_phase(
+        prepare_mesh_bind_groups_for_phase(
+            &mut cpu_preprocessing_mesh_bind_groups,
             instance_data_binding,
             &meshes,
             &mesh_pipeline,
@@ -3963,7 +3979,7 @@ pub fn prepare_mesh_bind_groups(
             &weights_uniform,
             &mesh_allocator,
             &render_morph_target_allocator,
-            &mut render_lightmaps,
+            &render_lightmaps,
             &metadata_fallback_buffer,
         );
 
@@ -3975,7 +3991,13 @@ pub fn prepare_mesh_bind_groups(
 
     // GPU mesh preprocessing path.
     if let Some(gpu_batched_instance_buffers) = gpu_batched_instance_buffers {
-        let mut gpu_preprocessing_mesh_bind_groups = TypeIdMap::default();
+        // Reuse allocations
+        let mut gpu_preprocessing_mesh_bind_groups = match mesh_bind_groups.as_deref_mut() {
+            None | Some(MeshBindGroups::CpuPreprocessing(_)) => TypeIdMap::default(),
+            Some(MeshBindGroups::GpuPreprocessing(gpu_preprocessing_mesh_bind_groups)) => {
+                core::mem::take(gpu_preprocessing_mesh_bind_groups)
+            }
+        };
 
         // Loop over each phase.
         for (phase_type_id, batched_phase_instance_buffers) in
@@ -3986,8 +4008,13 @@ pub fn prepare_mesh_bind_groups(
             else {
                 continue;
             };
+            let groups = gpu_preprocessing_mesh_bind_groups
+                .entry(*phase_type_id)
+                .or_insert(MeshPhaseBindGroups::new(&render_device));
+            groups.reset();
 
-            let mesh_phase_bind_groups = prepare_mesh_bind_groups_for_phase(
+            prepare_mesh_bind_groups_for_phase(
+                groups,
                 instance_data_binding,
                 &meshes,
                 &mesh_pipeline,
@@ -3997,11 +4024,9 @@ pub fn prepare_mesh_bind_groups(
                 &weights_uniform,
                 &mesh_allocator,
                 &render_morph_target_allocator,
-                &mut render_lightmaps,
+                &render_lightmaps,
                 &metadata_fallback_buffer,
             );
-
-            gpu_preprocessing_mesh_bind_groups.insert(*phase_type_id, mesh_phase_bind_groups);
         }
 
         commands.insert_resource(MeshBindGroups::GpuPreprocessing(
@@ -4012,6 +4037,7 @@ pub fn prepare_mesh_bind_groups(
 
 /// Creates the per-mesh bind groups for each type of mesh, for a single phase.
 fn prepare_mesh_bind_groups_for_phase(
+    groups: &mut MeshPhaseBindGroups,
     model: BindingResource,
     meshes: &RenderAssets<RenderMesh>,
     mesh_pipeline: &MeshPipeline,
@@ -4021,12 +4047,10 @@ fn prepare_mesh_bind_groups_for_phase(
     weights_uniform: &MorphUniforms,
     mesh_allocator: &MeshAllocator,
     render_morph_target_allocator: &RenderMorphTargetAllocator,
-    render_lightmaps: &mut RenderLightmaps,
+    render_lightmaps: &RenderLightmaps,
     metadata_fallback_buffer: &MeshMetadataFallbackBuffer,
-) -> MeshPhaseBindGroups {
+) {
     let layouts = &mesh_pipeline.mesh_layouts;
-    // TODO: Reuse allocations.
-    let mut groups = MeshPhaseBindGroups::new(render_device);
 
     for metadata_slab_id in mesh_allocator.metadata_slabs() {
         let metadata_buffer = mesh_allocator
@@ -4103,9 +4127,12 @@ fn prepare_mesh_bind_groups_for_phase(
 
         // Create lightmap bindgroups. There will be one bindgroup for each slab.
         let bindless_supported = render_lightmaps.bindless_supported;
-        for (lightmap_slab_id, lightmap_slab) in render_lightmaps.slabs.iter_mut().enumerate() {
+        for (lightmap_slab_id, lightmap_slab) in render_lightmaps.slabs.iter().enumerate() {
             groups.lightmaps.insert(
-                LightmapSlabIndex(NonMaxU32::new(lightmap_slab_id as u32).unwrap()),
+                (
+                    metadata_slab_id,
+                    LightmapSlabIndex(NonMaxU32::new(lightmap_slab_id as u32).unwrap()),
+                ),
                 layouts.lightmapped(
                     render_device,
                     pipeline_cache,
@@ -4154,8 +4181,6 @@ fn prepare_mesh_bind_groups_for_phase(
             }
         }
     }
-
-    groups
 }
 
 /// Creates per-mesh morph target bind groups for a single phase.