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

@@ -25,7 +25,7 @@ pub mod serde;
 /// Rusty Object Notation, a crate used to serialize and deserialize bevy scenes.
 pub use bevy_asset::ron;
 
-use bevy_ecs::schedule::IntoSystemConfigs;
+use bevy_ecs::{schedule::IntoSystemConfigs, world::World};
 pub use bundle::*;
 pub use dynamic_scene::*;
 pub use dynamic_scene_builder::*;
@@ -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)]
@@ -58,6 +58,7 @@ impl Plugin for ScenePlugin {
             .init_asset_loader::<SceneLoader>()
             .add_event::<SceneInstanceReady>()
             .init_resource::<SceneSpawner>()
+            .add_systems(Startup, setup)
             .add_systems(SpawnScene, (scene_spawner, scene_spawner_system).chain());
     }
 }
@@ -66,3 +67,26 @@ impl Plugin for ScenePlugin {
 impl Plugin for ScenePlugin {
     fn build(&self, _: &mut App) {}
 }
+
+fn setup(world: &mut World) {
+    world
+        .register_component_hooks::<Handle<DynamicScene>>()
+        .on_remove(|mut world, entity, _| {
+            let id = world.get::<Handle<DynamicScene>>(entity).unwrap().id();
+            if let Some(&SceneInstance(scene_instance)) = world.get::<SceneInstance>(entity) {
+                let mut scene_spawner = world.resource_mut::<SceneSpawner>();
+                if let Some(instance_ids) = scene_spawner.spawned_dynamic_scenes.get_mut(&id) {
+                    instance_ids.remove(&scene_instance);
+                }
+                scene_spawner.despawn_instance(scene_instance);
+            }
+        });
+    world
+        .register_component_hooks::<Handle<Scene>>()
+        .on_remove(|mut world, entity, _| {
+            if let Some(&SceneInstance(scene_instance)) = world.get::<SceneInstance>(entity) {
+                let mut scene_spawner = world.resource_mut::<SceneSpawner>();
+                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));