Despawn entities when the mod that spawned them is despawned

src/asset.rs

@@ -6,8 +6,10 @@ use wasmtime::component::{Component, InstancePre, Val};
 
 use crate::{
     access::ModAccess,
+    cleanup::DespawnModEntities,
     engine::{Engine, Linker},
     host::WasmHost,
+    mods::ModDespawnBehaviour,
     runner::{Config, ConfigRunSystem, ConfigSetup, Runner},
 };
 
@@ -69,9 +71,25 @@ impl ModAsset {
             }
         };
 
-        // This is very cheap, since it's all Arcs
+        // This is very cheap, since it's just Arcs
         let instance_pre = asset.instance_pre.clone();
 
+        // The mod might have reloaded. It's necessary we perform cleanup
+        // if the mod has spawned entities before.
+        if ModDespawnBehaviour::is_despawn_entities(world) {
+            let (entities, mut commands) = world.entities_and_commands();
+            let despawn = entities
+                .get(mod_id)
+                .expect("Mod entity exists")
+                .get::<DespawnModEntities>()
+                .expect(
+                    "DespawnModEntities should have been registered as a required componet for Mod",
+                );
+            for source_entity in despawn.iter() {
+                commands.entity(source_entity).try_despawn();
+            }
+        }
+
         let engine = world
             .get_resource::<Engine>()
             .expect("Engine should never be removed from world");

src/cleanup.rs

@@ -1,12 +1,16 @@
 use bevy_ecs::{
+    entity::EntityHashSet,
     prelude::*,
     schedule::{ScheduleCleanupPolicy, ScheduleError},
     system::SystemState,
 };
 use bevy_log::prelude::*;
 use bevy_platform::collections::{HashMap, HashSet};
 
-use crate::{mods::ModSystemSet, prelude::ModSchedules};
+use crate::{
+    mods::{ModDespawnBehaviour, ModSystemSet},
+    prelude::ModSchedules,
+};
 
 /// A [Message] that triggers disabling of scheduled [ModSystemSets](ModSystemSet).
 ///
@@ -76,3 +80,29 @@ pub(crate) fn disable_mod_system_sets(
             .insert(schedule);
     }
 }
+
+/// A component that tracks all of the entities spawned by a mod (and considered belonging to it).
+///
+/// When a mod is despawned, so will all the entities it spawned due to the `linked_spawn` clause of the relationship.
+#[derive(Component, Default)]
+#[relationship_target(relationship = DespawnModEntity, linked_spawn)]
+pub(crate) struct DespawnModEntities(EntityHashSet);
+
+/// A component that tracks the mod responsible for spawning an entity.
+#[derive(Component)]
+#[relationship(relationship_target = DespawnModEntities)]
+pub(crate) struct DespawnModEntity(pub(crate) Entity);
+
+/// Determines whether [DespawnModEntity] should be inserted to entities spawned by mods
+#[derive(Clone, Copy)]
+pub(crate) struct InsertDespawnComponent(pub(crate) Option<Entity>);
+
+impl InsertDespawnComponent {
+    pub(crate) fn new(mod_id: Entity, world: &World) -> Self {
+        Self(if ModDespawnBehaviour::is_despawn_entities(world) {
+            Some(mod_id)
+        } else {
+            None
+        })
+    }
+}

src/host/app.rs

