From 0a13a9631d527e811d84920b575f4592fce958ae Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Sun, 7 Jun 2020 18:48:22 +0100
Subject: [PATCH] Fix LuaJIT 2.1 libraries loading

---
 src/ffi/compat53.rs | 32 ++++++++++++++++++++++----------
 src/lua.rs          |  7 +++++++
 tests/tests.rs      | 12 ++++++++++++
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/src/ffi/compat53.rs b/src/ffi/compat53.rs
index 72a4f39..fb5574d 100644
--- a/src/ffi/compat53.rs
+++ b/src/ffi/compat53.rs
@@ -44,18 +44,18 @@ use super::lua::{
     lua_createtable, lua_dump_old, lua_error, lua_getfield_old, lua_getstack, lua_gettable_old,
     lua_gettop, lua_insert, lua_isstring, lua_istable, lua_newuserdata, lua_pop, lua_pushboolean,
     lua_pushcfunction, lua_pushfstring, lua_pushinteger, lua_pushliteral, lua_pushlstring_old,
-    lua_pushnumber, lua_pushthread, lua_pushvalue, lua_rawget_old, lua_rawgeti_old, lua_rawset,
-    lua_replace, lua_setfield, lua_setglobal, lua_setmetatable, lua_settable, lua_toboolean,
-    lua_tointeger, lua_tolstring, lua_tonumber, lua_topointer, lua_tostring, lua_touserdata,
-    lua_type, lua_typename,
+    lua_pushnil, lua_pushnumber, lua_pushthread, lua_pushvalue, lua_rawget_old, lua_rawgeti_old,
+    lua_rawset, lua_replace, lua_setfield, lua_setglobal, lua_setmetatable, lua_settable,
+    lua_toboolean, lua_tointeger, lua_tolstring, lua_tonumber, lua_topointer, lua_tostring,
+    lua_touserdata, lua_type, lua_typename,
 };
 
 #[cfg(any(feature = "lua51", feature = "luajit"))]
 use super::lua::{
     lua_checkstack, lua_concat, lua_equal, lua_getfenv, lua_getinfo, lua_getmetatable,
     lua_isnumber, lua_lessthan, lua_newtable, lua_next, lua_objlen, lua_pushcclosure,
-    lua_pushlightuserdata, lua_pushnil, lua_pushstring_old, lua_rawequal, lua_remove,
-    lua_resume_old, lua_setfenv, lua_settop, LUA_OPADD, LUA_OPUNM,
+    lua_pushlightuserdata, lua_pushstring_old, lua_rawequal, lua_remove, lua_resume_old,
+    lua_setfenv, lua_settop, LUA_OPADD, LUA_OPUNM,
 };
 
 #[cfg(feature = "lua52")]
@@ -785,13 +785,25 @@ pub unsafe fn luaL_requiref(
         lua_pop(L, 1);
         lua_pushcfunction(L, openf);
         lua_pushstring(L, modname);
-        lua_call(L, 1, 1);
-        lua_pushvalue(L, -1);
-        lua_setfield(L, -3, modname);
+        #[cfg(any(feature = "lua52", feature = "lua51"))]
+        {
+            lua_call(L, 1, 1);
+            lua_pushvalue(L, -1);
+            lua_setfield(L, -3, modname);
+        }
+        #[cfg(feature = "luajit")]
+        {
+            lua_call(L, 1, 0);
+            lua_getfield(L, -1, modname);
+        }
     }
-    if glb != 0 {
+    if cfg!(any(feature = "lua52", feature = "lua51")) && glb != 0 {
         lua_pushvalue(L, -1);
         lua_setglobal(L, modname);
     }
+    if cfg!(feature = "luajit") && glb == 0 {
+        lua_pushnil(L);
+        lua_setglobal(L, modname);
+    }
     lua_replace(L, -2);
 }
diff --git a/src/lua.rs b/src/lua.rs
index 9782e70..ad81558 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -1975,6 +1975,10 @@ impl<'lua, 'a> Chunk<'lua, 'a> {
 }
 
 unsafe fn load_from_std_lib(state: *mut ffi::lua_State, libs: StdLib) {
+    #[cfg(feature = "luajit")]
+    // Stop collector during library initialization
+    ffi::lua_gc(state, ffi::LUA_GCSTOP, 0);
+
     #[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))]
     {
         if libs.contains(StdLib::COROUTINE) {
@@ -2067,6 +2071,9 @@ unsafe fn load_from_std_lib(state: *mut ffi::lua_State, libs: StdLib) {
             ffi::lua_pop(state, 1);
         }
     }
+
+    #[cfg(feature = "luajit")]
+    ffi::lua_gc(state, ffi::LUA_GCRESTART, -1);
 }
 
 unsafe fn ref_stack_pop(extra: &mut ExtraData) -> c_int {
diff --git a/tests/tests.rs b/tests/tests.rs
index b62a68d..64b8d2c 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -885,3 +885,15 @@ fn context_thread_51() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+#[cfg(feature = "luajit")]
+fn test_jit_version() -> Result<()> {
+    let lua = Lua::new();
+    let jit: Table = lua.globals().get("jit")?;
+    assert!(jit
+        .get::<_, String>("version")?
+        .to_str()?
+        .contains("LuaJIT"));
+    Ok(())
+}