Add new feature flag `luau-jit` to enable experimental Luau codegen backend

.github/workflows/main.yml

@@ -9,7 +9,7 @@ jobs:
       matrix:
         os: [ubuntu-22.04, macos-latest, windows-latest]
         rust: [stable]
-        lua: [lua54, lua53, lua52, lua51, luajit, luau]
+        lua: [lua54, lua53, lua52, lua51, luajit, luau, luau-jit]
         include:
           - os: ubuntu-22.04
             target: x86_64-unknown-linux-gnu
@@ -104,7 +104,7 @@ jobs:
       matrix:
         os: [ubuntu-22.04, macos-latest, windows-latest]
         rust: [stable, nightly]
-        lua: [lua54, lua53, lua52, lua51, luajit, luajit52, luau]
+        lua: [lua54, lua53, lua52, lua51, luajit, luajit52, luau, luau-jit]
         include:
           - os: ubuntu-22.04
             target: x86_64-unknown-linux-gnu
@@ -140,7 +140,7 @@ jobs:
       matrix:
         os: [ubuntu-22.04]
         rust: [nightly]
-        lua: [lua54, lua53, lua52, lua51, luajit, luau]
+        lua: [lua54, lua53, lua52, lua51, luajit, luau, luau-jit]
         include:
           - os: ubuntu-22.04
             target: x86_64-unknown-linux-gnu
@@ -222,7 +222,7 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        lua: [lua54, lua53, lua52, lua51, luajit, luau]
+        lua: [lua54, lua53, lua52, lua51, luajit, luau, luau-jit]
     steps:
       - uses: actions/checkout@v3
       - uses: dtolnay/rust-toolchain@stable

Cargo.toml

@@ -32,6 +32,7 @@ lua51 = ["ffi/lua51"]
 luajit = ["ffi/luajit"]
 luajit52 = ["luajit", "ffi/luajit52"]
 luau = ["ffi/luau"]
+luau-jit = ["luau", "ffi/luau-codegen"]
 vendored = ["ffi/vendored"]
 module = ["mlua_derive", "ffi/module"]
 async = ["futures-core", "futures-task", "futures-util"]

README.md

@@ -45,6 +45,7 @@ Below is a list of the available feature flags. By default `mlua` does not enabl
 * `luajit`: activate [LuaJIT] support
 * `luajit52`: activate [LuaJIT] support with partial compatibility with Lua 5.2
 * `luau`: activate [Luau] support (auto vendored mode)
+* `luau-jit`: activate [Luau] support with experimental jit backend. This is unstable feature and not recommended to use.
 * `vendored`: build static Lua(JIT) library from sources during `mlua` compilation using [lua-src] or [luajit-src] crates
 * `module`: enable module mode (building loadable `cdylib` library for Lua)
 * `async`: enable async/await support (any executor can be used, eg. [tokio] or [async-std])

mlua-sys/Cargo.toml

@@ -26,6 +26,7 @@ lua51 = []
 luajit = []
 luajit52 = ["luajit"]
 luau = ["luau0-src"]
+luau-codegen = []
 vendored = ["lua-src", "luajit-src"]
 module = []
 
@@ -37,4 +38,4 @@ cfg-if = "1.0"
 pkg-config = "0.3.17"
 lua-src = { version = ">= 546.0.0, < 550.0.0", optional = true }
 luajit-src = { version = ">= 210.4.0, < 220.0.0", optional = true }
-luau0-src = { version = "0.5.6", optional = true }
+luau0-src = { version = "0.5.8", optional = true }

mlua-sys/build/find_vendored.rs

@@ -5,22 +5,25 @@ use std::path::PathBuf;
 pub fn probe_lua() -> Option<PathBuf> {
     #[cfg(feature = "lua54")]
     let artifacts = lua_src::Build::new().build(lua_src::Lua54);
+
     #[cfg(feature = "lua53")]
     let artifacts = lua_src::Build::new().build(lua_src::Lua53);
+
     #[cfg(feature = "lua52")]
     let artifacts = lua_src::Build::new().build(lua_src::Lua52);
+
     #[cfg(feature = "lua51")]
     let artifacts = lua_src::Build::new().build(lua_src::Lua51);
+
     #[cfg(feature = "luajit")]
-    let artifacts = {
-        let mut builder = luajit_src::Build::new();
-        if cfg!(feature = "luajit52") {
-            builder.lua52compat(true);
-        }
-        builder.build()
-    };
+    let artifacts = luajit_src::Build::new()
+        .lua52compat(cfg!(feature = "luajit52"))
+        .build();
+
     #[cfg(feature = "luau")]
-    let artifacts = luau0_src::Build::new().build();
+    let artifacts = luau0_src::Build::new()
+        .enable_codegen(cfg!(feature = "luau-codegen"))
+        .build();
 
     artifacts.print_cargo_metadata();
 

mlua-sys/src/luau/luacodegen.rs

@@ -0,0 +1,11 @@
+//! Contains definitions from `luacodegen.h`.
+
+use std::os::raw::c_int;
+
+use super::lua::lua_State;
+
+extern "C" {
+    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);
+}