@@ -73,7 +73,7 @@ impl HostApp for WasmHost {
             for system in systems.iter() {
                 let schedule_config = table
                     .get_mut(system)?
-                    .schedule(world, mod_name, asset_id, asset_version, &access)?
+                    .schedule(world, mod_id, mod_name, asset_id, asset_version, &access)?
                     .in_set(ModSystemSet::All)
                     .in_set(ModSystemSet::Mod(mod_id))
                     .in_set(ModSystemSet::Access(*access));

src/host/commands.rs

@@ -4,8 +4,8 @@ use bevy_log::prelude::*;
 use wasmtime::component::Resource;
 
 use crate::{
-    access::ModAccess, bindings::wasvy::ecs::app::HostCommands, component::insert_component,
-    host::WasmHost, runner::State,
+    access::ModAccess, bindings::wasvy::ecs::app::HostCommands, cleanup::DespawnModEntity,
+    component::insert_component, host::WasmHost, runner::State,
 };
 
 pub struct Commands;
@@ -20,22 +20,29 @@ impl HostCommands for WasmHost {
             mut commands,
             type_registry,
             access,
+            insert_despawn_component,
             ..
         } = self.access()
         else {
             bail!("commands resource is only accessible when running systems")
         };
 
+        let mut entity_commands = commands.spawn_empty();
+
         // Make sure the entity is not spawned outside the sandbox
         // The mod can still override the ChildOf with its own value
-        // Note: We can't currently prevent a mod from creating a component that has a relation to a component outside the sadnbox
+        // Note: We can't currently prevent a mod from creating a component that has a relation to a component outside the sandbox
         // TODO: Restrict what entities a mod can reference via permissions
-        let entity = if let ModAccess::Sandbox(entity) = access {
-            commands.spawn(ChildOf(*entity)).id()
-        } else {
-            commands.spawn_empty().id()
+        if let ModAccess::Sandbox(entity) = access {
+            entity_commands.insert(ChildOf(*entity));
         };
 
+        // Make sure this entity is despawned when the mod is despawned. See [ModDespawnBehaviour]
+        if let Some(mod_id) = insert_despawn_component.0 {
+            entity_commands.insert(DespawnModEntity(mod_id));
+        }
+
+        let entity = entity_commands.id();
         trace!("Spawn empty {entity}, with components:");
 
         for (type_path, serialized_component) in components {

src/host/system.rs

@@ -19,6 +19,7 @@ use crate::{
     access::ModAccess,
     asset::ModAsset,
     bindings::wasvy::ecs::app::{HostSystem, QueryFor},
+    cleanup::InsertDespawnComponent,
     engine::Engine,
     host::{Commands, Query, QueryForComponent, WasmHost, create_query_builder},
     runner::{ConfigRunSystem, Runner, State},
@@ -46,6 +47,7 @@ impl System {
     pub(crate) fn schedule(
         &mut self,
         mut world: &mut World,
+        mod_id: Entity,
         mod_name: &str,
         asset_id: &AssetId<ModAsset>,
         asset_version: &Tick,
@@ -66,6 +68,7 @@ impl System {
             asset_version: asset_version.clone(),
             built_params,
             access: *access,
+            insert_despawn_component: InsertDespawnComponent::new(mod_id, world),
         };
 
         // Generate the queries necessary to run this system
@@ -135,6 +138,7 @@ struct Input {
     asset_version: Tick,
     built_params: Vec<BuiltParam>,
     access: ModAccess,
+    insert_despawn_component: InsertDespawnComponent,
 }
 
 impl FromWorld for Input {
@@ -178,6 +182,7 @@ fn system_runner(
             type_registry: &type_registry,
             queries: &mut queries,
             access: input.access.clone(),
+            insert_despawn_component: input.insert_despawn_component.clone(),
         },
         &params,
     )?;

src/mods.rs

@@ -255,3 +255,33 @@ impl ModSystemSet {
         Self::Access(ModAccess::Sandbox(sandbox_id))
     }
 }
+
+/// An enum that defines what happens when a mod is despawned (or reloaded)
+///
+/// Set this value during plugin instantiation via
+/// [ModloaderPlugin::set_despawn_behaviour](crate::plugin::ModloaderPlugin::set_despawn_behaviour).
+///
+/// The default behaviour is to despawn all entities this mod spawned.
+/// See [DespawnEntities](ModDespawnBehaviour::DespawnEntities).
+#[derive(Resource, Debug, Default, PartialEq, Eq)]
+pub enum ModDespawnBehaviour {
+    /// Do nothing when a mod is despawned
+    None,
+
+    /// The default. Despawn all entities this mod spawned.
+    ///
+    /// So for example if your mod spawns a cube in the center of the scene,
+    /// when this mod is hot reloaded the cube is despawned, and the newest
+    /// version of the mod spawns a new cube in its place.
+    #[default]
+    DespawnEntities,
+}
+
+impl ModDespawnBehaviour {
+    pub(crate) fn is_despawn_entities(world: &World) -> bool {
+        match world.get_resource() {
+            None | Some(ModDespawnBehaviour::DespawnEntities) => true,
+            Some(ModDespawnBehaviour::None) => false,
+        }
+    }
+}

src/plugin.rs

@@ -7,9 +7,10 @@ use bevy_log::prelude::*;
 
 use crate::{
     asset::{ModAsset, ModAssetLoader},
-    cleanup::{DisableSystemSet, disable_mod_system_sets},
+    cleanup::{DespawnModEntities, DisableSystemSet, disable_mod_system_sets},
     component::WasmComponentRegistry,
     engine::{Engine, Linker, create_linker},
+    mods::{Mod, ModDespawnBehaviour},
     sandbox::Sandboxed,
     schedule::{ModSchedule, ModSchedules, ModStartup},
     setup::run_setup,
@@ -79,6 +80,7 @@ struct Inner {
     linker: Linker,
     schedules: ModSchedules,
     setup_schedule: Interned<dyn ScheduleLabel>,
+    despawn_behaviour: ModDespawnBehaviour,
 }
 
 impl Default for ModloaderPlugin {
@@ -93,11 +95,13 @@ impl ModloaderPlugin {
         let engine = Engine::new();
         let linker = create_linker(&engine);
         let setup_schedule = First.intern();
+        let despawn_behaviour = ModDespawnBehaviour::default();
         let inner = Inner {
             engine,
             linker,
             schedules,
             setup_schedule,
+            despawn_behaviour,
         };
         ModloaderPlugin(Mutex::new(Some(inner)))
     }
@@ -111,6 +115,16 @@ impl ModloaderPlugin {
         Self::new(ModSchedules::empty())
     }
 
+    /// Sets the despawn behaviour for when mods are despawned (or reloaded).
+    ///
+    /// The default behaviour is to despawn all entities the mod spawned.
+    /// See [DespawnEntities](ModDespawnBehaviour::DespawnEntities).
+    pub fn set_despawn_behaviour(mut self, despawn_behaviour: ModDespawnBehaviour) -> Self {
+        let inner = self.inner();
+        inner.despawn_behaviour = despawn_behaviour;
+        self
+    }
+
     /// Enables a new schedule with the modloader.
     ///
     /// When mods add a system to this schedule, then wasvy will automatically add them to the schedule.
@@ -161,16 +175,23 @@ impl Plugin for ModloaderPlugin {
             linker,
             schedules,
             setup_schedule,
+            despawn_behaviour,
         } = self
             .0
             .lock()
             .expect("ModloaderPlugin is not locked")
             .take()
             .expect("ModloaderPlugin is not built");
 
+        if despawn_behaviour == ModDespawnBehaviour::DespawnEntities {
+            // Registers a component that tracks mod entities and despawns them when the mod despawns
+            app.register_required_components::<Mod, DespawnModEntities>();
+        }
+
         app.init_asset::<ModAsset>()
             .register_asset_loader(ModAssetLoader { linker })
             .insert_resource(engine)
+            .insert_resource(despawn_behaviour)
             .init_resource::<WasmComponentRegistry>()
             .insert_resource(schedules)
             .add_schedule(ModStartup::new_schedule())

src/prelude.rs

@@ -1,6 +1,6 @@
 pub use crate::access::ModAccess;
 pub use crate::asset::ModAsset;
-pub use crate::mods::{Mod, ModSystemSet, Mods};
+pub use crate::mods::{Mod, ModDespawnBehaviour, ModSystemSet, Mods};
 pub use crate::plugin::ModloaderPlugin;
 pub use crate::sandbox::Sandbox;
 pub use crate::schedule::{ModSchedule, ModSchedules};

src/runner.rs

@@ -13,7 +13,8 @@ use wasmtime::component::ResourceAny;
 use wasmtime_wasi::ResourceTable;
 
 use crate::{
-    access::ModAccess, asset::ModAsset, engine::Engine, host::WasmHost, send_sync_ptr::SendSyncPtr,
+    access::ModAccess, asset::ModAsset, cleanup::InsertDespawnComponent, engine::Engine,
+    host::WasmHost, send_sync_ptr::SendSyncPtr,
 };
 
 pub(crate) type Store = wasmtime::Store<WasmHost>;
@@ -73,11 +74,13 @@ impl Runner {
                 type_registry,
                 queries,
                 access,
+                insert_despawn_component,
             }) => Inner::RunSystem {
                 commands: SendSyncPtr::new(NonNull::from_mut(commands).cast()),
                 type_registry: SendSyncPtr::new(NonNull::from_ref(type_registry)),
                 queries: SendSyncPtr::new(NonNull::from_ref(queries).cast()),
                 access,
+                insert_despawn_component,
             },
         }));
 
@@ -110,6 +113,7 @@ enum Inner {
         type_registry: SendSyncPtr<AppTypeRegistry>,
         queries: SendSyncPtr<Queries<'static, 'static>>,
         access: ModAccess,
+        insert_despawn_component: InsertDespawnComponent,
     },
 }
 
@@ -151,6 +155,7 @@ impl Data {
                 type_registry,
                 queries,
                 access,
+                insert_despawn_component,
             } =>
             // Safety: Runner::use_store ensures that this always contains a valid reference
             // See the rules here: https://doc.rust-lang.org/stable/core/ptr/index.html#pointer-to-reference-conversion
@@ -159,6 +164,7 @@ impl Data {
                     commands: commands.cast().as_mut(),
                     type_registry: type_registry.as_ref(),
                     queries: queries.cast().as_mut(),
+                    insert_despawn_component,
                     access,
                     table,
                 })
@@ -185,6 +191,7 @@ pub(crate) enum State<'a> {
         type_registry: &'a AppTypeRegistry,
         queries: &'a mut Queries<'a, 'a>,
         access: &'a ModAccess,
+        insert_despawn_component: &'a InsertDespawnComponent,
     },
 }
 
@@ -208,4 +215,5 @@ pub(crate) struct ConfigRunSystem<'a, 'b, 'c, 'd, 'e, 'f, 'g> {
     pub(crate) queries:
         &'a mut ParamSet<'d, 'e, Vec<Query<'f, 'g, FilteredEntityMut<'static, 'static>>>>,
     pub(crate) access: ModAccess,
+    pub(crate) insert_despawn_component: InsertDespawnComponent,
 }