Add Lua::set_jit_options with support of Luau JIT inliner

Requires Luau 0.727+

Author: khvzak

mlua-sys/Cargo.toml

@@ -43,7 +43,7 @@ cfg-if = "1.0"
 pkg-config = "0.3.17"
 lua-src = { version = ">= 550.1.0, < 550.2.0", optional = true }
 luajit-src = { version = ">= 210.7.0, < 210.8.0", optional = true }
-luau0-src = { version = "0.20.0", optional = true }
+luau0-src = { version = "0.20.6", optional = true }
 
 [lints.rust]
 unexpected_cfgs = { level = "warn", check-cfg = ['cfg(raw_dylib)'] }

mlua-sys/src/luau/compat.rs

@@ -12,6 +12,11 @@ use super::luacode::*;
 
 pub const LUA_RESUMEERROR: c_int = -1;
 
+// Keep in sync with Bytecode.h
+const LBC_VERSION_MAX: u8 = 11;
+const LBC_TYPE_VERSION_MIN: u8 = 1;
+const LBC_TYPE_VERSION_MAX: u8 = 3;
+
 unsafe fn compat53_reverse(L: *mut lua_State, mut a: c_int, mut b: c_int) {
     while a < b {
         lua_pushvalue(L, a);
@@ -368,6 +373,24 @@ pub unsafe fn luaL_newmetatable(L: *mut lua_State, tname: *const c_char) -> c_in
     }
 }
 
