fix!: script contexts being completely overwritten on a re-load (#345)

Author: makspll

crates/languages/bevy_mod_scripting_lua/src/lib.rs

@@ -156,35 +156,54 @@ impl Plugin for LuaScriptingPlugin {
         self.scripting_plugin.finish(app);
     }
 }
-#[profiling::function]
-/// Load a lua context from a script
-pub fn lua_context_load(
+
+fn load_lua_content_into_context(
+    context: &mut Lua,
     script_id: &ScriptId,
     content: &[u8],
     initializers: &[ContextInitializer<LuaScriptingPlugin>],
     pre_handling_initializers: &[ContextPreHandlingInitializer<LuaScriptingPlugin>],
-    _: &mut (),
-) -> Result<Lua, ScriptError> {
-    #[cfg(feature = "unsafe_lua_modules")]
-    let mut context = unsafe { Lua::unsafe_new() };
-    #[cfg(not(feature = "unsafe_lua_modules"))]
-    let mut context = Lua::new();
-
+) -> Result<(), ScriptError> {
     initializers
         .iter()
-        .try_for_each(|init| init(script_id, &mut context))?;
+        .try_for_each(|init| init(script_id, context))?;
 
     pre_handling_initializers
         .iter()
-        .try_for_each(|init| init(script_id, Entity::from_raw(0), &mut context))?;
+        .try_for_each(|init| init(script_id, Entity::from_raw(0), context))?;
 
     context
         .load(content)
         .exec()
         .map_err(ScriptError::from_mlua_error)?;
 
+    Ok(())
+}
+
+#[profiling::function]
+/// Load a lua context from a script
+pub fn lua_context_load(
+    script_id: &ScriptId,
+    content: &[u8],
+    initializers: &[ContextInitializer<LuaScriptingPlugin>],
+    pre_handling_initializers: &[ContextPreHandlingInitializer<LuaScriptingPlugin>],
+    _: &mut (),
+) -> Result<Lua, ScriptError> {
+    #[cfg(feature = "unsafe_lua_modules")]
+    let mut context = unsafe { Lua::unsafe_new() };
+    #[cfg(not(feature = "unsafe_lua_modules"))]
+    let mut context = Lua::new();
+
+    load_lua_content_into_context(
+        &mut context,
+        script_id,
+        content,
+        initializers,
+        pre_handling_initializers,
+    )?;
     Ok(context)
 }
+
 #[profiling::function]
 /// Reload a lua context from a script
 pub fn lua_context_reload(
@@ -195,12 +214,12 @@ pub fn lua_context_reload(
     pre_handling_initializers: &[ContextPreHandlingInitializer<LuaScriptingPlugin>],
     _: &mut (),
 ) -> Result<(), ScriptError> {
-    *old_ctxt = lua_context_load(
+    load_lua_content_into_context(
+        old_ctxt,
         script,
         content,
         initializers,
         pre_handling_initializers,
-        &mut (),
     )?;
     Ok(())
 }
@@ -243,3 +262,49 @@ pub fn lua_handler(
     let out = handler.call::<LuaScriptValue>(input)?;
     Ok(out.into())
 }
+
+#[cfg(test)]
+mod test {
+    use mlua::Value;
+
+    use super::*;
+
+    #[test]
+    fn test_reload_doesnt_overwrite_old_context() {
+        let lua = Lua::new();
+        let script_id = ScriptId::from("asd.lua");
+        let initializers = vec![];
+        let pre_handling_initializers = vec![];
+        let mut old_ctxt = lua.clone();
+
+        lua_context_load(
+            &script_id,
+            "function hello_world_from_first_load()
+            
+            end"
+            .as_bytes(),
+            &initializers,
+            &pre_handling_initializers,
+            &mut (),
+        )
+        .unwrap();
+
+        lua_context_reload(
+            &script_id,
+            "function hello_world_from_second_load()
+            
+            end"
+            .as_bytes(),
+            &mut old_ctxt,
+            &initializers,
+            &pre_handling_initializers,
+            &mut (),
+        )
+        .unwrap();
+
+        // assert both functions exist in globals
+        let globals = lua.globals();
+        assert!(globals.get::<Value>("hello_world_from_first_load").is_ok());
+        assert!(globals.get::<Value>("hello_world_from_second_load").is_ok());
+    }
+}

crates/languages/bevy_mod_scripting_rhai/src/lib.rs

@@ -179,37 +179,53 @@ impl Plugin for RhaiScriptingPlugin {
     }
 }
 
-/// Load a rhai context from a script.
-pub fn rhai_context_load(
+// NEW helper function to load content into an existing context without clearing previous definitions.
+fn load_rhai_content_into_context(
+    context: &mut RhaiScriptContext,
     script: &ScriptId,
     content: &[u8],
     initializers: &[ContextInitializer<RhaiScriptingPlugin>],
     pre_handling_initializers: &[ContextPreHandlingInitializer<RhaiScriptingPlugin>],
     runtime: &mut RhaiRuntime,
-) -> Result<RhaiScriptContext, ScriptError> {
+) -> Result<(), ScriptError> {
     let mut ast = runtime.compile(std::str::from_utf8(content)?)?;
     ast.set_source(script.to_string());
-
-    let mut context = RhaiScriptContext {
-        ast,
-        scope: Scope::new(),
-    };
     initializers
         .iter()
-        .try_for_each(|init| init(script, &mut context))?;
-
+        .try_for_each(|init| init(script, context))?;
     pre_handling_initializers
         .iter()
-        .try_for_each(|init| init(script, Entity::from_raw(0), &mut context))?;
-
-    runtime.eval_ast_with_scope(&mut context.scope, &context.ast)?;
-    // do not invoke top level statements after the first time we run the script
-    context.ast.clear_statements();
+        .try_for_each(|init| init(script, Entity::from_raw(0), context))?;
+    runtime.eval_ast_with_scope(&mut context.scope, &ast)?;
+    // Unlike before, do not clear statements so that definitions persist.
+    Ok(())
+}
 
+/// Load a rhai context from a script.
+pub fn rhai_context_load(
+    script: &ScriptId,
+    content: &[u8],
+    initializers: &[ContextInitializer<RhaiScriptingPlugin>],
+    pre_handling_initializers: &[ContextPreHandlingInitializer<RhaiScriptingPlugin>],
+    runtime: &mut RhaiRuntime,
+) -> Result<RhaiScriptContext, ScriptError> {
+    let mut context = RhaiScriptContext {
+        // Using an empty AST as a placeholder.
+        ast: AST::empty(),
+        scope: Scope::new(),
+    };
+    load_rhai_content_into_context(
+        &mut context,
+        script,
+        content,
+        initializers,
+        pre_handling_initializers,
+        runtime,
+    )?;
     Ok(context)
 }
 
-/// Reload a rhai context from a script.
+/// Reload a rhai context from a script. New content is appended to the existing context.
 pub fn rhai_context_reload(
     script: &ScriptId,
     content: &[u8],
@@ -218,14 +234,14 @@ pub fn rhai_context_reload(
     pre_handling_initializers: &[ContextPreHandlingInitializer<RhaiScriptingPlugin>],
     runtime: &mut RhaiRuntime,
 ) -> Result<(), ScriptError> {
-    *context = rhai_context_load(
+    load_rhai_content_into_context(
+        context,
         script,
         content,
         initializers,
         pre_handling_initializers,
         runtime,
-    )?;
-    Ok(())
+    )
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -278,3 +294,45 @@ pub fn rhai_callback_handler(
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_reload_doesnt_overwrite_old_context() {
+        let mut runtime = RhaiRuntime::new();
+        let script_id = ScriptId::from("asd.rhai");
+        let initializers: Vec<ContextInitializer<RhaiScriptingPlugin>> = vec![];
+        let pre_handling_initializers: Vec<ContextPreHandlingInitializer<RhaiScriptingPlugin>> =
+            vec![];
+
+        // Load first content defining a function that returns 42.
+        let mut context = rhai_context_load(
+            &script_id,
+            b"let hello = 2;",
+            &initializers,
+            &pre_handling_initializers,
+            &mut runtime,
+        )
+        .unwrap();
+
+        // Reload with additional content defining a second function that returns 24.
+        rhai_context_reload(
+            &script_id,
+            b"let hello2 = 3",
+            &mut context,
+            &initializers,
+            &pre_handling_initializers,
+            &mut runtime,
+        )
+        .unwrap();
+
+        // get first var
+        let hello = context.scope.get_value::<i64>("hello").unwrap();
+        assert_eq!(hello, 2);
+        // get second var
+        let hello2 = context.scope.get_value::<i64>("hello2").unwrap();
+        assert_eq!(hello2, 3);
+    }
+}