Add new module attribule skip_memory_check to improve performance

in module mode (by skipping memory allocations check) with extra risks.

mlua_derive/src/lib.rs

@@ -13,6 +13,7 @@ use {
 #[derive(Default)]
 struct ModuleAttributes {
     name: Option<Ident>,
+    skip_memory_check: bool,
 }
 
 impl ModuleAttributes {
@@ -26,6 +27,11 @@ impl ModuleAttributes {
                     return Err(meta.error("`name` attribute must have a value"));
                 }
             }
+        } else if meta.path.is_ident("skip_memory_check") {
+            if meta.value().is_ok() {
+                return Err(meta.error("`skip_memory_check` attribute have no values"));
+            }
+            self.skip_memory_check = true;
         } else {
             return Err(meta.error("unsupported module attribute"));
         }
@@ -45,6 +51,7 @@ pub fn lua_module(attr: TokenStream, item: TokenStream) -> TokenStream {
     let func_name = &func.sig.ident;
     let module_name = args.name.unwrap_or_else(|| func_name.clone());
     let ext_entrypoint_name = Ident::new(&format!("luaopen_{module_name}"), Span::call_site());
+    let skip_memory_check = args.skip_memory_check;
 
     let wrapped = quote! {
         ::mlua::require_module_feature!();
@@ -53,7 +60,11 @@ pub fn lua_module(attr: TokenStream, item: TokenStream) -> TokenStream {
 
         #[no_mangle]
         unsafe extern "C" fn #ext_entrypoint_name(state: *mut ::mlua::lua_State) -> ::std::os::raw::c_int {
-            ::mlua::Lua::init_from_ptr(state)
+            let lua = ::mlua::Lua::init_from_ptr(state);
+            if #skip_memory_check {
+                lua.skip_memory_check(true);
+            }
+            lua
                 .entrypoint1(#func_name)
                 .expect("cannot initialize module")
         }

src/lib.rs

@@ -235,7 +235,7 @@ pub use mlua_derive::chunk;
 ///
 /// You can also pass options to the attribute:
 ///
-/// name - name of the module, defaults to the name of the function
+/// * name - name of the module, defaults to the name of the function
 ///
 /// ```ignore
 /// #[mlua::lua_module(name = "alt_module")]
@@ -244,6 +244,20 @@ pub use mlua_derive::chunk;
 /// }
 /// ```
 ///
+/// * skip_memory_check - skip memory allocation checks for some operations.
+///
+/// In module mode, mlua runs in unknown environment and cannot say are there any memory
+/// limits or not. As result, some operations that require memory allocation runs in
+/// protected mode. Setting this mode will improve performance of such operations
+/// with risk of having uncaught exceptions and memory leaks.
+///
+/// ```ignore
+/// #[mlua::lua_module(skip_memory_check)]
+/// fn my_module(lua: &Lua) -> Result<Table> {
+///     ...
+/// }
+/// ```
+///
 #[cfg(any(feature = "module", docsrs))]
 #[cfg_attr(docsrs, doc(cfg(feature = "module")))]
 pub use mlua_derive::lua_module;

src/lua.rs

@@ -95,6 +95,8 @@ pub(crate) struct ExtraData {
     safe: bool,
     libs: StdLib,
     mem_state: Option<NonNull<MemoryState>>,
+    #[cfg(feature = "module")]
+    skip_memory_check: bool,
 
     // Auxiliary thread to store references
     ref_thread: *mut ffi::lua_State,
@@ -515,6 +517,8 @@ impl Lua {
             safe: false,
             libs: StdLib::NONE,
             mem_state: None,
+            #[cfg(feature = "module")]
+            skip_memory_check: false,
             ref_thread,
             // We need 1 extra stack space to move values in and out of the ref stack.
             ref_stack_size: ffi::LUA_MINSTACK - 1,
@@ -773,6 +777,13 @@ impl Lua {
         self.entrypoint(move |lua, _: ()| func(lua))
     }
 
+    /// Skips memory checks for some operations.
+    #[doc(hidden)]
+    #[cfg(feature = "module")]
+    pub fn skip_memory_check(&self, skip: bool) {
+        unsafe { (*self.extra.get()).skip_memory_check = skip };
+    }
+
     /// Enables (or disables) sandbox mode on this Lua instance.
     ///
     /// This method, in particular:
@@ -3079,7 +3090,13 @@ impl Lua {
         (*self.extra.get())
             .mem_state
             .map(|x| x.as_ref().memory_limit() == 0)
-            .unwrap_or_default()
+            .unwrap_or_else(|| {
+                // Alternatively, check the special flag (only for module mode)
+                #[cfg(feature = "module")]
+                return (*self.extra.get()).skip_memory_check;
+                #[cfg(not(feature = "module"))]
+                return false;
+            })
     }
 
     #[cfg(feature = "unstable")]

tests/module/src/lib.rs

@@ -26,7 +26,7 @@ struct MyUserData(i32);
 
 impl LuaUserData for MyUserData {}
 
-#[mlua::lua_module(name = "test_module_second")]
+#[mlua::lua_module(name = "test_module_second", skip_memory_check)]
 fn test_module2(lua: &Lua) -> LuaResult<LuaTable> {
     let exports = lua.create_table()?;
     exports.set("userdata", MyUserData(123))?;