+// Detects whether a chunk is Luau bytecode or text source.
+pub unsafe fn luaL_isbytecode(data: *const c_char, size: usize) -> bool {
+    if size == 0 {
+        return false;
+    }
+    match *data as u8 {
+        b if b < b'\t' => true, // bytecode
+        b if b <= LBC_VERSION_MAX => {
+            let types_version = (size >= 2).then(|| *data.add(1) as u8);
+            match types_version {
+                Some(LBC_TYPE_VERSION_MIN..=LBC_TYPE_VERSION_MAX) => true, // bytecode
+                _ => false,                                                // text
+            }
+        }
+        _ => false, // text
+    }
+}
+
 pub unsafe fn luaL_loadbufferenv(
     L: *mut lua_State,
     data: *const c_char,
@@ -384,19 +407,21 @@ pub unsafe fn luaL_loadbufferenv(
         free(*(data as *mut *mut c_char) as *mut c_void);
     }
 
-    let chunk_is_text = size == 0 || (*data as u8) >= b'\t';
+    let is_bytecode = luaL_isbytecode(data, size);
     if !mode.is_null() {
         let modeb = CStr::from_ptr(mode).to_bytes();
-        if !chunk_is_text && !modeb.contains(&b'b') {
+        let allow_binary = modeb.contains(&b'b');
+        let allow_text = modeb.contains(&b't');
+        if is_bytecode && !allow_binary {
             lua_pushfstring(L, cstr!("attempt to load a binary chunk (mode is '%s')"), mode);
             return LUA_ERRSYNTAX;
-        } else if chunk_is_text && !modeb.contains(&b't') {
+        } else if !is_bytecode && !allow_text {
             lua_pushfstring(L, cstr!("attempt to load a text chunk (mode is '%s')"), mode);
             return LUA_ERRSYNTAX;
         }
     }
 
-    let status = if chunk_is_text {
+    let status = if !is_bytecode {
         if env < 0 {
             env -= 1;
         }

mlua-sys/src/luau/luacodegen.rs

@@ -8,4 +8,7 @@ unsafe extern "C-unwind" {
     pub fn luau_codegen_supported() -> c_int;
     pub fn luau_codegen_create(state: *mut lua_State);
     pub fn luau_codegen_compile(state: *mut lua_State, idx: c_int);
+
+    pub fn luau_enable_jit_inliner(state: *mut lua_State);
+    pub fn luau_disable_jit_inliner(state: *mut lua_State);
 }

src/chunk.rs

@@ -751,7 +751,7 @@ impl Chunk<'_> {
                 return ChunkMode::Binary;
             }
             #[cfg(feature = "luau")]
-            if *source.first().unwrap_or(&u8::MAX) < b'\t' {
+            if unsafe { ffi::luaL_isbytecode(source.as_ptr().cast(), source.len()) } {
                 return ChunkMode::Binary;
             }
         }

src/lib.rs

@@ -153,6 +153,9 @@ pub use crate::userdata::{
 #[doc(inline)]
 pub use crate::debug::HookTriggers;
 
+#[cfg(any(feature = "luau-jit", doc))]
+#[doc(inline)]
+pub use crate::state::JitOptions;
 #[cfg(any(feature = "luau", doc))]
 #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
 pub use crate::{buffer::Buffer, vector::Vector};

src/prelude.rs

@@ -27,7 +27,7 @@ pub use crate::HookTriggers as LuaHookTriggers;
 #[doc(no_inline)]
 pub use crate::state::GcGenParams as LuaGcGenParams;
 
-#[cfg(feature = "luau")]
+#[cfg(any(feature = "luau", doc))]
 #[doc(no_inline)]
 pub use crate::{
     Buffer as LuaBuffer, Vector as LuaVector,
@@ -38,6 +38,10 @@ pub use crate::{
     },
 };
 
+#[cfg(any(feature = "luau-jit", doc))]
+#[doc(no_inline)]
+pub use crate::state::JitOptions as LuaJitOptions;
+
 #[cfg(feature = "async")]
 #[doc(no_inline)]
 pub use crate::{function::LuaNativeAsyncFn, thread::AsyncThread as LuaAsyncThread};

src/state.rs

@@ -255,6 +255,38 @@ impl LuaOptions {
     }
 }
 
+/// Luau JIT options
+#[cfg(any(feature = "luau-jit", doc))]
+#[cfg_attr(docsrs, doc(cfg(feature = "luau-jit")))]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct JitOptions {
+    inliner: bool,
+}
+
+#[cfg(any(feature = "luau-jit", doc))]
+impl Default for JitOptions {
+    fn default() -> Self {
+        const { Self::new() }
+    }
+}
+
+#[cfg(any(feature = "luau-jit", doc))]
+impl JitOptions {
+    /// Creates default JIT options.
+    pub const fn new() -> Self {
+        JitOptions { inliner: false }
+    }
+
+    /// Toggles the runtime bytecode inliner.
+    ///
+    /// Disabled by default. Changing this option does not affect already loaded functions.
+    #[must_use]
+    pub const fn set_inliner(mut self, enabled: bool) -> Self {
+        self.inliner = enabled;
+        self
+    }
+}
+
 impl Drop for Lua {
     fn drop(&mut self) {
         if self.collect_garbage {
@@ -1251,6 +1283,23 @@ impl Lua {
         unsafe { (*lua.extra.get()).enable_jit = enable };
     }
 
+    /// Configures JIT options for this Lua VM.
+    #[cfg(any(feature = "luau-jit", doc))]
+    #[cfg_attr(docsrs, doc(cfg(feature = "luau-jit")))]
+    pub fn set_jit_options(&self, options: JitOptions) {
+        let lua = self.lock();
+        unsafe {
+            let state = lua.main_state();
+            if options.inliner {
+                let _ = Self::set_fflag("LuauCallFeedback", true);
+                let _ = Self::set_fflag("LuauEmitCallFeedback", true);
+                ffi::luau_enable_jit_inliner(state);
+            } else {
+                ffi::luau_disable_jit_inliner(state);
+            }
+        }
+    }
+
     /// Sets Luau feature flag (global setting).
     ///
     /// See https://github.com/luau-lang/luau/blob/master/CONTRIBUTING.md#feature-flags for details.

tests/luau.rs

@@ -358,6 +358,32 @@ fn test_fflags() {
     assert!(Lua::set_fflag("UnknownFlag", true).is_err());
 }
 
+#[cfg(feature = "luau-jit")]
+#[test]
+fn test_jit_inliner() -> Result<()> {
+    let lua = Lua::new();
+    lua.set_jit_options(mlua::JitOptions::new().set_inliner(true));
+
+    // An inlinable helper called in a hot loop.
+    let sum = lua
+        .load(
+            r#"
+            local function add(a, b)
+                return a + b
+            end
+            local sum = 0
+            for i = 1, 1000 do
+                sum = add(sum, i)
+            end
+            return sum
+        "#,
+        )
+        .eval::<i64>()?;
+    assert_eq!(sum, 500500);
+
+    Ok(())
+}
+
 #[test]
 fn test_loadstring() -> Result<()> {
     let lua = Lua::new();