Fix removal of script contexts

assets/scripts/multiple_events_rhai.rhai

@@ -4,4 +4,6 @@ fn on_init(name) {
 
 fn on_update(name, delta) {
     print(`Hello World! From "${name}" in Update: ${delta}`);
+
+    world.despawn_recursive(entity);
 }

bevy_mod_scripting_core/src/systems.rs

@@ -95,11 +95,20 @@ pub fn script_remove_synchronizer<H: ScriptHost>(
     mut query: RemovedComponents<ScriptCollection<H::ScriptAsset>>,
     mut contexts: ResMut<ScriptContexts<H::ScriptContext>>,
 ) {
-    query.iter().for_each(|v| {
+    for v in query.iter() {
         // we know that this entity used to have a script component
         // ergo a script context must exist in ctxts, remove all scripts on the entity
-        contexts.remove_context(v.index());
-    })
+        let mut script_ids = Vec::new();
+        for (script_id, (entity, ..)) in contexts.context_entities.iter() {
+            if entity.index() == v.index() {
+                script_ids.push(*script_id);
+            }
+        }
+
+        for script_id in script_ids {
+            contexts.remove_context(script_id);
+        }
+    }
 }
 
 /// Reloads hot-reloaded scripts, or loads missing contexts for scripts which were added but not loaded

examples/rhai/multiple_events_rhai.rs

@@ -1,5 +1,6 @@
 use bevy::{prelude::*, reflect::Reflect};
 use bevy_mod_scripting_core::{
+    event::ScriptLoaded,
     prelude::{APIProvider, PriorityEventWriter, Recipients, Script, ScriptCollection},
     AddScriptApiProvider, AddScriptHost, AddScriptHostHandler, ScriptingPlugin,
 };
@@ -17,7 +18,7 @@ fn main() {
         .add_systems(Update, (call_init, call_update))
         .add_script_host::<RhaiScriptHost<ScriptArgs>>(PostUpdate)
         .add_api_provider::<RhaiScriptHost<ScriptArgs>>(Box::new(RhaiBevyAPIProvider))
-        .add_script_handler::<RhaiScriptHost<ScriptArgs>, 0, 1>(PostUpdate)
+        .add_script_handler::<RhaiScriptHost<ScriptArgs>, 0, 2>(PostUpdate)
         .run();
 }
 
@@ -106,28 +107,37 @@ fn call_update(
 
 fn call_init(
     mut events: PriorityEventWriter<RhaiEvent<ScriptArgs>>,
-    mut commands: Commands,
+    mut loaded_scripts: EventReader<ScriptLoaded>,
     entity_query: Query<
         (Entity, Option<&Name>, Option<&ScriptCollection<RhaiFile>>),
-        Added<NewlyAddedEntityCallInit>,
+        With<NewlyAddedEntityCallInit>,
     >,
+    mut commands: Commands,
 ) {
-    entity_query.for_each(|(entity, name, scripts)| {
-        if scripts.is_some() {
-            events.send(
-                RhaiEvent {
-                    hook_name: "on_init".to_owned(),
-                    args: ScriptArgs {
-                        delta_time: None,
-                        entity_name: name.map(|n| n.to_string()),
-                    },
-                    recipients: Recipients::Entity(entity),
-                },
-                0,
-            );
-            commands.entity(entity).remove::<NewlyAddedEntityCallInit>();
-        } else {
-            commands.entity(entity).remove::<NewlyAddedEntityCallInit>();
+    // Call init on all entities that have a script that was just loaded, and remove the Marker component, so that Update can be called
+    'outer: for loaded_script in loaded_scripts.iter() {
+        for (entity, name, scripts) in entity_query.iter() {
+            if let Some(scripts) = scripts {
+                if scripts.scripts.iter().any(|s| s.id() == loaded_script.sid) {
+                    events.send(
+                        RhaiEvent {
+                            hook_name: "on_init".to_owned(),
+                            args: ScriptArgs {
+                                entity_name: name.map(|n| n.to_string()),
+                                ..Default::default()
+                            },
+                            recipients: Recipients::ScriptID(loaded_script.sid),
+                        },
+                        0,
+                    );
+
+                    commands.entity(entity).remove::<NewlyAddedEntityCallInit>();
+                    continue 'outer;
+                }
+            } else {
+                commands.entity(entity).remove::<NewlyAddedEntityCallInit>();
+                continue 'outer;
+            }
         }
-    });
+    }
 }