From 3e37378f9cc65629723ad3fe653e8b05397675e4 Mon Sep 17 00:00:00 2001
From: Jeremiah Senkpiel <fishrock123@rocketmail.com>
Date: Sat, 22 Aug 2020 10:30:55 -0700
Subject: [PATCH] SceneSpawner: add unload() and unload_sync()

This allows you to unload scenes.

Refs: https://github.com/bevyengine/bevy/issues/255
---
 crates/bevy_scene/src/scene_spawner.rs | 57 ++++++++++++++++++++++----
 1 file changed, 49 insertions(+), 8 deletions(-)

diff --git a/crates/bevy_scene/src/scene_spawner.rs b/crates/bevy_scene/src/scene_spawner.rs
index 9453ba13138c4..99cb703bb3ded 100644
--- a/crates/bevy_scene/src/scene_spawner.rs
+++ b/crates/bevy_scene/src/scene_spawner.rs
@@ -28,6 +28,7 @@ pub struct SceneSpawner {
     scene_asset_event_reader: EventReader<AssetEvent<Scene>>,
     scenes_to_spawn: Vec<Handle<Scene>>,
     scenes_to_load: Vec<Handle<Scene>>,
+    scenes_to_unload: Vec<Handle<Scene>>,
 }
 
 #[derive(Error, Debug)]
@@ -47,6 +48,10 @@ impl SceneSpawner {
         self.scenes_to_load.push(scene_handle);
     }
 
+    pub fn unload(&mut self, scene_handle: Handle<Scene>) {
+        self.scenes_to_unload.push(scene_handle);
+    }
+
     pub fn load_sync(
         &mut self,
         world: &mut World,
@@ -58,6 +63,28 @@ impl SceneSpawner {
         Ok(())
     }
 
+    pub fn unload_sync(
+        &mut self,
+        world: &mut World,
+        scene_handle: Handle<Scene>,
+    ) -> Result<(), SceneSpawnError> {
+        if let Some(instance_ids) = self.spawned_scenes.get(&scene_handle) {
+            for instance_id in instance_ids {
+                if let Some(instance) = self.spawned_instances.get(&instance_id) {
+                    for entity in instance.entity_map.values() {
+                        let _ = world.despawn(*entity); // Ignore the result, unload only cares if it exists.
+                    }
+                }
+            }
+
+            self.loaded_scenes
+                .retain(|handle| handle.id != scene_handle.id);
+            self.spawned_scenes
+                .retain(|handle, _| handle.id != scene_handle.id);
+        }
+        Ok(())
+    }
+
     pub fn spawn_sync(
         &mut self,
         world: &mut World,
@@ -152,19 +179,32 @@ impl SceneSpawner {
         world: &mut World,
         resources: &Resources,
     ) -> Result<(), SceneSpawnError> {
-        let scenes_to_load = self.scenes_to_load.drain(..).collect::<Vec<_>>();
-        let mut non_existent_scenes = Vec::new();
+        let mut scenes_to_load = Vec::new();
+        std::mem::swap(&mut self.scenes_to_load, &mut scenes_to_load);
+
         for scene_handle in scenes_to_load {
             match self.load_sync(world, resources, scene_handle) {
                 Ok(_) => {}
                 Err(SceneSpawnError::NonExistentScene { .. }) => {
-                    non_existent_scenes.push(scene_handle)
+                    self.scenes_to_spawn.push(scene_handle)
                 }
                 Err(err) => return Err(err),
             }
         }
 
-        self.scenes_to_load = non_existent_scenes;
+        Ok(())
+    }
+
+    pub fn unload_queued_scenes(&mut self, world: &mut World) -> Result<(), SceneSpawnError> {
+        let mut scenes_to_unload = Vec::new();
+        std::mem::swap(&mut self.scenes_to_unload, &mut scenes_to_unload);
+
+        for scene_handle in scenes_to_unload {
+            match self.unload_sync(world, scene_handle) {
+                Ok(_) => {}
+                Err(err) => return Err(err),
+            }
+        }
         Ok(())
     }
 
@@ -173,19 +213,19 @@ impl SceneSpawner {
         world: &mut World,
         resources: &Resources,
     ) -> Result<(), SceneSpawnError> {
-        let scenes_to_spawn = self.scenes_to_spawn.drain(..).collect::<Vec<_>>();
-        let mut non_existent_scenes = Vec::new();
+        let mut scenes_to_spawn = Vec::new();
+        std::mem::swap(&mut self.scenes_to_spawn, &mut scenes_to_spawn);
+
         for scene_handle in scenes_to_spawn {
             match self.spawn_sync(world, resources, scene_handle) {
                 Ok(_) => {}
                 Err(SceneSpawnError::NonExistentScene { .. }) => {
-                    non_existent_scenes.push(scene_handle)
+                    self.scenes_to_spawn.push(scene_handle)
                 }
                 Err(err) => return Err(err),
             }
         }
 
-        self.scenes_to_spawn = non_existent_scenes;
         Ok(())
     }
 }
@@ -209,6 +249,7 @@ pub fn scene_spawner_system(world: &mut World, resources: &mut Resources) {
         }
     }
 
+    scene_spawner.unload_queued_scenes(world).unwrap();
     scene_spawner.load_queued_scenes(world, resources).unwrap();
     scene_spawner.spawn_queued_scenes(world, resources).unwrap();
     scene_spawner