Split single lua module to multiple submodules under state

src/chunk.rs

@@ -7,7 +7,7 @@ use std::string::String as StdString;
 
 use crate::error::{Error, ErrorContext, Result};
 use crate::function::Function;
-use crate::lua::{Lua, WeakLua};
+use crate::state::{Lua, WeakLua};
 use crate::table::Table;
 use crate::value::{FromLuaMulti, IntoLua, IntoLuaMulti};
 

src/conversion.rs

@@ -12,7 +12,7 @@ use num_traits::cast;
 
 use crate::error::{Error, Result};
 use crate::function::Function;
-use crate::lua::{Lua, LuaInner};
+use crate::state::{Lua, RawLua};
 use crate::string::String;
 use crate::table::Table;
 use crate::thread::Thread;
@@ -34,7 +34,7 @@ impl IntoLua for &Value {
     }
 
     #[inline]
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         lua.push_value(self)
     }
 }
@@ -60,7 +60,7 @@ impl IntoLua for &String {
     }
 
     #[inline]
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         lua.push_ref(&self.0);
         Ok(())
     }
@@ -93,7 +93,7 @@ impl IntoLua for &Table {
     }
 
     #[inline]
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         lua.push_ref(&self.0);
         Ok(())
     }
@@ -127,7 +127,7 @@ impl IntoLua for &Function {
     }
 
     #[inline]
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         lua.push_ref(&self.0);
         Ok(())
     }
@@ -161,7 +161,7 @@ impl IntoLua for &Thread {
     }
 
     #[inline]
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         lua.push_ref(&self.0);
         Ok(())
     }
@@ -195,7 +195,7 @@ impl IntoLua for &AnyUserData {
     }
 
     #[inline]
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         lua.push_ref(&self.0);
         Ok(())
     }
@@ -250,7 +250,7 @@ impl IntoLua for RegistryKey {
     }
 
     #[inline]
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         <&RegistryKey>::push_into_stack(&self, lua)
     }
 }
@@ -261,7 +261,7 @@ impl IntoLua for &RegistryKey {
         lua.registry_value(self)
     }
 
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         if !lua.owns_registry_value(self) {
             return Err(Error::MismatchedRegistryKey);
         }
@@ -290,7 +290,7 @@ impl IntoLua for bool {
     }
 
     #[inline]
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         ffi::lua_pushboolean(lua.state(), self as c_int);
         Ok(())
     }
@@ -307,7 +307,7 @@ impl FromLua for bool {
     }
 
     #[inline]
-    unsafe fn from_stack(idx: c_int, lua: &LuaInner) -> Result<Self> {
+    unsafe fn from_stack(idx: c_int, lua: &RawLua) -> Result<Self> {
         Ok(ffi::lua_toboolean(lua.state(), idx) != 0)
     }
 }
@@ -363,7 +363,7 @@ impl IntoLua for StdString {
     }
 
     #[inline]
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         push_bytes_into_stack(self, lua)
     }
 }
@@ -384,7 +384,7 @@ impl FromLua for StdString {
     }
 
     #[inline]
-    unsafe fn from_stack(idx: c_int, lua: &LuaInner) -> Result<Self> {
+    unsafe fn from_stack(idx: c_int, lua: &RawLua) -> Result<Self> {
         let state = lua.state();
         if ffi::lua_type(state, idx) == ffi::LUA_TSTRING {
             let mut size = 0;
@@ -410,7 +410,7 @@ impl IntoLua for &str {
     }
 
     #[inline]
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         push_bytes_into_stack(self, lua)
     }
 }
@@ -522,7 +522,7 @@ impl FromLua for BString {
         }
     }
 
