Remove InstanceId when Scene Despawn

crates/bevy_scene/src/bundle.rs

@@ -16,7 +16,7 @@ use crate::{DynamicScene, InstanceId, Scene, SceneSpawner};
 /// [`InstanceId`] of a spawned scene. It can be used with the [`SceneSpawner`] to
 /// interact with the spawned scene.
 #[derive(Component, Deref, DerefMut)]
-pub struct SceneInstance(InstanceId);
+pub struct SceneInstance(pub(crate) InstanceId);
 
 /// A component bundle for a [`Scene`] root.
 ///

crates/bevy_scene/src/lib.rs

@@ -44,7 +44,7 @@ pub mod prelude {
 }
 
 use bevy_app::prelude::*;
-use bevy_asset::AssetApp;
+use bevy_asset::{AssetApp, Handle};
 
 /// Plugin that provides scene functionality to an [`App`].
 #[derive(Default)]
@@ -59,6 +59,37 @@ impl Plugin for ScenePlugin {
             .add_event::<SceneInstanceReady>()
             .init_resource::<SceneSpawner>()
             .add_systems(SpawnScene, (scene_spawner, scene_spawner_system).chain());
+
+        // Register component hooks for DynamicScene
+        app.world
+            .register_component_hooks::<Handle<DynamicScene>>()
+            .on_remove(|mut world, entity, _| {
+                let Some(handle) = world.get::<Handle<DynamicScene>>(entity) else {
+                    return;
+                };
+                let id = handle.id();
+                if let Some(&SceneInstance(scene_instance)) = world.get::<SceneInstance>(entity) {
+                    let Some(mut scene_spawner) = world.get_resource_mut::<SceneSpawner>() else {
+                        return;
+                    };
+                    if let Some(instance_ids) = scene_spawner.spawned_dynamic_scenes.get_mut(&id) {
+                        instance_ids.remove(&scene_instance);
+                    }
+                    scene_spawner.despawn_instance(scene_instance);
+                }
+            });
+
+        // Register component hooks for Scene
+        app.world
+            .register_component_hooks::<Handle<Scene>>()
+            .on_remove(|mut world, entity, _| {
+                if let Some(&SceneInstance(scene_instance)) = world.get::<SceneInstance>(entity) {
+                    let Some(mut scene_spawner) = world.get_resource_mut::<SceneSpawner>() else {
+                        return;
+                    };
+                    scene_spawner.despawn_instance(scene_instance);
+                }
+            });
     }
 }
 

crates/bevy_scene/src/scene_spawner.rs

@@ -60,9 +60,8 @@ impl InstanceId {
 /// - [`despawn_instance`](Self::despawn_instance)
 #[derive(Default, Resource)]
 pub struct SceneSpawner {
-    spawned_scenes: HashMap<AssetId<Scene>, Vec<InstanceId>>,
-    spawned_dynamic_scenes: HashMap<AssetId<DynamicScene>, Vec<InstanceId>>,
-    spawned_instances: HashMap<InstanceId, InstanceInfo>,
+    pub(crate) spawned_dynamic_scenes: HashMap<AssetId<DynamicScene>, HashSet<InstanceId>>,
+    pub(crate) spawned_instances: HashMap<InstanceId, InstanceInfo>,
     scene_asset_event_reader: ManualEventReader<AssetEvent<DynamicScene>>,
     dynamic_scenes_to_spawn: Vec<(Handle<DynamicScene>, InstanceId)>,
     scenes_to_spawn: Vec<(Handle<Scene>, InstanceId)>,
@@ -210,7 +209,7 @@ impl SceneSpawner {
         self.spawned_instances
             .insert(instance_id, InstanceInfo { entity_map });
         let spawned = self.spawned_dynamic_scenes.entry(id).or_default();
-        spawned.push(instance_id);
+        spawned.insert(instance_id);
         Ok(instance_id)
     }
 
@@ -251,8 +250,6 @@ impl SceneSpawner {
                 scene.write_to_world_with(world, &world.resource::<AppTypeRegistry>().clone())?;
 
             self.spawned_instances.insert(instance_id, instance_info);
-            let spawned = self.spawned_scenes.entry(id).or_default();
-            spawned.push(instance_id);
             Ok(instance_id)
         })
     }
@@ -310,8 +307,8 @@ impl SceneSpawner {
                     let spawned = self
                         .spawned_dynamic_scenes
                         .entry(handle.id())
-                        .or_insert_with(Vec::new);
-                    spawned.push(instance_id);
+                        .or_insert_with(HashSet::new);
+                    spawned.insert(instance_id);
                 }
                 Err(SceneSpawnError::NonExistentScene { .. }) => {
                     self.dynamic_scenes_to_spawn.push((handle, instance_id));
@@ -443,17 +440,14 @@ pub fn scene_spawner_system(world: &mut World) {
 mod tests {
     use bevy_app::App;
     use bevy_asset::{AssetPlugin, AssetServer};
-    use bevy_ecs::component::Component;
-    use bevy_ecs::entity::Entity;
     use bevy_ecs::event::EventReader;
     use bevy_ecs::prelude::ReflectComponent;
     use bevy_ecs::query::With;
-    use bevy_ecs::reflect::AppTypeRegistry;
     use bevy_ecs::system::{Commands, Res, ResMut, RunSystemOnce};
-    use bevy_ecs::world::World;
+    use bevy_ecs::{component::Component, system::Query};
     use bevy_reflect::Reflect;
 
-    use crate::{DynamicScene, DynamicSceneBuilder, SceneInstanceReady, ScenePlugin, SceneSpawner};
+    use crate::{DynamicSceneBuilder, ScenePlugin};
 
     use super::*;
 
@@ -554,4 +548,47 @@ mod tests {
             },
         );
     }
+
+    #[test]
+    fn despawn_scene() {
+        let mut app = App::new();
+        app.add_plugins((AssetPlugin::default(), ScenePlugin));
+        app.register_type::<ComponentA>();
+
+        let asset_server = app.world.resource::<AssetServer>();
+
+        // Build scene.
+        let scene = asset_server.add(DynamicScene::default());
+        let count = 10;
+
+        // Checks the number of scene instances stored in `SceneSpawner`.
+        let check = |world: &mut World, expected_count: usize| {
+            let scene_spawner = world.resource::<SceneSpawner>();
+            assert_eq!(
+                scene_spawner.spawned_dynamic_scenes[&scene.id()].len(),
+                expected_count
+            );
+            assert_eq!(scene_spawner.spawned_instances.len(), expected_count);
+        };
+
+        // Spawn scene.
+        for _ in 0..count {
+            app.world.spawn((ComponentA, scene.clone()));
+        }
+
+        app.update();
+        check(&mut app.world, count);
+
+        // Despawn scene.
+        app.world.run_system_once(
+            |mut commands: Commands, query: Query<Entity, With<ComponentA>>| {
+                for entity in query.iter() {
+                    commands.entity(entity).despawn_recursive();
+                }
+            },
+        );
+
+        app.update();
+        check(&mut app.world, 0);
+    }
 }