From e8b28547bf27153e323a5f9406a9840c3cedcd7c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois?= <mockersf@gmail.com>
Date: Tue, 20 Dec 2022 16:16:58 +0000
Subject: [PATCH] Cleanup dynamic scene before building (#6254)

# Objective

- Dynamic scene builder can build scenes without components, if they didn't have any matching the type registry
- Those entities are not really useful in the final `DynamicScene`

## Solution

- Add a method `remove_empty_entities` that will remove empty entities. It's not called by default when calling `build`, I'm not sure if that's a good idea or not.
---
 .../bevy_scene/src/dynamic_scene_builder.rs   | 58 +++++++++++++++----
 1 file changed, 46 insertions(+), 12 deletions(-)

diff --git a/crates/bevy_scene/src/dynamic_scene_builder.rs b/crates/bevy_scene/src/dynamic_scene_builder.rs
index 841af6756f..5ea65a8714 100644
--- a/crates/bevy_scene/src/dynamic_scene_builder.rs
+++ b/crates/bevy_scene/src/dynamic_scene_builder.rs
@@ -31,9 +31,9 @@ use std::collections::BTreeMap;
 /// let dynamic_scene = builder.build();
 /// ```
 pub struct DynamicSceneBuilder<'w> {
-    entities: BTreeMap<u32, DynamicEntity>,
+    extracted_scene: BTreeMap<u32, DynamicEntity>,
     type_registry: AppTypeRegistry,
-    world: &'w World,
+    original_world: &'w World,
 }
 
 impl<'w> DynamicSceneBuilder<'w> {
@@ -41,9 +41,9 @@ impl<'w> DynamicSceneBuilder<'w> {
     /// All components registered in that world's [`AppTypeRegistry`] resource will be extracted.
     pub fn from_world(world: &'w World) -> Self {
         Self {
-            entities: default(),
+            extracted_scene: default(),
             type_registry: world.resource::<AppTypeRegistry>().clone(),
-            world,
+            original_world: world,
         }
     }
 
@@ -51,16 +51,19 @@ impl<'w> DynamicSceneBuilder<'w> {
     /// Only components registered in the given [`AppTypeRegistry`] will be extracted.
     pub fn from_world_with_type_registry(world: &'w World, type_registry: AppTypeRegistry) -> Self {
         Self {
-            entities: default(),
+            extracted_scene: default(),
             type_registry,
-            world,
+            original_world: world,
         }
     }
 
     /// Consume the builder, producing a [`DynamicScene`].
+    ///
+    /// To make sure the dynamic scene doesn't contain entities without any components, call
+    /// [`Self::remove_empty_entities`] before building the scene.
     pub fn build(self) -> DynamicScene {
         DynamicScene {
-            entities: self.entities.into_values().collect(),
+            entities: self.extracted_scene.into_values().collect(),
         }
     }
 
@@ -71,6 +74,16 @@ impl<'w> DynamicSceneBuilder<'w> {
         self.extract_entities(std::iter::once(entity))
     }
 
+    /// Despawns all enitities with no components.
+    ///
+    /// These were likely created because none of their components were present in the provided type registry upon extraction.
+    pub fn remove_empty_entities(&mut self) -> &mut Self {
+        self.extracted_scene
+            .retain(|_, entity| !entity.components.is_empty());
+
+        self
+    }
+
     /// Extract entities from the builder's [`World`].
     ///
     /// Re-extracting an entity that was already extracted will have no effect.
@@ -102,7 +115,7 @@ impl<'w> DynamicSceneBuilder<'w> {
         for entity in entities {
             let index = entity.index();
 
-            if self.entities.contains_key(&index) {
+            if self.extracted_scene.contains_key(&index) {
                 continue;
             }
 
@@ -111,21 +124,22 @@ impl<'w> DynamicSceneBuilder<'w> {
                 components: Vec::new(),
             };
 
-            for component_id in self.world.entity(entity).archetype().components() {
+            for component_id in self.original_world.entity(entity).archetype().components() {
                 let reflect_component = self
-                    .world
+                    .original_world
                     .components()
                     .get_info(component_id)
                     .and_then(|info| type_registry.get(info.type_id().unwrap()))
                     .and_then(|registration| registration.data::<ReflectComponent>());
 
                 if let Some(reflect_component) = reflect_component {
-                    if let Some(component) = reflect_component.reflect(self.world, entity) {
+                    if let Some(component) = reflect_component.reflect(self.original_world, entity)
+                    {
                         entry.components.push(component.clone_value());
                     }
                 }
             }
-            self.entities.insert(index, entry);
+            self.extracted_scene.insert(index, entry);
         }
 
         drop(type_registry);
@@ -270,4 +284,24 @@ mod tests {
         scene_entities.sort();
         assert_eq!(scene_entities, [entity_a_b.index(), entity_a.index()]);
     }
+
+    #[test]
+    fn remove_componentless_entity() {
+        let mut world = World::default();
+
+        let atr = AppTypeRegistry::default();
+        atr.write().register::<ComponentA>();
+        world.insert_resource(atr);
+
+        let entity_a = world.spawn(ComponentA).id();
+        let entity_b = world.spawn(ComponentB).id();
+
+        let mut builder = DynamicSceneBuilder::from_world(&world);
+        builder.extract_entities([entity_a, entity_b].into_iter());
+        builder.remove_empty_entities();
+        let scene = builder.build();
+
+        assert_eq!(scene.entities.len(), 1);
+        assert_eq!(scene.entities[0].entity, entity_a.index());
+    }
 }