mlua-sys/src/luau/mod.rs

@@ -4,10 +4,12 @@ pub use compat::*;
 pub use lauxlib::*;
 pub use lua::*;
 pub use luacode::*;
+pub use luacodegen::*;
 pub use lualib::*;
 
 pub mod compat;
 pub mod lauxlib;
 pub mod lua;
 pub mod luacode;
+pub mod luacodegen;
 pub mod lualib;

src/lua.rs

@@ -131,6 +131,8 @@ pub(crate) struct ExtraData {
     sandboxed: bool,
     #[cfg(feature = "luau")]
     compiler: Option<Compiler>,
+    #[cfg(feature = "luau-jit")]
+    enable_jit: bool,
 }
 
 /// Mode of the Lua garbage collector (GC).
@@ -395,6 +397,12 @@ impl Lua {
         ffi::luaL_requiref(state, cstr!("_G"), ffi::luaopen_base, 1);
         ffi::lua_pop(state, 1);
 
+        // Init Luau code generator (jit)
+        #[cfg(feature = "luau-jit")]
+        if ffi::luau_codegen_supported() != 0 {
+            ffi::luau_codegen_create(state);
+        }
+
         let lua = Lua::init_from_ptr(state);
         let extra = lua.extra.get();
         (*extra).mem_state = NonNull::new(mem_state);
@@ -532,6 +540,8 @@ impl Lua {
             sandboxed: false,
             #[cfg(feature = "luau")]
             compiler: None,
+            #[cfg(feature = "luau-jit")]
+            enable_jit: true,
         }));
 
         // Store it in the registry
@@ -1297,6 +1307,16 @@ impl Lua {
         unsafe { (*self.extra.get()).compiler = Some(compiler) };
     }
 
+    /// Toggles JIT compilation mode for new chunks of code.
+    ///
+    /// By default JIT is enabled. Changing this option does not have any effect on
+    /// already loaded functions.
+    #[cfg(any(feature = "luau-jit", docsrs))]
+    #[cfg_attr(docsrs, doc(cfg(feature = "luau-jit")))]
+    pub fn enable_jit(&self, enable: bool) {
+        unsafe { (*self.extra.get()).enable_jit = enable };
+    }
+
     /// Returns Lua source code as a `Chunk` builder type.
     ///
     /// In order to actually compile or run the resulting code, you must call [`Chunk::exec`] or
@@ -1351,6 +1371,12 @@ impl Lua {
                         #[cfg(any(feature = "lua51", feature = "luajit", feature = "luau"))]
                         ffi::lua_setfenv(state, -2);
                     }
+
+                    #[cfg(feature = "luau-jit")]
+                    if (*self.extra.get()).enable_jit && ffi::luau_codegen_supported() != 0 {
+                        ffi::luau_codegen_compile(state, -1);
+                    }
+
                     Ok(Function(self.pop_ref()))
                 }
                 err => Err(pop_error(state, err)),

src/util.rs

@@ -460,6 +460,11 @@ unsafe fn init_userdata_metatable_index(state: *mut ffi::lua_State) -> Result<()
         ffi::lua_pushcfunction(state, lua_isfunction_impl);
         ffi::lua_call(state, 2, 1);
 
+        #[cfg(feature = "luau-jit")]
+        if ffi::luau_codegen_supported() != 0 {
+            ffi::luau_codegen_compile(state, -1);
+        }
+
         // Store in the registry
         ffi::lua_pushvalue(state, -1);
         ffi::lua_rawsetp(state, ffi::LUA_REGISTRYINDEX, index_key);
@@ -508,6 +513,11 @@ pub unsafe fn init_userdata_metatable_newindex(state: *mut ffi::lua_State) -> Re
         ffi::lua_pushcfunction(state, lua_isfunction_impl);
         ffi::lua_call(state, 2, 1);
 
+        #[cfg(feature = "luau-jit")]
+        if ffi::luau_codegen_supported() != 0 {
+            ffi::luau_codegen_compile(state, -1);
+        }
+
         // Store in the registry
         ffi::lua_pushvalue(state, -1);
         ffi::lua_rawsetp(state, ffi::LUA_REGISTRYINDEX, newindex_key);