Add package module to Luau

Introduce module loaders
Support loading binary modules

Author: khvzak

Cargo.toml

@@ -32,7 +32,7 @@ lua52 = ["ffi/lua52"]
 lua51 = ["ffi/lua51"]
 luajit = ["ffi/luajit"]
 luajit52 = ["luajit", "ffi/luajit52"]
-luau = ["ffi/luau"]
+luau = ["ffi/luau", "libloading"]
 luau-jit = ["luau", "ffi/luau-codegen"]
 luau-vector4 = ["luau", "ffi/luau-vector4"]
 vendored = ["ffi/vendored"]
@@ -57,6 +57,9 @@ parking_lot = { version = "0.12", optional = true }
 
 ffi = { package = "mlua-sys", version = "0.3.2", path = "mlua-sys" }
 
+[target.'cfg(unix)'.dependencies]
+libloading = { version = "0.8", optional = true }
+
 [dev-dependencies]
 rustyline = "12.0"
 criterion = { version = "0.5", features = ["async_tokio"] }

mlua-sys/Cargo.toml

@@ -40,4 +40,4 @@ cfg-if = "1.0"
 pkg-config = "0.3.17"
 lua-src = { version = ">= 546.0.0, < 546.1.0", optional = true }
 luajit-src = { version = ">= 210.5.0, < 210.6.0", optional = true }
-luau0-src = { version = "0.7.0", optional = true }
+luau0-src = { version = "0.7.7", optional = true }

mlua-sys/build/main_inner.rs

@@ -9,8 +9,8 @@ cfg_if::cfg_if! {
 }
 
 fn main() {
-    #[cfg(all(feature = "luau", feature = "module"))]
-    compile_error!("Luau does not support `module` mode");
+    #[cfg(all(feature = "luau", feature = "module", windows))]
+    compile_error!("Luau does not support `module` mode on Windows");
 
     #[cfg(all(feature = "module", feature = "vendored"))]
     compile_error!("`vendored` and `module` features are mutually exclusive");

mlua-sys/src/luau/lua.rs

@@ -2,7 +2,7 @@
 
 use std::marker::{PhantomData, PhantomPinned};
 use std::os::raw::{c_char, c_double, c_float, c_int, c_uint, c_void};
-use std::ptr;
+use std::{mem, ptr};
 
 // Option for multiple returns in 'lua_pcall' and 'lua_call'
 pub const LUA_MULTRET: c_int = -1;
@@ -278,6 +278,7 @@ extern "C-unwind" {
     pub fn lua_getuserdatadtor(L: *mut lua_State, tag: c_int) -> Option<lua_Destructor>;
     pub fn lua_clonefunction(L: *mut lua_State, idx: c_int);
     pub fn lua_cleartable(L: *mut lua_State, idx: c_int);
+    pub fn lua_getallocf(L: *mut lua_State, ud: *mut *mut c_void) -> lua_Alloc;
 }
 
 //
@@ -325,6 +326,15 @@ pub unsafe fn lua_newuserdata(L: *mut lua_State, sz: usize) -> *mut c_void {
     lua_newuserdatatagged(L, sz, 0)
 }
 
+#[inline(always)]
+pub unsafe fn lua_newuserdata_t<T>(L: *mut lua_State) -> *mut T {
+    unsafe extern "C-unwind" fn destructor<T>(ud: *mut c_void) {
+        ptr::drop_in_place(ud as *mut T);
+    }
+
+    lua_newuserdatadtor(L, mem::size_of::<T>(), destructor::<T>) as *mut T
+}
+
 // TODO: lua_strlen
 
 #[inline(always)]

src/lua.rs

@@ -1,17 +1,15 @@
 use std::any::TypeId;
 use std::cell::{RefCell, UnsafeCell};
 use std::ffi::{CStr, CString};
-use std::fmt;
 use std::marker::PhantomData;
 use std::mem::MaybeUninit;
 use std::ops::Deref;
 use std::os::raw::{c_char, c_int, c_void};
 use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe, Location};
-use std::ptr::NonNull;
 use std::result::Result as StdResult;
 use std::sync::atomic::{AtomicPtr, Ordering};
 use std::sync::{Arc, Mutex};
-use std::{mem, ptr, str};
+use std::{fmt, mem, ptr, str};
 
 use rustc_hash::FxHashMap;
 
@@ -60,6 +58,7 @@ use {
     crate::types::{AsyncCallback, AsyncCallbackUpvalue, AsyncPollUpvalue},
     futures_util::future::{self, Future},
     futures_util::task::{noop_waker_ref, Context, Poll, Waker},
+    std::ptr::NonNull,
 };
 
 #[cfg(feature = "serialize")]
