diff --git a/mlua-sys/src/lua51/compat.rs b/mlua-sys/src/lua51/compat.rs index 8c2d7bf..29837f2 100644 --- a/mlua-sys/src/lua51/compat.rs +++ b/mlua-sys/src/lua51/compat.rs @@ -548,7 +548,7 @@ pub unsafe fn luaL_getsubtable(L: *mut lua_State, idx: c_int, fname: *const c_ch pub unsafe fn luaL_requiref(L: *mut lua_State, modname: *const c_char, openf: lua_CFunction, glb: c_int) { luaL_checkstack(L, 3, cstr!("not enough stack slots available")); - luaL_getsubtable(L, LUA_REGISTRYINDEX, cstr!("_LOADED")); + luaL_getsubtable(L, LUA_REGISTRYINDEX, LUA_LOADED_TABLE); if lua_getfield(L, -1, modname) == LUA_TNIL { lua_pop(L, 1); lua_pushcfunction(L, openf); diff --git a/mlua-sys/src/lua51/lauxlib.rs b/mlua-sys/src/lua51/lauxlib.rs index 5423885..78cef29 100644 --- a/mlua-sys/src/lua51/lauxlib.rs +++ b/mlua-sys/src/lua51/lauxlib.rs @@ -8,6 +8,9 @@ use super::lua::{self, lua_CFunction, lua_Integer, lua_Number, lua_State}; // Extra error code for 'luaL_load' pub const LUA_ERRFILE: c_int = lua::LUA_ERRERR + 1; +// Key, in the registry, for table of loaded modules +pub const LUA_LOADED_TABLE: *const c_char = cstr!("_LOADED"); + #[repr(C)] pub struct luaL_Reg { pub name: *const c_char, diff --git a/mlua-sys/src/lua52/compat.rs b/mlua-sys/src/lua52/compat.rs index 68c7029..0482914 100644 --- a/mlua-sys/src/lua52/compat.rs +++ b/mlua-sys/src/lua52/compat.rs @@ -232,7 +232,7 @@ pub unsafe fn luaL_tolstring(L: *mut lua_State, mut idx: c_int, len: *mut usize) pub unsafe fn luaL_requiref(L: *mut lua_State, modname: *const c_char, openf: lua_CFunction, glb: c_int) { luaL_checkstack(L, 3, cstr!("not enough stack slots available")); - luaL_getsubtable(L, LUA_REGISTRYINDEX, cstr!("_LOADED")); + luaL_getsubtable(L, LUA_REGISTRYINDEX, LUA_LOADED_TABLE); if lua_getfield(L, -1, modname) == LUA_TNIL { lua_pop(L, 1); lua_pushcfunction(L, openf); diff --git a/mlua-sys/src/lua52/lauxlib.rs b/mlua-sys/src/lua52/lauxlib.rs index d5cdf66..fad19eb 100644 --- a/mlua-sys/src/lua52/lauxlib.rs +++ b/mlua-sys/src/lua52/lauxlib.rs @@ -8,6 +8,12 @@ use super::lua::{self, lua_CFunction, lua_Integer, lua_Number, lua_State, lua_Un // Extra error code for 'luaL_load' pub const LUA_ERRFILE: c_int = lua::LUA_ERRERR + 1; +// Key, in the registry, for table of loaded modules +pub const LUA_LOADED_TABLE: *const c_char = cstr!("_LOADED"); + +// Key, in the registry, for table of preloaded loaders +pub const LUA_PRELOAD_TABLE: *const c_char = cstr!("_PRELOAD"); + #[repr(C)] pub struct luaL_Reg { pub name: *const c_char, diff --git a/src/state.rs b/src/state.rs index ffed018..e23faa3 100644 --- a/src/state.rs +++ b/src/state.rs @@ -2,7 +2,7 @@ use std::any::TypeId; use std::cell::{BorrowError, BorrowMutError, RefCell}; use std::marker::PhantomData; use std::ops::Deref; -use std::os::raw::c_int; +use std::os::raw::{c_char, c_int}; use std::panic::Location; use std::result::Result as StdResult; use std::{fmt, mem, ptr}; @@ -347,40 +347,78 @@ impl Lua { unsafe { self.lock().load_std_libs(libs) } } - /// Loads module `modname` into an existing Lua state using the specified entrypoint - /// function. + /// Registers module into an existing Lua state using the specified value. /// - /// Internally calls the Lua function `func` with the string `modname` as an argument, - /// sets the call result to `package.loaded[modname]` and returns copy of the result. + /// After registration, the given value will always be immediately returned when the + /// given module is [required]. /// - /// If `package.loaded[modname]` value is not nil, returns copy of the value without - /// calling the function. + /// [required]: https://www.lua.org/manual/5.4/manual.html#pdf-require + pub fn register_module(&self, modname: &str, value: impl IntoLua) -> Result<()> { + #[cfg(not(feature = "luau"))] + const LOADED_MODULES_KEY: *const c_char = ffi::LUA_LOADED_TABLE; + #[cfg(feature = "luau")] + const LOADED_MODULES_KEY: *const c_char = cstr!("_REGISTEREDMODULES"); + + if cfg!(feature = "luau") && !modname.starts_with('@') { + return Err(Error::runtime("module name must begin with '@'")); + } + unsafe { + self.exec_raw::<()>(value, |state| { + ffi::luaL_getsubtable(state, ffi::LUA_REGISTRYINDEX, LOADED_MODULES_KEY); + ffi::lua_pushlstring(state, modname.as_ptr() as *const c_char, modname.len() as _); + ffi::lua_pushvalue(state, -3); + ffi::lua_rawset(state, -3); + }) + } + } + + /// Preloads module into an existing Lua state using the specified loader function. /// - /// If the function does not return a non-nil value then this method assigns true to - /// `package.loaded[modname]`. + /// When the module is required, the loader function will be called with module name as the + /// first argument. /// - /// Behavior is similar to Lua's [`require`] function. + /// This is similar to setting the [`package.preload[modname]`] field. /// - /// [`require`]: https://www.lua.org/manual/5.4/manual.html#pdf-require - pub fn load_from_function(&self, modname: &str, func: Function) -> Result - where - T: FromLua, - { - let lua = self.lock(); - let state = lua.state(); + /// [`package.preload[modname]`]: https://www.lua.org/manual/5.4/manual.html#pdf-package.preload + #[cfg(not(feature = "luau"))] + #[cfg_attr(docsrs, doc(cfg(not(feature = "luau"))))] + pub fn preload_module(&self, modname: &str, func: Function) -> Result<()> { + #[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))] + let preload = unsafe { + self.exec_raw::>((), |state| { + ffi::lua_getfield(state, ffi::LUA_REGISTRYINDEX, ffi::LUA_PRELOAD_TABLE); + })? + }; + #[cfg(any(feature = "lua51", feature = "luajit"))] + let preload = unsafe { + self.exec_raw::>((), |state| { + if ffi::lua_getfield(state, ffi::LUA_REGISTRYINDEX, ffi::LUA_LOADED_TABLE) != ffi::LUA_TNIL { + ffi::luaL_getsubtable(state, -1, ffi::LUA_LOADLIBNAME); + ffi::luaL_getsubtable(state, -1, cstr!("preload")); + ffi::lua_rotate(state, 1, 1); + } + })? + }; + if let Some(preload) = preload { + preload.raw_set(modname, func)?; + } + Ok(()) + } + + #[doc(hidden)] + #[deprecated(since = "0.11.0", note = "Use `register_module` instead")] + #[cfg(not(feature = "luau"))] + #[cfg(not(tarpaulin_include))] + pub fn load_from_function(&self, modname: &str, func: Function) -> Result { let loaded = unsafe { - let _sg = StackGuard::new(state); - check_stack(state, 2)?; - protect_lua!(state, 0, 1, fn(state) { - ffi::luaL_getsubtable(state, ffi::LUA_REGISTRYINDEX, cstr!("_LOADED")); - })?; - Table(lua.pop_ref()) + self.exec_raw::((), |state| { + ffi::luaL_getsubtable(state, ffi::LUA_REGISTRYINDEX, ffi::LUA_LOADED_TABLE); + })? }; - let modname = unsafe { lua.create_string(modname)? }; - let value = match loaded.raw_get(&modname)? { + let value = match loaded.raw_get(modname)? { Value::Nil => { - let result = match func.call(&modname)? { + let result = match func.call(modname)? { Value::Nil => Value::Boolean(true), res => res, }; @@ -394,24 +432,14 @@ impl Lua { /// Unloads module `modname`. /// - /// Removes module from the [`package.loaded`] table which allows to load it again. - /// It does not support unloading binary Lua modules since they are internally cached and can be - /// unloaded only by closing Lua state. + /// This method does not support unloading binary Lua modules since they are internally cached + /// and can be unloaded only by closing Lua state. + /// + /// This is similar to calling [`Lua::register_module`] with `Nil` value. /// /// [`package.loaded`]: https://www.lua.org/manual/5.4/manual.html#pdf-package.loaded - pub fn unload(&self, modname: &str) -> Result<()> { - let lua = self.lock(); - let state = lua.state(); - let loaded = unsafe { - let _sg = StackGuard::new(state); - check_stack(state, 2)?; - protect_lua!(state, 0, 1, fn(state) { - ffi::luaL_getsubtable(state, ffi::LUA_REGISTRYINDEX, cstr!("_LOADED")); - })?; - Table(lua.pop_ref()) - }; - - loaded.raw_set(modname, Nil) + pub fn unload_module(&self, modname: &str) -> Result<()> { + self.register_module(modname, Nil) } // Executes module entrypoint function, which returns only one Value. diff --git a/tests/tests.rs b/tests/tests.rs index 9444299..61021ad 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -3,7 +3,6 @@ use std::collections::HashMap; use std::iter::FromIterator; use std::panic::{catch_unwind, AssertUnwindSafe}; use std::string::String as StdString; -use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::Arc; use std::{error, f32, f64, fmt}; @@ -1168,36 +1167,79 @@ fn test_jit_version() -> Result<()> { } #[test] -fn test_load_from_function() -> Result<()> { +fn test_register_module() -> Result<()> { let lua = Lua::new(); - let i = Arc::new(AtomicU32::new(0)); - let i2 = i.clone(); - let func = lua.create_function(move |lua, modname: String| { - i2.fetch_add(1, Ordering::Relaxed); + let t = lua.create_table()?; + t.set("name", "my_module")?; + lua.register_module("@my_module", &t)?; + + lua.load( + r#" + local my_module = require("@my_module") + assert(my_module.name == "my_module") + "#, + ) + .exec()?; + + lua.unload_module("@my_module")?; + lua.load( + r#" + local ok, err = pcall(function() return require("@my_module") end) + assert(not ok) + "#, + ) + .exec()?; + + #[cfg(feature = "luau")] + { + // Luau registered modules must have '@' prefix + let res = lua.register_module("my_module", 123); + assert!(res.is_err()); + assert_eq!( + res.unwrap_err().to_string(), + "runtime error: module name must begin with '@'" + ); + } + + Ok(()) +} + +#[test] +#[cfg(not(feature = "luau"))] +fn test_preload_module() -> Result<()> { + let lua = Lua::new(); + + let loader = lua.create_function(move |lua, modname: String| { let t = lua.create_table()?; - t.set("__name", modname)?; + t.set("name", modname)?; Ok(t) })?; - let t: Table = lua.load_from_function("my_module", func.clone())?; - assert_eq!(t.get::("__name")?, "my_module"); - assert_eq!(i.load(Ordering::Relaxed), 1); - - let _: Value = lua.load_from_function("my_module", func.clone())?; - assert_eq!(i.load(Ordering::Relaxed), 1); - - let func_nil = lua.create_function(move |_, _: String| Ok(Value::Nil))?; - let v: Value = lua.load_from_function("my_module2", func_nil)?; - assert_eq!(v, Value::Boolean(true)); + lua.preload_module("@my_module", loader.clone())?; + lua.load( + r#" + -- `my_module` is global for purposes of next test + my_module = require("@my_module") + assert(my_module.name == "@my_module") + local my_module2 = require("@my_module") + assert(my_module == my_module2) + "#, + ) + .exec() + .unwrap(); // Test unloading and loading again - lua.unload("my_module")?; - let _: Value = lua.load_from_function("my_module", func)?; - assert_eq!(i.load(Ordering::Relaxed), 2); - - // Unloading nonexistent module must not fail - lua.unload("my_module2")?; + lua.unload_module("@my_module")?; + lua.load( + r#" + local my_module3 = require("@my_module") + -- `my_module` is not equal to `my_module3` because it was reloaded + assert(my_module ~= my_module3) + "#, + ) + .exec() + .unwrap(); Ok(()) }