-    unsafe fn from_stack(idx: c_int, lua: &LuaInner) -> Result<Self> {
+    unsafe fn from_stack(idx: c_int, lua: &RawLua) -> Result<Self> {
         let state = lua.state();
         match ffi::lua_type(state, idx) {
             ffi::LUA_TSTRING => {
@@ -553,7 +553,7 @@ impl IntoLua for &BStr {
 }
 
 #[inline]
-unsafe fn push_bytes_into_stack<T>(this: T, lua: &LuaInner) -> Result<()>
+unsafe fn push_bytes_into_stack<T>(this: T, lua: &RawLua) -> Result<()>
 where
     T: IntoLua + AsRef<[u8]>,
 {
@@ -584,7 +584,7 @@ macro_rules! lua_convert_int {
             }
 
             #[inline]
-            unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+            unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
                 match cast(self) {
                     Some(i) => ffi::lua_pushinteger(lua.state(), i),
                     None => ffi::lua_pushnumber(lua.state(), self as ffi::lua_Number),
@@ -881,7 +881,7 @@ impl<T: IntoLua> IntoLua for Option<T> {
     }
 
     #[inline]
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         match self {
             Some(val) => val.push_into_stack(lua)?,
             None => ffi::lua_pushnil(lua.state()),
@@ -900,7 +900,7 @@ impl<T: FromLua> FromLua for Option<T> {
     }
 
     #[inline]
-    unsafe fn from_stack(idx: c_int, lua: &LuaInner) -> Result<Self> {
+    unsafe fn from_stack(idx: c_int, lua: &RawLua) -> Result<Self> {
         if ffi::lua_isnil(lua.state(), idx) != 0 {
             Ok(None)
         } else {

src/function.rs

@@ -5,7 +5,7 @@ use std::ptr;
 use std::slice;
 
 use crate::error::{Error, Result};
-use crate::lua::Lua;
+use crate::state::Lua;
 use crate::table::Table;
 use crate::types::{Callback, MaybeSend, ValueRef};
 use crate::util::{
@@ -161,11 +161,13 @@ impl Function {
         R: FromLuaMulti,
     {
         let lua = self.0.lua.lock();
-        let thread_res = lua.create_recycled_thread(self).map(|th| {
-            let mut th = th.into_async(args);
-            th.set_recyclable(true);
-            th
-        });
+        let thread_res = unsafe {
+            lua.create_recycled_thread(self).map(|th| {
+                let mut th = th.into_async(args);
+                th.set_recyclable(true);
+                th
+            })
+        };
         async move { thread_res?.await }
     }
 

src/hook.rs

@@ -1,6 +1,6 @@
 use std::borrow::Cow;
 use std::cell::UnsafeCell;
-use std::mem::ManuallyDrop;
+use std::ops::Deref;
 #[cfg(not(feature = "luau"))]
 use std::ops::{BitOr, BitOrAssign};
 use std::os::raw::c_int;
@@ -8,7 +8,7 @@ use std::os::raw::c_int;
 use ffi::lua_Debug;
 use parking_lot::ReentrantMutexGuard;
 
-use crate::lua::{Lua, LuaInner};
+use crate::state::RawLua;
 use crate::util::{linenumber_to_usize, ptr_to_lossy_str, ptr_to_str};
 
 /// Contains information about currently executing Lua code.
@@ -20,38 +20,46 @@ use crate::util::{linenumber_to_usize, ptr_to_lossy_str, ptr_to_str};
 ///
 /// [lua_doc]: https://www.lua.org/manual/5.4/manual.html#lua_Debug
 /// [`Lua::set_hook`]: crate::Lua::set_hook
-pub struct Debug<'lua> {
-    lua: ManuallyDrop<ReentrantMutexGuard<'lua, LuaInner>>,
+pub struct Debug<'a> {
+    lua: EitherLua<'a>,
     ar: ActivationRecord,
     #[cfg(feature = "luau")]
     level: c_int,
 }
 
-impl<'lua> Drop for Debug<'lua> {
-    fn drop(&mut self) {
-        if let ActivationRecord::Owned(_) = self.ar {
-            unsafe { ManuallyDrop::drop(&mut self.lua) }
+enum EitherLua<'a> {
+    Owned(ReentrantMutexGuard<'a, RawLua>),
+    Borrowed(&'a RawLua),
+}
+
+impl Deref for EitherLua<'_> {
+    type Target = RawLua;
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            EitherLua::Owned(guard) => &*guard,
+            EitherLua::Borrowed(lua) => lua,
         }
     }
 }
 
-impl<'lua> Debug<'lua> {
+impl<'a> Debug<'a> {
     // We assume the lock is held when this function is called.
     #[cfg(not(feature = "luau"))]
-    pub(crate) fn new(lua: &'lua Lua, ar: *mut lua_Debug) -> Self {
+    pub(crate) fn new(lua: &'a RawLua, ar: *mut lua_Debug) -> Self {
         Debug {
-            lua: unsafe { lua.guard_unchecked() },
+            lua: EitherLua::Borrowed(lua),
             ar: ActivationRecord::Borrowed(ar),
         }
     }
 
     pub(crate) fn new_owned(
-        guard: ReentrantMutexGuard<'lua, LuaInner>,
+        guard: ReentrantMutexGuard<'a, RawLua>,
         _level: c_int,
         ar: lua_Debug,
     ) -> Self {
         Debug {
-            lua: ManuallyDrop::new(guard),
+            lua: EitherLua::Owned(guard),
             ar: ActivationRecord::Owned(UnsafeCell::new(ar)),
             #[cfg(feature = "luau")]
             level: _level,

src/lib.rs

@@ -84,12 +84,12 @@ mod conversion;
 mod error;
 mod function;
 mod hook;
-mod lua;
 #[cfg(feature = "luau")]
 mod luau;
 mod memory;
 mod multi;
 // mod scope;
+mod state;
 mod stdlib;
 mod string;
 mod table;
@@ -107,7 +107,7 @@ pub use crate::chunk::{AsChunk, Chunk, ChunkMode};
 pub use crate::error::{Error, ErrorContext, ExternalError, ExternalResult, Result};
 pub use crate::function::{Function, FunctionInfo};
 pub use crate::hook::{Debug, DebugEvent, DebugNames, DebugSource, DebugStack};
-pub use crate::lua::{GCMode, Lua, LuaOptions};
+pub use crate::state::{GCMode, Lua, LuaOptions};
 pub use crate::multi::Variadic;
 // pub use crate::scope::Scope;
 pub use crate::stdlib::StdLib;

src/luau/mod.rs

@@ -2,7 +2,7 @@ use std::ffi::CStr;
 use std::os::raw::{c_float, c_int};
 
 use crate::error::Result;
-use crate::lua::Lua;
+use crate::state::Lua;
 
 // Since Luau has some missing standard functions, we re-implement them here
 

src/luau/package.rs

@@ -7,7 +7,7 @@ use std::{env, fs};
 
 use crate::chunk::ChunkMode;
 use crate::error::Result;
-use crate::lua::Lua;
+use crate::state::Lua;
 use crate::table::Table;
 use crate::types::RegistryKey;
 use crate::value::{IntoLua, Value};

src/multi.rs

@@ -5,7 +5,8 @@ use std::os::raw::c_int;
 use std::result::Result as StdResult;
 
 use crate::error::Result;
-use crate::lua::{Lua, LuaInner};
+use crate::state::Lua;
+use crate::state::RawLua;
 use crate::util::check_stack;
 use crate::value::{FromLua, FromLuaMulti, IntoLua, IntoLuaMulti, MultiValue, Nil};
 
@@ -21,7 +22,7 @@ impl<T: IntoLua, E: IntoLua> IntoLuaMulti for StdResult<T, E> {
     }
 
     #[inline]
-    unsafe fn push_into_stack_multi(self, lua: &LuaInner) -> Result<c_int> {
+    unsafe fn push_into_stack_multi(self, lua: &RawLua) -> Result<c_int> {
         match self {
             Ok(val) => (val,).push_into_stack_multi(lua),
             Err(err) => (Nil, err).push_into_stack_multi(lua),
@@ -39,7 +40,7 @@ impl<E: IntoLua> IntoLuaMulti for StdResult<(), E> {
     }
 
     #[inline]
-    unsafe fn push_into_stack_multi(self, lua: &LuaInner) -> Result<c_int> {
+    unsafe fn push_into_stack_multi(self, lua: &RawLua) -> Result<c_int> {
         match self {
             Ok(_) => Ok(0),
             Err(err) => (Nil, err).push_into_stack_multi(lua),
@@ -56,7 +57,7 @@ impl<T: IntoLua> IntoLuaMulti for T {
     }
 
     #[inline]
-    unsafe fn push_into_stack_multi(self, lua: &LuaInner) -> Result<c_int> {
+    unsafe fn push_into_stack_multi(self, lua: &RawLua) -> Result<c_int> {
         self.push_into_stack(lua)?;
         Ok(1)
     }
@@ -74,7 +75,7 @@ impl<T: FromLua> FromLuaMulti for T {
     }
 
     #[inline]
-    unsafe fn from_stack_multi(nvals: c_int, lua: &LuaInner) -> Result<Self> {
+    unsafe fn from_stack_multi(nvals: c_int, lua: &RawLua) -> Result<Self> {
         if nvals == 0 {
             return T::from_lua(Nil, lua.lua());
         }
@@ -86,7 +87,7 @@ impl<T: FromLua> FromLuaMulti for T {
         nargs: c_int,
         i: usize,
         to: Option<&str>,
-        lua: &LuaInner,
+        lua: &RawLua,
     ) -> Result<Self> {
         if nargs == 0 {
             return T::from_lua_arg(Nil, i, to, lua.lua());
@@ -209,7 +210,7 @@ macro_rules! impl_tuple {
             }
 
             #[inline]
-            unsafe fn push_into_stack_multi(self, _lua: &LuaInner) -> Result<c_int> {
+            unsafe fn push_into_stack_multi(self, _lua: &RawLua) -> Result<c_int> {
                 Ok(0)
             }
         }
@@ -221,7 +222,7 @@ macro_rules! impl_tuple {
             }
 
             #[inline]
-            unsafe fn from_stack_multi(nvals: c_int, lua: &LuaInner) -> Result<Self> {
+            unsafe fn from_stack_multi(nvals: c_int, lua: &RawLua) -> Result<Self> {
                 if nvals > 0 {
                     ffi::lua_pop(lua.state(), nvals);
                 }
@@ -247,7 +248,7 @@ macro_rules! impl_tuple {
 
             #[allow(non_snake_case)]
             #[inline]
-            unsafe fn push_into_stack_multi(self, lua: &LuaInner) -> Result<c_int> {
+            unsafe fn push_into_stack_multi(self, lua: &RawLua) -> Result<c_int> {
                 let ($($name,)* $last,) = self;
                 let mut nresults = 0;
                 $(
@@ -288,7 +289,7 @@ macro_rules! impl_tuple {
 
             #[allow(unused_mut, non_snake_case)]
             #[inline]
-            unsafe fn from_stack_multi(mut nvals: c_int, lua: &LuaInner) -> Result<Self> {
+            unsafe fn from_stack_multi(mut nvals: c_int, lua: &RawLua) -> Result<Self> {
                 $(
                     let $name = if nvals > 0 {
                         nvals -= 1;
@@ -303,7 +304,7 @@ macro_rules! impl_tuple {
 
             #[allow(unused_mut, non_snake_case)]
             #[inline]
-            unsafe fn from_stack_args(mut nargs: c_int, mut i: usize, to: Option<&str>, lua: &LuaInner) -> Result<Self> {
+            unsafe fn from_stack_args(mut nargs: c_int, mut i: usize, to: Option<&str>, lua: &RawLua) -> Result<Self> {
                 $(
                     let $name = if nargs > 0 {
                         nargs -= 1;

src/serde/mod.rs

@@ -5,7 +5,7 @@ use std::os::raw::c_void;
 use serde::{de::DeserializeOwned, ser::Serialize};
 
 use crate::error::Result;
-use crate::lua::Lua;
+use crate::state::Lua;
 use crate::private::Sealed;
 use crate::table::Table;
 use crate::util::check_stack;

src/serde/ser.rs

@@ -2,7 +2,7 @@ use serde::{ser, Serialize};
 
 use super::LuaSerdeExt;
 use crate::error::{Error, Result};
-use crate::lua::Lua;
+use crate::state::Lua;
 use crate::table::Table;
 use crate::value::{IntoLua, Value};
 

src/state/extra.rs

@@ -0,0 +1,236 @@
+use std::any::TypeId;
+use std::cell::UnsafeCell;
+// use std::collections::VecDeque;
+use std::mem::{self, MaybeUninit};
+use std::os::raw::{c_int, c_void};
+use std::ptr;
+use std::sync::{Arc, Weak};
+
+use parking_lot::{Mutex, ReentrantMutex};
+use rustc_hash::FxHashMap;
+
+use crate::error::Result;
+use crate::state::RawLua;
+use crate::stdlib::StdLib;
+use crate::types::AppData;
+use crate::util::{get_gc_metatable, push_gc_userdata, WrappedFailure};
+
+#[cfg(any(feature = "luau", doc))]
+use crate::chunk::Compiler;
+
+#[cfg(feature = "async")]
+use {futures_util::task::noop_waker_ref, std::ptr::NonNull, std::task::Waker};
+
+use super::{Lua, WeakLua};
+
+// Unique key to store `ExtraData` in the registry
+static EXTRA_REGISTRY_KEY: u8 = 0;
+
+const WRAPPED_FAILURE_POOL_SIZE: usize = 64;
+// const MULTIVALUE_POOL_SIZE: usize = 64;
+const REF_STACK_RESERVE: c_int = 1;
+
+/// Data associated with the Lua state.
+pub(crate) struct ExtraData {
+    // Same layout as `Lua`
+    pub(super) lua: MaybeUninit<Arc<ReentrantMutex<RawLua>>>,
+    // Same layout as `WeakLua`
+    pub(super) weak: MaybeUninit<Weak<ReentrantMutex<RawLua>>>,
+
+    pub(super) registered_userdata: FxHashMap<TypeId, c_int>,
+    pub(super) registered_userdata_mt: FxHashMap<*const c_void, Option<TypeId>>,
+    pub(super) last_checked_userdata_mt: (*const c_void, Option<TypeId>),
+
+    // When Lua instance dropped, setting `None` would prevent collecting `RegistryKey`s
+    pub(super) registry_unref_list: Arc<Mutex<Option<Vec<c_int>>>>,
+
+    // Container to store arbitrary data (extensions)
+    pub(super) app_data: AppData,
+
+    pub(super) safe: bool,
+    pub(super) libs: StdLib,
+    #[cfg(feature = "module")]
+    pub(super) skip_memory_check: bool,
+
+    // Auxiliary thread to store references
+    pub(super) ref_thread: *mut ffi::lua_State,
+    pub(super) ref_stack_size: c_int,
+    pub(super) ref_stack_top: c_int,
+    pub(super) ref_free: Vec<c_int>,
+
+    // Pool of `WrappedFailure` enums in the ref thread (as userdata)
+    pub(super) wrapped_failure_pool: Vec<c_int>,
+    // Pool of `MultiValue` containers
+    // multivalue_pool: Vec<VecDeque<Value>>,
+    // Pool of `Thread`s (coroutines) for async execution
+    #[cfg(feature = "async")]
+    pub(super) thread_pool: Vec<c_int>,
+
+    // Address of `WrappedFailure` metatable
+    pub(super) wrapped_failure_mt_ptr: *const c_void,
+
+    // Waker for polling futures
+    #[cfg(feature = "async")]
+    pub(super) waker: NonNull<Waker>,
+
+    #[cfg(not(feature = "luau"))]
+    pub(super) hook_callback: Option<crate::types::HookCallback>,
+    #[cfg(not(feature = "luau"))]
+    pub(super) hook_thread: *mut ffi::lua_State,
+    #[cfg(feature = "lua54")]
+    pub(super) warn_callback: Option<crate::types::WarnCallback>,
+    #[cfg(feature = "luau")]
+    pub(super) interrupt_callback: Option<crate::types::InterruptCallback>,
+
+    #[cfg(feature = "luau")]
+    pub(super) sandboxed: bool,
+    #[cfg(feature = "luau")]
+    pub(super) compiler: Option<Compiler>,
+    #[cfg(feature = "luau-jit")]
+    pub(super) enable_jit: bool,
+}
+
+impl Drop for ExtraData {
+    fn drop(&mut self) {
+        #[cfg(feature = "module")]
+        unsafe {
+            self.inner.assume_init_drop();
+        }
+        unsafe { self.weak.assume_init_drop() };
+        *self.registry_unref_list.lock() = None;
+    }
+}
+
+impl ExtraData {
+    // Index of `error_traceback` function in auxiliary thread stack
+    #[cfg(any(feature = "lua51", feature = "luajit", feature = "luau"))]
+    pub(super) const ERROR_TRACEBACK_IDX: c_int = 1;
+
+    pub(super) unsafe fn init(state: *mut ffi::lua_State) -> Arc<UnsafeCell<Self>> {
+        // Create ref stack thread and place it in the registry to prevent it
+        // from being garbage collected.
+        let ref_thread = mlua_expect!(
+            protect_lua!(state, 0, 0, |state| {
+                let thread = ffi::lua_newthread(state);
+                ffi::luaL_ref(state, ffi::LUA_REGISTRYINDEX);
+                thread
+            }),
+            "Error while creating ref thread",
+        );
+
+        let wrapped_failure_mt_ptr = {
+            get_gc_metatable::<WrappedFailure>(state);
+            let ptr = ffi::lua_topointer(state, -1);
+            ffi::lua_pop(state, 1);
+            ptr
+        };
+
+        // Store `error_traceback` function on the ref stack
+        #[cfg(any(feature = "lua51", feature = "luajit", feature = "luau"))]
+        {
+            ffi::lua_pushcfunction(ref_thread, crate::util::error_traceback);
+            assert_eq!(ffi::lua_gettop(ref_thread), Self::ERROR_TRACEBACK_IDX);
+        }
+
+        let extra = Arc::new(UnsafeCell::new(ExtraData {
+            lua: MaybeUninit::uninit(),
+            weak: MaybeUninit::uninit(),
+            registered_userdata: FxHashMap::default(),
+            registered_userdata_mt: FxHashMap::default(),
+            last_checked_userdata_mt: (ptr::null(), None),
+            registry_unref_list: Arc::new(Mutex::new(Some(Vec::new()))),
+            app_data: AppData::default(),
+            safe: false,
+            libs: StdLib::NONE,
+            #[cfg(feature = "module")]
+            skip_memory_check: false,
+            ref_thread,
+            // We need some reserved stack space to move values in and out of the ref stack.
+            ref_stack_size: ffi::LUA_MINSTACK - REF_STACK_RESERVE,
+            ref_stack_top: ffi::lua_gettop(ref_thread),
+            ref_free: Vec::new(),
+            wrapped_failure_pool: Vec::with_capacity(WRAPPED_FAILURE_POOL_SIZE),
+            // multivalue_pool: Vec::with_capacity(MULTIVALUE_POOL_SIZE),
+            #[cfg(feature = "async")]
+            thread_pool: Vec::new(),
+            wrapped_failure_mt_ptr,
+            #[cfg(feature = "async")]
+            waker: NonNull::from(noop_waker_ref()),
+            #[cfg(not(feature = "luau"))]
+            hook_callback: None,
+            #[cfg(not(feature = "luau"))]
+            hook_thread: ptr::null_mut(),
+            #[cfg(feature = "lua54")]
+            warn_callback: None,
+            #[cfg(feature = "luau")]
+            interrupt_callback: None,
+            #[cfg(feature = "luau")]
+            sandboxed: false,
+            #[cfg(feature = "luau")]
+            compiler: None,
+            #[cfg(feature = "luau-jit")]
+            enable_jit: true,
+        }));
+
+        // Store it in the registry
+        mlua_expect!(Self::store(&extra, state), "Error while storing extra data");
+
+        extra
+    }
+
+    pub(super) unsafe fn set_lua(&mut self, lua: &Arc<ReentrantMutex<RawLua>>) {
+        self.lua.write(Arc::clone(lua));
+        if cfg!(not(feature = "module")) {
+            Arc::decrement_strong_count(Arc::as_ptr(lua));
+        }
+        self.weak.write(Arc::downgrade(lua));
+    }
+
+    pub(super) unsafe fn get(state: *mut ffi::lua_State) -> *mut Self {
+        #[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 _;
+        }
+
+        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.
+            // This case in used in `Lua::try_from_ptr()`.
+            ffi::lua_pop(state, 1);
+            return ptr::null_mut();
+        }
+        let extra_ptr = ffi::lua_touserdata(state, -1) as *mut Arc<UnsafeCell<ExtraData>>;
+        ffi::lua_pop(state, 1);
+        (*extra_ptr).get()
+    }
+
+    unsafe fn store(extra: &Arc<UnsafeCell<Self>>, state: *mut ffi::lua_State) -> 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);
+        })
+    }
+
+    #[inline(always)]
+    pub(super) unsafe fn lua(&self) -> &Lua {
+        mem::transmute(self.lua.assume_init_ref())
+    }
+
+    #[inline(always)]
+    pub(super) unsafe fn raw_lua(&self) -> &RawLua {
+        &*self.lua.assume_init_ref().data_ptr()
+    }
+
+    #[inline(always)]
+    pub(super) unsafe fn weak(&self) -> &WeakLua {
+        mem::transmute(self.weak.assume_init_ref())
+    }
+}

src/state/util.rs

@@ -0,0 +1,187 @@
+use std::os::raw::c_int;
+use std::panic::{catch_unwind, AssertUnwindSafe};
+use std::ptr;
+use std::sync::Arc;
+
+use crate::error::{Error, Result};
+use crate::state::{ExtraData, RawLua};
+use crate::util::{self, get_gc_metatable, WrappedFailure};
+
+const WRAPPED_FAILURE_POOL_SIZE: usize = 64;
+// const MULTIVALUE_POOL_SIZE: usize = 64;
+
+pub(super) struct StateGuard<'a>(&'a RawLua, *mut ffi::lua_State);
+
+impl<'a> StateGuard<'a> {
+    pub(super) fn new(inner: &'a RawLua, mut state: *mut ffi::lua_State) -> Self {
+        state = inner.state.replace(state);
+        Self(inner, state)
+    }
+}
+
+impl<'a> Drop for StateGuard<'a> {
+    fn drop(&mut self) {
+        self.0.state.set(self.1);
+    }
+}
+
+// An optimized version of `callback_error` that does not allocate `WrappedFailure` userdata
+// and instead reuses unsed values from previous calls (or allocates new).
+pub(super) unsafe fn callback_error_ext<F, R>(
+    state: *mut ffi::lua_State,
+    mut extra: *mut ExtraData,
+    f: F,
+) -> R
+where
+    F: FnOnce(c_int) -> Result<R>,
+{
+    if extra.is_null() {
+        extra = ExtraData::get(state);
+    }
+
+    let nargs = ffi::lua_gettop(state);
+
+    enum PreallocatedFailure {
+        New(*mut WrappedFailure),
+        Existing(i32),
+    }
+
+    impl PreallocatedFailure {
+        unsafe fn reserve(state: *mut ffi::lua_State, extra: *mut ExtraData) -> Self {
+            match (*extra).wrapped_failure_pool.pop() {
+                Some(index) => PreallocatedFailure::Existing(index),
+                None => {
+                    // We need to check stack for Luau in case when callback is called from interrupt
+                    // See https://github.com/Roblox/luau/issues/446 and mlua #142 and #153
+                    #[cfg(feature = "luau")]
+                    ffi::lua_rawcheckstack(state, 2);
+                    // Place it to the beginning of the stack
+                    let ud = WrappedFailure::new_userdata(state);
+                    ffi::lua_insert(state, 1);
+                    PreallocatedFailure::New(ud)
+                }
+            }
+        }
+
+        unsafe fn r#use(
+            &self,
+            state: *mut ffi::lua_State,
+            extra: *mut ExtraData,
+        ) -> *mut WrappedFailure {
+            let ref_thread = (*extra).ref_thread;
+            match *self {
+                PreallocatedFailure::New(ud) => {
+                    ffi::lua_settop(state, 1);
+                    ud
+                }
+                PreallocatedFailure::Existing(index) => {
+                    ffi::lua_settop(state, 0);
+                    #[cfg(feature = "luau")]
+                    ffi::lua_rawcheckstack(state, 2);
+                    ffi::lua_pushvalue(ref_thread, index);
+                    ffi::lua_xmove(ref_thread, state, 1);
+                    ffi::lua_pushnil(ref_thread);
+                    ffi::lua_replace(ref_thread, index);
+                    (*extra).ref_free.push(index);
+                    ffi::lua_touserdata(state, -1) as *mut WrappedFailure
+                }
+            }
+        }
+
+        unsafe fn release(self, state: *mut ffi::lua_State, extra: *mut ExtraData) {
+            let ref_thread = (*extra).ref_thread;
+            match self {
+                PreallocatedFailure::New(_) => {
+                    if (*extra).wrapped_failure_pool.len() < WRAPPED_FAILURE_POOL_SIZE {
+                        ffi::lua_rotate(state, 1, -1);
+                        ffi::lua_xmove(state, ref_thread, 1);
+                        let index = ref_stack_pop(extra);
+                        (*extra).wrapped_failure_pool.push(index);
+                    } else {
+                        ffi::lua_remove(state, 1);
+                    }
+                }
+                PreallocatedFailure::Existing(index) => {
+                    if (*extra).wrapped_failure_pool.len() < WRAPPED_FAILURE_POOL_SIZE {
+                        (*extra).wrapped_failure_pool.push(index);
+                    } else {
+                        ffi::lua_pushnil(ref_thread);
+                        ffi::lua_replace(ref_thread, index);
+                        (*extra).ref_free.push(index);
+                    }
+                }
+            }
+        }
+    }
+
+    // We cannot shadow Rust errors with Lua ones, so we need to reserve pre-allocated memory
+    // to store a wrapped failure (error or panic) *before* we proceed.
+    let prealloc_failure = PreallocatedFailure::reserve(state, extra);
+
+    match catch_unwind(AssertUnwindSafe(|| f(nargs))) {
+        Ok(Ok(r)) => {
+            // Return unused `WrappedFailure` to the pool
+            prealloc_failure.release(state, extra);
+            r
+        }
+        Ok(Err(err)) => {
+            let wrapped_error = prealloc_failure.r#use(state, extra);
+
+            // Build `CallbackError` with traceback
+            let traceback = if ffi::lua_checkstack(state, ffi::LUA_TRACEBACK_STACK) != 0 {
+                ffi::luaL_traceback(state, state, ptr::null(), 0);
+                let traceback = util::to_string(state, -1);
+                ffi::lua_pop(state, 1);
+                traceback
+            } else {
+                "<not enough stack space for traceback>".to_string()
+            };
+            let cause = Arc::new(err);
+            ptr::write(
+                wrapped_error,
+                WrappedFailure::Error(Error::CallbackError { traceback, cause }),
+            );
+            get_gc_metatable::<WrappedFailure>(state);
+            ffi::lua_setmetatable(state, -2);
+
+            ffi::lua_error(state)
+        }
+        Err(p) => {
+            let wrapped_panic = prealloc_failure.r#use(state, extra);
+            ptr::write(wrapped_panic, WrappedFailure::Panic(Some(p)));
+            get_gc_metatable::<WrappedFailure>(state);
+            ffi::lua_setmetatable(state, -2);
+            ffi::lua_error(state)
+        }
+    }
+}
+
+pub(super) unsafe fn ref_stack_pop(extra: *mut ExtraData) -> c_int {
+    let extra = &mut *extra;
+    if let Some(free) = extra.ref_free.pop() {
+        ffi::lua_replace(extra.ref_thread, free);
+        return free;
+    }
+
+    // Try to grow max stack size
+    if extra.ref_stack_top >= extra.ref_stack_size {
+        let mut inc = extra.ref_stack_size; // Try to double stack size
+        while inc > 0 && ffi::lua_checkstack(extra.ref_thread, inc) == 0 {
+            inc /= 2;
+        }
+        if inc == 0 {
+            // Pop item on top of the stack to avoid stack leaking and successfully run destructors
+            // during unwinding.
+            ffi::lua_pop(extra.ref_thread, 1);
+            let top = extra.ref_stack_top;
+            // It is a user error to create enough references to exhaust the Lua max stack size for
+            // the ref thread.
+            panic!(
+                "cannot create a Lua reference, out of auxiliary stack space (used {top} slots)"
+            );
+        }
+        extra.ref_stack_size += inc;
+    }
+    extra.ref_stack_top += 1;
+    extra.ref_stack_top
+}

src/thread.rs

@@ -2,8 +2,8 @@ use std::os::raw::{c_int, c_void};
 
 use crate::error::{Error, Result};
 #[allow(unused)]
-use crate::lua::Lua;
-use crate::lua::LuaInner;
+use crate::state::Lua;
+use crate::state::RawLua;
 use crate::types::ValueRef;
 use crate::util::{check_stack, error_traceback_thread, pop_error, StackGuard};
 use crate::value::{FromLuaMulti, IntoLuaMulti};
@@ -64,11 +64,6 @@ pub struct AsyncThread<R> {
 
 impl Thread {
     #[inline(always)]
-    pub(crate) fn new(lua: &LuaInner, r#ref: ValueRef) -> Self {
-        let state = unsafe { ffi::lua_tothread(lua.ref_thread(), r#ref.index) };
-        Thread(r#ref, state)
-    }
-
     const fn state(&self) -> *mut ffi::lua_State {
         self.1
     }
@@ -502,7 +497,7 @@ unsafe fn is_poll_pending(state: *mut ffi::lua_State) -> bool {
 
 #[cfg(feature = "async")]
 struct WakerGuard<'lua, 'a> {
-    lua: &'lua LuaInner,
+    lua: &'lua RawLua,
     prev: NonNull<Waker>,
     _phantom: PhantomData<&'a ()>,
 }
@@ -510,7 +505,7 @@ struct WakerGuard<'lua, 'a> {
 #[cfg(feature = "async")]
 impl<'lua, 'a> WakerGuard<'lua, 'a> {
     #[inline]
-    pub fn new(lua: &'lua LuaInner, waker: &'a Waker) -> Result<WakerGuard<'lua, 'a>> {
+    pub fn new(lua: &'lua RawLua, waker: &'a Waker) -> Result<WakerGuard<'lua, 'a>> {
         let prev = unsafe { lua.set_waker(NonNull::from(waker)) };
         Ok(WakerGuard {
             lua,

src/types.rs

@@ -14,7 +14,7 @@ use rustc_hash::FxHashMap;
 use crate::error::Result;
 #[cfg(not(feature = "luau"))]
 use crate::hook::Debug;
-use crate::lua::{ExtraData, Lua, LuaGuard, LuaInner, WeakLua};
+use crate::state::{ExtraData, Lua, LuaGuard, RawLua, WeakLua};
 
 #[cfg(feature = "async")]
 use {crate::value::MultiValue, futures_util::future::LocalBoxFuture};
@@ -41,7 +41,7 @@ pub(crate) enum SubtypeId {
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub struct LightUserData(pub *mut c_void);
 
-pub(crate) type Callback<'a> = Box<dyn Fn(&'a LuaInner, c_int) -> Result<c_int> + 'static>;
+pub(crate) type Callback<'a> = Box<dyn Fn(&'a RawLua, c_int) -> Result<c_int> + 'static>;
 
 pub(crate) struct Upvalue<T> {
     pub(crate) data: T,
@@ -52,7 +52,7 @@ pub(crate) type CallbackUpvalue = Upvalue<Callback<'static>>;
 
 #[cfg(feature = "async")]
 pub(crate) type AsyncCallback<'a> =
-    Box<dyn Fn(&'a LuaInner, MultiValue) -> LocalBoxFuture<'a, Result<c_int>> + 'static>;
+    Box<dyn Fn(&'a RawLua, MultiValue) -> LocalBoxFuture<'a, Result<c_int>> + 'static>;
 
 #[cfg(feature = "async")]
 pub(crate) type AsyncCallbackUpvalue = Upvalue<AsyncCallback<'static>>;
@@ -281,7 +281,8 @@ pub(crate) struct ValueRef {
 }
 
 impl ValueRef {
-    pub(crate) fn new(lua: &LuaInner, index: c_int) -> Self {
+    #[inline]
+    pub(crate) fn new(lua: &RawLua, index: c_int) -> Self {
         ValueRef {
             lua: lua.weak().clone(),
             index,
@@ -304,7 +305,7 @@ impl fmt::Debug for ValueRef {
 
 impl Clone for ValueRef {
     fn clone(&self) -> Self {
-        self.lua.lock().clone_ref(self)
+        unsafe { self.lua.lock().clone_ref(self) }
     }
 }
 
@@ -312,7 +313,7 @@ impl Drop for ValueRef {
     fn drop(&mut self) {
         if self.drop {
             if let Some(lua) = self.lua.try_lock() {
-                lua.drop_ref(self);
+                unsafe { lua.drop_ref(self) };
             }
         }
     }

src/userdata.rs

@@ -16,7 +16,7 @@ use {
 
 use crate::error::{Error, Result};
 use crate::function::Function;
-use crate::lua::{Lua, LuaGuard};
+use crate::state::{Lua, LuaGuard};
 use crate::string::String;
 use crate::table::{Table, TablePairs};
 use crate::types::{MaybeSend, SubtypeId, ValueRef};

src/userdata/cell.rs

@@ -9,7 +9,8 @@ use std::rc::Rc;
 use serde::ser::{Serialize, Serializer};
 
 use crate::error::{Error, Result};
-use crate::lua::{Lua, LuaGuard, LuaInner};
+use crate::state::{Lua, LuaGuard};
+use crate::state::RawLua;
 use crate::userdata::AnyUserData;
 use crate::util::get_userdata;
 use crate::value::{FromLua, Value};
@@ -199,7 +200,7 @@ impl<T: 'static> FromLua for UserDataRef<T> {
         try_value_to_userdata::<T>(value)?.borrow()
     }
 
-    unsafe fn from_stack(idx: c_int, lua: &LuaInner) -> Result<Self> {
+    unsafe fn from_stack(idx: c_int, lua: &RawLua) -> Result<Self> {
         let type_id = lua.get_userdata_type_id(idx)?;
         match type_id {
             Some(type_id) if type_id == TypeId::of::<T>() => {
@@ -268,7 +269,7 @@ impl<T: 'static> FromLua for UserDataRefMut<T> {
         try_value_to_userdata::<T>(value)?.borrow_mut()
     }
 
-    unsafe fn from_stack(idx: c_int, lua: &LuaInner) -> Result<Self> {
+    unsafe fn from_stack(idx: c_int, lua: &RawLua) -> Result<Self> {
         let type_id = lua.get_userdata_type_id(idx)?;
         match type_id {
             Some(type_id) if type_id == TypeId::of::<T>() => {

src/userdata/registry.rs

@@ -7,7 +7,7 @@ use std::os::raw::c_int;
 use std::string::String as StdString;
 
 use crate::error::{Error, Result};
-use crate::lua::Lua;
+use crate::state::Lua;
 use crate::types::{Callback, MaybeSend};
 use crate::userdata::{AnyUserData, MetaMethod, UserData, UserDataFields, UserDataMethods};
 use crate::util::{get_userdata, short_type_name};

src/util/mod.rs

@@ -1,5 +1,6 @@
 use std::any::{Any, TypeId};
 use std::borrow::Cow;
+use std::cell::UnsafeCell;
 use std::ffi::CStr;
 use std::fmt::Write;
 use std::mem::MaybeUninit;
@@ -18,7 +19,19 @@ pub(crate) use short_names::short_type_name;
 
 static METATABLE_CACHE: Lazy<FxHashMap<TypeId, u8>> = Lazy::new(|| {
     let mut map = FxHashMap::with_capacity_and_hasher(32, Default::default());
-    crate::lua::init_metatable_cache(&mut map);
+
+    map.insert(TypeId::of::<Arc<UnsafeCell<crate::state::ExtraData>>>(), 0);
+    map.insert(TypeId::of::<crate::types::Callback>(), 0);
+    map.insert(TypeId::of::<crate::types::CallbackUpvalue>(), 0);
+
+    #[cfg(feature = "async")]
+    {
+        map.insert(TypeId::of::<crate::types::AsyncCallback>(), 0);
+        map.insert(TypeId::of::<crate::types::AsyncCallbackUpvalue>(), 0);
+        map.insert(TypeId::of::<crate::types::AsyncPollUpvalue>(), 0);
+        map.insert(TypeId::of::<Option<std::task::Waker>>(), 0);
+    }
+
     map.insert(TypeId::of::<WrappedFailure>(), 0);
     map.insert(TypeId::of::<String>(), 0);
     map

src/value.rs

@@ -19,7 +19,7 @@ use {
 
 use crate::error::{Error, Result};
 use crate::function::Function;
-use crate::lua::{Lua, LuaInner};
+use crate::state::{Lua, RawLua};
 use crate::string::String;
 use crate::table::Table;
 use crate::thread::Thread;
@@ -698,7 +698,7 @@ pub trait IntoLua: Sized {
     /// This method does not check Lua stack space.
     #[doc(hidden)]
     #[inline]
-    unsafe fn push_into_stack(self, lua: &LuaInner) -> Result<()> {
+    unsafe fn push_into_stack(self, lua: &RawLua) -> Result<()> {
         lua.push_value(&self.into_lua(lua.lua())?)
     }
 }
@@ -726,19 +726,14 @@ pub trait FromLua: Sized {
     /// Performs the conversion for a value in the Lua stack at index `idx`.
     #[doc(hidden)]
     #[inline]
-    unsafe fn from_stack(idx: c_int, lua: &LuaInner) -> Result<Self> {
+    unsafe fn from_stack(idx: c_int, lua: &RawLua) -> Result<Self> {
         Self::from_lua(lua.stack_value(idx), lua.lua())
     }
 
     /// Same as `from_lua_arg` but for a value in the Lua stack at index `idx`.
     #[doc(hidden)]
     #[inline]
-    unsafe fn from_stack_arg(
-        idx: c_int,
-        i: usize,
-        to: Option<&str>,
-        lua: &LuaInner,
-    ) -> Result<Self> {
+    unsafe fn from_stack_arg(idx: c_int, i: usize, to: Option<&str>, lua: &RawLua) -> Result<Self> {
         Self::from_stack(idx, lua).map_err(|err| Error::BadArgument {
             to: to.map(|s| s.to_string()),
             pos: i,
@@ -876,7 +871,7 @@ pub trait IntoLuaMulti: Sized {
     /// Returns number of pushed values.
     #[doc(hidden)]
     #[inline]
-    unsafe fn push_into_stack_multi(self, lua: &LuaInner) -> Result<c_int> {
+    unsafe fn push_into_stack_multi(self, lua: &RawLua) -> Result<c_int> {
         let values = self.into_lua_multi(lua.lua())?;
         let len: c_int = values.len().try_into().unwrap();
         unsafe {
@@ -916,7 +911,7 @@ pub trait FromLuaMulti: Sized {
     /// Performs the conversion for a number of values in the Lua stack.
     #[doc(hidden)]
     #[inline]
-    unsafe fn from_stack_multi(nvals: c_int, lua: &LuaInner) -> Result<Self> {
+    unsafe fn from_stack_multi(nvals: c_int, lua: &RawLua) -> Result<Self> {
         let mut values = MultiValue::with_lua_and_capacity(lua.lua(), nvals as usize);
         for idx in 0..nvals {
             values.push_back(lua.stack_value(-nvals + idx));
@@ -935,7 +930,7 @@ pub trait FromLuaMulti: Sized {
         nargs: c_int,
         i: usize,
         to: Option<&str>,
-        lua: &LuaInner,
+        lua: &RawLua,
     ) -> Result<Self> {
         let _ = (i, to);
         Self::from_stack_multi(nargs, lua)

tests/tests.rs

@@ -17,7 +17,7 @@ use mlua::{
 fn test_safety() -> Result<()> {
     let lua = Lua::new();
     assert!(lua.load(r#"require "debug""#).exec().is_err());
-    match lua.load_from_std_lib(StdLib::DEBUG) {
+    match lua.load_std_libs(StdLib::DEBUG) {
         Err(Error::SafetyError(_)) => {}
         Err(e) => panic!("expected SafetyError, got {:?}", e),
         Ok(_) => panic!("expected SafetyError, got no error"),
@@ -53,7 +53,7 @@ fn test_safety() -> Result<()> {
     // Test safety rules after dynamically loading `package` library
     let lua = Lua::new_with(StdLib::NONE, LuaOptions::default())?;
     assert!(lua.globals().get::<_, Option<Value>>("require")?.is_none());
-    lua.load_from_std_lib(StdLib::PACKAGE)?;
+    lua.load_std_libs(StdLib::PACKAGE)?;
     match lua.load(r#"package.loadlib()"#).exec() {
         Err(Error::CallbackError { ref cause, .. }) => match cause.as_ref() {
             Error::SafetyError(_) => {}
@@ -657,7 +657,7 @@ fn test_recursive_mut_callback_error() -> Result<()> {
     let lua = Lua::new();
 
     let mut v = Some(Box::new(123));
-    let f = lua.create_function_mut::<_, (), _>(move |lua, mutate: bool| {
+    let f = lua.create_function_mut(move |lua, mutate: bool| {
         if mutate {
             v = None;
         } else {