From a33c85b964810acac478eebca67ff717ee928eff Mon Sep 17 00:00:00 2001
From: Yoh Deadfall <yoh.deadfall@hotmail.com>
Date: Tue, 19 Jul 2022 14:52:16 +0300
Subject: [PATCH] Freeing unused memory held by visible entities

---
 crates/bevy_pbr/src/light.rs | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/crates/bevy_pbr/src/light.rs b/crates/bevy_pbr/src/light.rs
index 5033b273b1b32..10f9e5be4bf8c 100644
--- a/crates/bevy_pbr/src/light.rs
+++ b/crates/bevy_pbr/src/light.rs
@@ -1551,6 +1551,22 @@ pub fn check_light_mesh_visibility(
         (Without<NotShadowCaster>, Without<DirectionalLight>),
     >,
 ) {
+    fn shrink_entities(visible_entities: &mut VisibleEntities) {
+        // Check that visible entities capacity() is no more than two times greater than len()
+        let capacity = visible_entities.entities.capacity();
+        let reserved = capacity
+            .checked_div(visible_entities.entities.len())
+            .map_or(0, |reserve| {
+                if reserve > 2 {
+                    capacity / (reserve / 2)
+                } else {
+                    capacity
+                }
+            });
+
+        visible_entities.entities.shrink_to(reserved);
+    }
+
     // Directional lights
     for (
         directional_light,
@@ -1592,8 +1608,7 @@ pub fn check_light_mesh_visibility(
             visible_entities.entities.push(entity);
         }
 
-        // TODO: check for big changes in visible entities len() vs capacity() (ex: 2x) and resize
-        // to prevent holding unneeded memory
+        shrink_entities(&mut visible_entities);
     }
 
     for visible_lights in &visible_point_lights {
@@ -1664,11 +1679,12 @@ pub fn check_light_mesh_visibility(
                     }
                 }
 
-                // TODO: check for big changes in visible entities len() vs capacity() (ex: 2x) and resize
-                // to prevent holding unneeded memory
+                for visible_entities in cubemap_visible_entities.iter_mut() {
+                    shrink_entities(visible_entities);
+                }
             }
 
-            // spot lights
+            // Spot lights
             if let Ok((point_light, transform, frustum, mut visible_entities, maybe_view_mask)) =
                 spot_lights.get_mut(light_entity)
             {
@@ -1720,8 +1736,7 @@ pub fn check_light_mesh_visibility(
                     }
                 }
 
-                // TODO: check for big changes in visible entities len() vs capacity() (ex: 2x) and resize
-                // to prevent holding unneeded memory
+                shrink_entities(&mut visible_entities);
             }
         }
     }