Allow toggling Luau's JIT in the CLI

crates/lune-std-luau/src/lib.rs

@@ -2,7 +2,7 @@
 
 use mlua::prelude::*;
 
-use lune_utils::TableBuilder;
+use lune_utils::{jit::JitStatus, TableBuilder};
 
 mod options;
 
@@ -78,7 +78,13 @@ fn load_source<'lua>(
     // changed, otherwise disable JIT since it'll fall back anyways
     lua.enable_jit(options.codegen_enabled && !env_changed);
     let function = chunk.into_function()?;
-    lua.enable_jit(true);
+    lua.enable_jit(
+        lua.app_data_ref::<JitStatus>()
+            .ok_or(LuaError::runtime(
+                "Failed to get current JitStatus ref from AppData",
+            ))?
+            .enabled(),
+    );
 
     Ok(function)
 }

crates/lune-utils/src/jit.rs

@@ -0,0 +1,30 @@
+#[derive(Debug, Clone, Copy, Default)]
+pub struct JitStatus(bool);
+
+impl JitStatus {
+    #[must_use]
+    pub fn new(enabled: bool) -> Self {
+        Self(enabled)
+    }
+
+    pub fn set_status(&mut self, enabled: bool) {
+        self.0 = enabled;
+    }
+
+    #[must_use]
+    pub fn enabled(self) -> bool {
+        self.0
+    }
+}
+
+impl From<JitStatus> for bool {
+    fn from(val: JitStatus) -> Self {
+        val.enabled()
+    }
+}
+
+impl From<bool> for JitStatus {
+    fn from(val: bool) -> Self {
+        Self::new(val)
+    }
+}

crates/lune-utils/src/lib.rs

@@ -4,6 +4,7 @@ mod table_builder;
 mod version_string;
 
 pub mod fmt;
+pub mod jit;
 pub mod path;
 
 pub use self::table_builder::TableBuilder;

crates/lune/src/cli/run.rs

@@ -1,4 +1,4 @@
-use std::process::ExitCode;
+use std::{env, process::ExitCode};
 
 use anyhow::{Context, Result};
 use clap::Parser;
@@ -41,7 +41,15 @@ impl RunCommand {
         };
 
         // Create a new lune runtime with all globals & run the script
-        let mut rt = Runtime::new().with_args(self.script_args);
+        let mut rt = Runtime::new()
+            .with_args(self.script_args)
+            // Enable JIT compilation unless it was requested to be disabled
+            .with_jit(
+                !matches!(
+                    env::var("LUNE_LUAU_JIT").ok(), 
+                    Some(jit_enabled) if jit_enabled == "0" || jit_enabled == "false" || jit_enabled == "off"
+                )
+             );
 
         let result = rt
             .run(&script_display_name, strip_shebang(script_contents))

crates/lune/src/rt/runtime.rs

@@ -8,6 +8,7 @@ use std::{
     },
 };
 
+use lune_utils::jit::JitStatus;
 use mlua::prelude::*;
 use mlua_luau_scheduler::{Functions, Scheduler};
 use self_cell::self_cell;
@@ -100,6 +101,7 @@ impl RuntimeInner {
 */
 pub struct Runtime {
     inner: RuntimeInner,
+    jit_status: JitStatus,
 }
 
 impl Runtime {
@@ -113,6 +115,7 @@ impl Runtime {
     pub fn new() -> Self {
         Self {
             inner: RuntimeInner::create().expect("Failed to create runtime"),
+            jit_status: JitStatus::default(),
         }
     }
 
@@ -130,6 +133,15 @@ impl Runtime {
         self
     }
 
+    /**
+        Enables or disables JIT compilation.
+    */
+    #[must_use]
+    pub fn with_jit(mut self, jit_status: impl Into<JitStatus>) -> Self {
+        self.jit_status = jit_status.into();
+        self
+    }
+
     /**
         Runs a Lune script inside of the current runtime.
 
@@ -155,6 +167,10 @@ impl Runtime {
             eprintln!("{}", RuntimeError::from(e));
         });
 
+        // Enable / disable the JIT as requested and store the current status as AppData
+        lua.set_app_data(self.jit_status);
+        lua.enable_jit(self.jit_status.enabled());
+
         // Load our "main" thread
         let main = lua
             .load(script_contents.as_ref())

crates/lune/src/tests.rs

@@ -31,13 +31,15 @@ macro_rules! create_tests {
             // The rest of the test logic can continue as normal
             let full_name = format!("{}/tests/{}.luau", workspace_dir.display(), $value);
             let script = read_to_string(&full_name).await?;
-            let mut lune = Runtime::new().with_args(
-                ARGS
-                    .clone()
-                    .iter()
-                    .map(ToString::to_string)
-                    .collect::<Vec<_>>()
-            );
+            let mut lune = Runtime::new()
+                .with_jit(true)
+                .with_args(
+                    ARGS
+                        .clone()
+                        .iter()
+                        .map(ToString::to_string)
+                        .collect::<Vec<_>>()
+                );
             let script_name = full_name
 				.trim_end_matches(".luau")
 				.trim_end_matches(".lua")