Allow opening jitlib in luajit-sys

crates/luajit-sys/build.rs

@@ -214,6 +214,11 @@ fn build_runtime(src_path: &Path) {
     let status = panic_err!(make.status(), "failed to execute make");
     (!status.success()).then(|| panic!("failed to compile luajit: {status}: {make:?}"));
 
+    println!(
+        "cargo::rustc-env=LUAJIT_SYS_JITLIB={}",
+        src_path.join("jit").display(),
+    );
+
     if feature!("runtime") {
         println!("cargo::rustc-link-search=native={}", src_path.display());
         println!("cargo::rustc-link-lib=static=luajit");

crates/luajit-sys/lib.rs

@@ -1,14 +1,53 @@
 #![allow(nonstandard_style)]
 use std::{ffi::*, ptr};
 
-include!(env!("LUAJIT_SYS_BINDGEN"));
-
 // #[cfg(all(panic = "abort", feature = "unwind"))]
 // compile_error!(r#"feature "unwind" cannot be enabled if panic = "abort""#);
 
 // #[cfg(all(panic = "unwind", not(feature = "unwind")))]
 // compile_error!(r#"feature "unwind" must be enabled if panic = "unwind""#);
 
+include!(env!("LUAJIT_SYS_BINDGEN"));
+
+pub unsafe extern "C" fn luaJIT_openlibs(L: *mut lua_State) {
+    unsafe {
+        lua_getglobal(L, c"package".as_ptr());
+        lua_getfield(L, -1, c"preload".as_ptr());
+        lua_replace(L, -2);
+        macro_rules! load {
+            ($n:literal, $f:literal) => {{
+                let n: &'static CStr = $n;
+                let f = include_bytes!(concat!(env!("LUAJIT_SYS_JITLIB"), "/", $f));
+                if luaL_loadbuffer(L, f.as_ptr().cast(), f.len(), n.as_ptr()) == 0 {
+                    lua_setfield(L, -2, n.as_ptr());
+                } else {
+                    lua_error(L);
+                }
+            }};
+        }
+        load!(c"jit.vmdef", "vmdef.lua");
+        load!(c"jit.dis_x86", "dis_x86.lua");
+        load!(c"jit.dis_x64", "dis_x64.lua");
+        load!(c"jit.dis_arm", "dis_arm.lua");
+        load!(c"jit.dis_arm64", "dis_arm64.lua");
+        load!(c"jit.dis_arm64be", "dis_arm64be.lua");
+        load!(c"jit.dis_ppc", "dis_ppc.lua");
+        load!(c"jit.dis_mips", "dis_mips.lua");
+        load!(c"jit.dis_mipsel", "dis_mipsel.lua");
+        load!(c"jit.dis_mips64", "dis_mips64.lua");
+        load!(c"jit.dis_mips64el", "dis_mips64el.lua");
+        load!(c"jit.dis_mips64r6", "dis_mips64r6.lua");
+        load!(c"jit.dis_mips64r6el", "dis_mips64r6el.lua");
+        load!(c"jit.bc", "bc.lua");
+        load!(c"jit.bcsave", "bcsave.lua");
+        load!(c"jit.v", "v.lua");
+        load!(c"jit.p", "p.lua");
+        load!(c"jit.dump", "dump.lua");
+        load!(c"jit.zone", "zone.lua");
+        lua_pop(L, 1);
+    }
+}
+
 // constants not exposed by lua.h
 pub const LUA_TPROTO: c_int = LUA_TTHREAD + 1;
 pub const LUA_TCDATA: c_int = LUA_TTHREAD + 2;

crates/luajit/src/lib.rs

@@ -505,9 +505,11 @@ impl State {
     }
 
     unsafe extern "C" fn open_cb(L: *mut lua_State) -> c_int {
-        // SAFETY: we wrap luaL_openlibs so we can potentially catch any library open errors
-        unsafe { luaL_openlibs(L) }
-        0
+        unsafe {
+            luaL_openlibs(L);
+            luaJIT_openlibs(L); // luajit-sys extension to open jitlibs
+            0
+        }
     }
 
     /// Creates a new empty thread (coroutine) associated with this state.
@@ -930,6 +932,21 @@ impl Stack {
         n
     }
 
+    /// Pushes the result of a Lua `require(...)` call onto the stack.
+    ///
+    /// Any return values from the library `name` are pushed. Lua libraries are allowed to return
+    /// multiple values. If `nret` is not [`LUA_MULTRET`], then the number of return values pushed
+    /// will be exactly `nret`, filling with nils if necessary. The number values pushed to the
+    /// stack is returned.
+    ///
+    /// Equivalent to `require(name)`.
+    pub fn require(&mut self, name: impl AsRef<[u8]>, nret: c_int) -> Result<c_int> {
+        self.push("require");
+        self.get(LUA_GLOBALSINDEX);
+        self.push(name.as_ref());
+        self.call(1, nret)
+    }
+
     /// Pushes the given chunk as a function at the top of the stack.
     ///
     /// Equivalent to [`lua_loadx`].
@@ -1039,8 +1056,8 @@ impl Stack {
     /// the index `top - narg` (i.e. the function is pushed first and then `narg` values as
     /// arguments). All arguments and the function are popped from the stack and then any return
     /// values are pushed. If `nret` is not [`LUA_MULTRET`], then the number of return values pushed
-    /// will be exactly `nret`, filling with nils if necessary. Finally, the number of returned
-    /// values pushed to the stack is returned.
+    /// will be exactly `nret`, filling with nils if necessary. Finally, the number of values pushed
+    /// to the stack is returned.
     ///
     /// The current thread status must not be suspended or dead.
     ///
@@ -1096,8 +1113,8 @@ impl Stack {
     /// the index `top - narg` (i.e. the function is pushed first and then `narg` values as
     /// arguments). All arguments and the function are popped from the stack and then any return
     /// values are pushed. If `nret` is not [`LUA_MULTRET`], then the number of return values pushed
-    /// will be exactly `nret`, filling with nils if necessary. Finally, the number of returned
-    /// values pushed to the stack is returned.
+    /// will be exactly `nret`, filling with nils if necessary. Finally, the number values pushed to
+    /// the stack is returned.
     ///
     /// If the thread yields a Rust [`Future`] value, then it will be polled to completion before
     /// the thread is resumed with the output of the [`Future`] as the argument. If the thread