From ed796fff51ebf12b3d100005d5159bf945840eb5 Mon Sep 17 00:00:00 2001 From: Alex Orlenko Date: Sun, 4 May 2025 11:00:59 +0100 Subject: [PATCH] Support async require loaders for Luau --- mlua-sys/src/luau/luarequire.rs | 2 + src/luau/mod.rs | 22 +------ src/luau/require.rs | 110 ++++++++++++++++++++++++++++++-- src/state.rs | 2 +- tests/luau/require.rs | 43 ++++++++++++- 5 files changed, 150 insertions(+), 29 deletions(-) diff --git a/mlua-sys/src/luau/luarequire.rs b/mlua-sys/src/luau/luarequire.rs index 30a7267..dcdd0b8 100644 --- a/mlua-sys/src/luau/luarequire.rs +++ b/mlua-sys/src/luau/luarequire.rs @@ -4,6 +4,8 @@ use std::os::raw::{c_char, c_int, c_void}; use super::lua::lua_State; +pub const LUA_REGISTERED_MODULES_TABLE: *const c_char = cstr!("_REGISTEREDMODULES"); + #[repr(C)] pub enum luarequire_NavigateResult { Success, diff --git a/src/luau/mod.rs b/src/luau/mod.rs index 592059c..635eebf 100644 --- a/src/luau/mod.rs +++ b/src/luau/mod.rs @@ -1,5 +1,4 @@ use std::ffi::CStr; -use std::mem; use std::os::raw::c_int; use crate::error::Result; @@ -16,26 +15,7 @@ impl Lua { #[cfg(any(feature = "luau", doc))] #[cfg_attr(docsrs, doc(cfg(feature = "luau")))] pub fn create_require_function(&self, require: R) -> Result { - unsafe extern "C-unwind" fn mlua_require(state: *mut ffi::lua_State) -> c_int { - let mut ar: ffi::lua_Debug = mem::zeroed(); - if ffi::lua_getinfo(state, 1, cstr!("s"), &mut ar) == 0 { - ffi::luaL_error(state, cstr!("require is not supported in this context")); - } - let top = ffi::lua_gettop(state); - ffi::lua_pushvalue(state, ffi::lua_upvalueindex(2)); // the "proxy" require function - ffi::lua_pushvalue(state, 1); // require path - ffi::lua_pushstring(state, ar.source); // current file - ffi::lua_call(state, 2, ffi::LUA_MULTRET); - ffi::lua_gettop(state) - top - } - - unsafe { - self.exec_raw((), move |state| { - let requirer_ptr = ffi::lua_newuserdata_t::>(state, Box::new(require)); - ffi::luarequire_pushproxyrequire(state, require::init_config, requirer_ptr as *mut _); - ffi::lua_pushcclosured(state, mlua_require, cstr!("require"), 2); - }) - } + require::create_require_function(self, require) } pub(crate) unsafe fn configure_luau(&self) -> Result<()> { diff --git a/src/luau/require.rs b/src/luau/require.rs index f290d53..152114d 100644 --- a/src/luau/require.rs +++ b/src/luau/require.rs @@ -5,13 +5,13 @@ use std::io::Result as IoResult; use std::os::raw::{c_char, c_int, c_void}; use std::path::{Component, Path, PathBuf}; use std::result::Result as StdResult; -use std::{env, fmt, fs, ptr}; +use std::{env, fmt, fs, mem, ptr}; use crate::error::Result; use crate::function::Function; use crate::state::{callback_error_ext, Lua}; +use crate::table::Table; use crate::types::MaybeSend; -use crate::value::Value; /// An error that can occur during navigation in the Luau `require` system. pub enum NavigateError { @@ -87,6 +87,8 @@ pub trait Require: MaybeSend { fn config(&self) -> IoResult>; /// Returns a loader that when called, loads the module and returns the result. + /// + /// Loader can be sync or async. fn loader(&self, lua: &Lua, path: &str, chunk_name: &str, content: &[u8]) -> Result { let _ = path; lua.load(content).set_name(chunk_name).into_function() @@ -425,10 +427,7 @@ pub(super) unsafe extern "C" fn init_config(config: *mut ffi::luarequire_Configu let contents = CStr::from_ptr(contents).to_bytes(); callback_error_ext(state, ptr::null_mut(), false, move |extra, _| { let rawlua = (*extra).raw_lua(); - match (this.loader(rawlua.lua(), &path, &chunk_name, contents)?).call(())? { - Value::Nil => rawlua.push(true)?, - value => rawlua.push_value(&value)?, - }; + rawlua.push(this.loader(rawlua.lua(), &path, &chunk_name, contents)?)?; Ok(1) }) } @@ -495,6 +494,105 @@ unsafe fn write_to_buffer( } } +#[cfg(feature = "luau")] +pub fn create_require_function(lua: &Lua, require: R) -> Result { + unsafe extern "C-unwind" fn find_current_file(state: *mut ffi::lua_State) -> c_int { + let mut ar: ffi::lua_Debug = mem::zeroed(); + for level in 2.. { + if ffi::lua_getinfo(state, level, cstr!("s"), &mut ar) == 0 { + ffi::luaL_error(state, cstr!("require is not supported in this context")); + } + if CStr::from_ptr(ar.what) != c"C" { + break; + } + } + ffi::lua_pushstring(state, ar.source); + 1 + } + + unsafe extern "C-unwind" fn get_cache_key(state: *mut ffi::lua_State) -> c_int { + let requirer = ffi::lua_touserdata(state, ffi::lua_upvalueindex(1)) as *const Box; + let cache_key = (*requirer).cache_key(); + ffi::lua_pushlstring(state, cache_key.as_ptr() as *const _, cache_key.len()); + 1 + } + + let (get_cache_key, find_current_file, proxyrequire, registered_modules, loader_cache) = unsafe { + lua.exec_raw::<(Function, Function, Function, Table, Table)>((), move |state| { + let requirer_ptr = ffi::lua_newuserdata_t::>(state, Box::new(require)); + ffi::lua_pushcclosured(state, get_cache_key, cstr!("get_cache_key"), 1); + ffi::lua_pushcfunctiond(state, find_current_file, cstr!("find_current_file")); + ffi::luarequire_pushproxyrequire(state, init_config, requirer_ptr as *mut _); + ffi::luaL_getsubtable(state, ffi::LUA_REGISTRYINDEX, ffi::LUA_REGISTERED_MODULES_TABLE); + ffi::luaL_getsubtable(state, ffi::LUA_REGISTRYINDEX, cstr!("__MLUA_LOADER_CACHE")); + }) + }?; + + unsafe extern "C-unwind" fn error(state: *mut ffi::lua_State) -> c_int { + ffi::luaL_where(state, 1); + ffi::lua_pushvalue(state, 1); + ffi::lua_concat(state, 2); + ffi::lua_error(state); + } + + unsafe extern "C-unwind" fn r#type(state: *mut ffi::lua_State) -> c_int { + ffi::lua_pushstring(state, ffi::lua_typename(state, ffi::lua_type(state, 1))); + 1 + } + + let (error, r#type) = unsafe { + lua.exec_raw::<(Function, Function)>((), move |state| { + ffi::lua_pushcfunctiond(state, error, cstr!("error")); + ffi::lua_pushcfunctiond(state, r#type, cstr!("type")); + }) + }?; + + // Prepare environment for the "require" function + let env = lua.create_table_with_capacity(0, 7)?; + env.raw_set("get_cache_key", get_cache_key)?; + env.raw_set("find_current_file", find_current_file)?; + env.raw_set("proxyrequire", proxyrequire)?; + env.raw_set("REGISTERED_MODULES", registered_modules)?; + env.raw_set("LOADER_CACHE", loader_cache)?; + env.raw_set("error", error)?; + env.raw_set("type", r#type)?; + + lua.load( + r#" + local path = ... + if type(path) ~= "string" then + error("bad argument #1 to 'require' (string expected, got " .. type(path) .. ")") + end + + -- Check if the module (path) is explicitly registered + local maybe_result = REGISTERED_MODULES[path] + if maybe_result ~= nil then + return maybe_result + end + + local loader = proxyrequire(path, find_current_file()) + local cache_key = get_cache_key() + -- Check if the loader result is already cached + local result = LOADER_CACHE[cache_key] + if result ~= nil then + return result + end + + -- Call the loader function and cache the result + result = loader() + if result == nil then + result = true + end + LOADER_CACHE[cache_key] = result + return result + "#, + ) + .try_cache() + .set_name("=__mlua_require") + .set_environment(env) + .into_function() +} + #[cfg(test)] mod tests { use std::path::Path; diff --git a/src/state.rs b/src/state.rs index 684eceb..69b39e6 100644 --- a/src/state.rs +++ b/src/state.rs @@ -356,7 +356,7 @@ impl Lua { #[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"); + const LOADED_MODULES_KEY: *const c_char = ffi::LUA_REGISTERED_MODULES_TABLE; if cfg!(feature = "luau") && !modname.starts_with('@') { return Err(Error::runtime("module name must begin with '@'")); diff --git a/tests/luau/require.rs b/tests/luau/require.rs index f17835b..a38b0cd 100644 --- a/tests/luau/require.rs +++ b/tests/luau/require.rs @@ -1,6 +1,6 @@ use mlua::{IntoLua, Lua, Result, Value}; -fn run_require(lua: &Lua, path: &str) -> Result { +fn run_require(lua: &Lua, path: impl IntoLua) -> Result { lua.load(r#"return require(...)"#).call(path) } @@ -26,6 +26,12 @@ fn test_require_errors() { assert!( (res.unwrap_err().to_string()).contains("require path must start with a valid prefix: ./, ../, or @") ); + + // Pass non-string to require + let res = run_require(&lua, true); + assert!(res.is_err()); + assert!((res.unwrap_err().to_string()) + .contains("bad argument #1 to 'require' (string expected, got boolean)")); } #[test] @@ -100,3 +106,38 @@ fn test_require_with_config() { assert!(res.is_err()); assert!((res.unwrap_err().to_string()).contains("@ is not a valid alias")); } + +#[cfg(feature = "async")] +#[tokio::test] +async fn test_async_require() -> Result<()> { + let lua = Lua::new(); + + let temp_dir = tempfile::tempdir().unwrap(); + let temp_path = temp_dir.path().join("async_chunk.luau"); + std::fs::write( + &temp_path, + r#" + sleep_ms(10) + return "result_after_async_sleep" + "#, + ) + .unwrap(); + + lua.globals().set( + "sleep_ms", + lua.create_async_function(|_, ms: u64| async move { + tokio::time::sleep(std::time::Duration::from_millis(ms)).await; + Ok(()) + })?, + )?; + + lua.load( + r#" + local result = require("./async_chunk") + assert(result == "result_after_async_sleep") + "#, + ) + .set_name(format!("@{}", temp_dir.path().join("require.rs").display())) + .exec_async() + .await +}