@@ -94,7 +93,6 @@ pub(crate) struct ExtraData {
 
     safe: bool,
     libs: StdLib,
-    mem_state: Option<NonNull<MemoryState>>,
     #[cfg(feature = "module")]
     skip_memory_check: bool,
 
@@ -244,11 +242,14 @@ impl Drop for Lua {
 impl Drop for LuaInner {
     fn drop(&mut self) {
         unsafe {
-            #[cfg(feature = "luau")]
-            {
-                (*ffi::lua_callbacks(self.state())).userdata = ptr::null_mut();
-            }
+            let mem_state = MemoryState::get(self.main_state);
+
             ffi::lua_close(self.main_state);
+
+            // Deallocate MemoryState if it was created by us
+            if !mem_state.is_null() {
+                drop(Box::from_raw(mem_state));
+            }
         }
     }
 }
@@ -261,9 +262,6 @@ impl Drop for ExtraData {
         }
 
         *mlua_expect!(self.registry_unref_list.lock(), "unref list poisoned") = None;
-        if let Some(mem_state) = self.mem_state {
-            drop(unsafe { Box::from_raw(mem_state.as_ptr()) });
-        }
     }
 }
 
@@ -359,36 +357,34 @@ impl Lua {
     ///
     /// [`StdLib`]: crate::StdLib
     pub unsafe fn unsafe_new_with(libs: StdLib, options: LuaOptions) -> Lua {
+        // Workaround to avoid stripping a few unused Lua symbols that could be imported
+        // by C modules in unsafe mode
+        let mut _symbols: Vec<*const extern "C-unwind" fn()> =
+            vec![ffi::lua_isuserdata as _, ffi::lua_tocfunction as _];
+
         #[cfg(not(feature = "luau"))]
+        _symbols.extend_from_slice(&[
+            ffi::lua_atpanic as _,
+            ffi::luaL_loadstring as _,
+            ffi::luaL_openlibs as _,
+        ]);
+        #[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))]
         {
-            // Workaround to avoid stripping a few unused Lua symbols that could be imported
-            // by C modules in unsafe mode
-            let mut _symbols: Vec<*const extern "C-unwind" fn()> = vec![
-                ffi::lua_atpanic as _,
-                ffi::lua_isuserdata as _,
-                ffi::lua_tocfunction as _,
-                ffi::luaL_loadstring as _,
-                ffi::luaL_openlibs as _,
-            ];
-            #[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))]
-            {
-                _symbols.push(ffi::lua_getglobal as _);
-                _symbols.push(ffi::lua_setglobal as _);
-                _symbols.push(ffi::luaL_setfuncs as _);
-            }
+            _symbols.push(ffi::lua_getglobal as _);
+            _symbols.push(ffi::lua_setglobal as _);
+            _symbols.push(ffi::luaL_setfuncs as _);
         }
 
         Self::inner_new(libs, options)
     }
 
     /// Creates a new Lua state with required `libs` and `options`
     unsafe fn inner_new(libs: StdLib, options: LuaOptions) -> Lua {
-        let mut mem_state: *mut MemoryState = Box::into_raw(Box::default());
+        let mem_state: *mut MemoryState = Box::into_raw(Box::default());
         let mut state = ffi::lua_newstate(ALLOCATOR, mem_state as *mut c_void);
         // If state is null then switch to Lua internal allocator
         if state.is_null() {
             drop(Box::from_raw(mem_state));
-            mem_state = ptr::null_mut();
             state = ffi::luaL_newstate();
         }
         assert!(!state.is_null(), "Failed to instantiate Lua VM");
@@ -404,7 +400,6 @@ impl Lua {
 
         let lua = Lua::init_from_ptr(state);
         let extra = lua.extra.get();
-        (*extra).mem_state = NonNull::new(mem_state);
 
         mlua_expect!(
             load_from_std_lib(state, libs),
@@ -440,7 +435,7 @@ impl Lua {
         }
 
         #[cfg(feature = "luau")]
-        mlua_expect!(lua.prepare_luau_state(), "Error preparing Luau state");
+        mlua_expect!(lua.prepare_luau_state(), "Error configuring Luau");
 
         lua
     }
@@ -514,7 +509,6 @@ impl Lua {
             app_data: AppData::default(),
             safe: false,
             libs: StdLib::NONE,
-            mem_state: None,
             #[cfg(feature = "module")]
             skip_memory_check: false,
             ref_thread,
@@ -547,14 +541,8 @@ impl Lua {
 
         // Store it in the registry
         mlua_expect!(
-            (|state| {
-                push_gc_userdata(state, Arc::clone(&extra), true)?;
-                protect_lua!(state, 1, 0, fn(state) {
-                    let extra_key = &EXTRA_REGISTRY_KEY as *const u8 as *const c_void;
-                    ffi::lua_rawsetp(state, ffi::LUA_REGISTRYINDEX, extra_key);
-                })
-            })(main_state),
-            "Error while storing extra data",
+            set_extra_data(main_state, &extra),
+            "Error while storing extra data"
         );
 
         // Register `DestructedUserdata` type
@@ -572,13 +560,6 @@ impl Lua {
         );
         assert_stack(main_state, ffi::LUA_MINSTACK);
 
-        // Set Luau callbacks userdata to extra data
-        // We can use global callbacks userdata since we don't allow C modules in Luau
-        #[cfg(feature = "luau")]
-        {
-            (*ffi::lua_callbacks(main_state)).userdata = extra.get() as *mut c_void;
-        }
-
         let inner = Arc::new(LuaInner {
             state: AtomicPtr::new(state),
             main_state,
@@ -1098,9 +1079,9 @@ impl Lua {
     /// Returns the amount of memory (in bytes) currently used inside this Lua state.
     pub fn used_memory(&self) -> usize {
         unsafe {
-            match (*self.extra.get()).mem_state.map(|x| x.as_ref()) {
-                Some(mem_state) => mem_state.used_memory(),
-                None => {
+            match MemoryState::get(self.main_state) {
+                mem_state if !mem_state.is_null() => (*mem_state).used_memory(),
+                _ => {
                     // Get data from the Lua GC
                     let used_kbytes = ffi::lua_gc(self.main_state, ffi::LUA_GCCOUNT, 0);
                     let used_kbytes_rem = ffi::lua_gc(self.main_state, ffi::LUA_GCCOUNTB, 0);
@@ -1119,9 +1100,9 @@ impl Lua {
     /// Does not work in module mode where Lua state is managed externally.
     pub fn set_memory_limit(&self, limit: usize) -> Result<usize> {
         unsafe {
-            match (*self.extra.get()).mem_state.map(|mut x| x.as_mut()) {
-                Some(mem_state) => Ok(mem_state.set_memory_limit(limit)),
-                None => Err(Error::MemoryLimitNotAvailable),
+            match MemoryState::get(self.main_state) {
+                mem_state if !mem_state.is_null() => Ok((*mem_state).set_memory_limit(limit)),
+                _ => Err(Error::MemoryLimitNotAvailable),
             }
         }
     }
@@ -3169,9 +3150,9 @@ impl Lua {
     #[inline]
     pub(crate) unsafe fn unlikely_memory_error(&self) -> bool {
         // MemoryInfo is empty in module mode so we cannot predict memory limits
-        (*self.extra.get())
-            .mem_state
-            .map(|x| x.as_ref().memory_limit() == 0)
+        MemoryState::get(self.main_state)
+            .as_ref()
+            .map(|x| x.memory_limit() == 0)
             .unwrap_or_else(|| {
                 // Alternatively, check the special flag (only for module mode)
                 #[cfg(feature = "module")]
@@ -3223,14 +3204,6 @@ impl LuaInner {
     }
 }
 
-impl ExtraData {
-    #[cfg(feature = "luau")]
-    #[inline]
-    pub(crate) fn mem_state(&self) -> NonNull<MemoryState> {
-        self.mem_state.unwrap()
-    }
-}
-
 struct StateGuard<'a>(&'a LuaInner, *mut ffi::lua_State);
 
 impl<'a> StateGuard<'a> {
@@ -3246,13 +3219,13 @@ impl<'a> Drop for StateGuard<'a> {
     }
 }
 
-#[cfg(feature = "luau")]
 unsafe fn extra_data(state: *mut ffi::lua_State) -> *mut ExtraData {
-    (*ffi::lua_callbacks(state)).userdata as *mut ExtraData
-}
+    #[cfg(feature = "luau")]
+    if cfg!(not(feature = "module")) {
+        // In the main app we can use `lua_callbacks` to access ExtraData
+        return (*ffi::lua_callbacks(state)).userdata as *mut _;
+    }
 
-#[cfg(not(feature = "luau"))]
-unsafe fn extra_data(state: *mut ffi::lua_State) -> *mut ExtraData {
     let extra_key = &EXTRA_REGISTRY_KEY as *const u8 as *const c_void;
     if ffi::lua_rawgetp(state, ffi::LUA_REGISTRYINDEX, extra_key) != ffi::LUA_TUSERDATA {
         // `ExtraData` can be null only when Lua state is foreign.
@@ -3265,6 +3238,23 @@ unsafe fn extra_data(state: *mut ffi::lua_State) -> *mut ExtraData {
     (*extra_ptr).get()
 }
 
+unsafe fn set_extra_data(
+    state: *mut ffi::lua_State,
+    extra: &Arc<UnsafeCell<ExtraData>>,
+) -> Result<()> {
+    #[cfg(feature = "luau")]
+    if cfg!(not(feature = "module")) {
+        (*ffi::lua_callbacks(state)).userdata = extra.get() as *mut _;
+        return Ok(());
+    }
+
+    push_gc_userdata(state, Arc::clone(extra), true)?;
+    protect_lua!(state, 1, 0, fn(state) {
+        let extra_key = &EXTRA_REGISTRY_KEY as *const u8 as *const c_void;
+        ffi::lua_rawsetp(state, ffi::LUA_REGISTRYINDEX, extra_key);
+    })
+}
+
 // Creates required entries in the metatable cache (see `util::METATABLE_CACHE`)
 pub(crate) fn init_metatable_cache(cache: &mut FxHashMap<TypeId, u8>) {
     cache.insert(TypeId::of::<Arc<UnsafeCell<ExtraData>>>(), 0);