From c70769cc2bd1036ad15b2bf9f824096acf45c3a8 Mon Sep 17 00:00:00 2001
From: Erica Marigold <hi@devcomp.xyz>
Date: Tue, 15 Oct 2024 13:10:44 +0100
Subject: [PATCH 1/5] feat(cli): allow toggling JIT compilation

* Lune now accepts the `LUNE_LUAU_JIT` to toggle JIT compilation of Luau
  code.
* The `Runtime` struct exposes the `with_jit_enabled` method to library
  consumers.
---
 crates/lune/src/cli/run.rs    | 9 ++++++++-
 crates/lune/src/rt/runtime.rs | 9 +++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/crates/lune/src/cli/run.rs b/crates/lune/src/cli/run.rs
index 35e523e6..c180e6ab 100644
--- a/crates/lune/src/cli/run.rs
+++ b/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;
@@ -43,6 +43,13 @@ impl RunCommand {
         // Create a new lune object with all globals & run the script
         let result = 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"
+                )
+            )
             .run(&script_display_name, strip_shebang(script_contents))
             .await;
         Ok(match result {
diff --git a/crates/lune/src/rt/runtime.rs b/crates/lune/src/rt/runtime.rs
index 1aa2acec..641e7bb2 100644
--- a/crates/lune/src/rt/runtime.rs
+++ b/crates/lune/src/rt/runtime.rs
@@ -131,6 +131,15 @@ impl Runtime {
         self
     }
 
+    /**
+        Enables or disables JIT compilation.
+    */
+    #[must_use]
+    pub fn with_jit(self, enabled: bool) -> Self {
+        self.inner.lua().enable_jit(enabled);
+        self
+    }
+
     /**
         Runs a Lune script inside of the current runtime.
 

From 6c391907a5ea5a05e25353324c1e91ac7a772cd3 Mon Sep 17 00:00:00 2001
From: Erica Marigold <hi@devcomp.xyz>
Date: Thu, 17 Oct 2024 06:57:21 +0100
Subject: [PATCH 2/5] chore: apply formatting

---
 crates/lune/src/cli/run.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/lune/src/cli/run.rs b/crates/lune/src/cli/run.rs
index e693e227..dd1fa63c 100644
--- a/crates/lune/src/cli/run.rs
+++ b/crates/lune/src/cli/run.rs
@@ -39,7 +39,7 @@ impl RunCommand {
             let file_display_name = file_path.with_extension("").display().to_string();
             (file_display_name, file_contents)
         };
-      
+
         // Create a new lune runtime with all globals & run the script
         let mut rt = Runtime::new()
             .with_args(self.script_args)

From 1025bb7a3ec0941b49060b6e78a0a30eff81a06e Mon Sep 17 00:00:00 2001
From: Erica Marigold <hi@devcomp.xyz>
Date: Thu, 17 Oct 2024 12:38:21 +0100
Subject: [PATCH 3/5] feat: abstract `JitStatus` struct for luau.load

* Creates a new struct called `JitStatus` in lune_utils, which is placed
  in `AppData` during initialization.
* The luau built-in now respects the JIT preference by using
  `JitStatus`.
---
 Cargo.lock                      | 94 ++++++++++++++++-----------------
 crates/lune-std-luau/src/lib.rs | 10 +++-
 crates/lune-utils/src/jit.rs    | 28 ++++++++++
 crates/lune-utils/src/lib.rs    |  1 +
 crates/lune/src/rt/runtime.rs   | 11 +++-
 5 files changed, 93 insertions(+), 51 deletions(-)
 create mode 100644 crates/lune-utils/src/jit.rs

diff --git a/Cargo.lock b/Cargo.lock
index 3470b8bc..5bc12bed 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -163,9 +163,9 @@ dependencies = [
 
 [[package]]
 name = "async-compression"
-version = "0.4.13"
+version = "0.4.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7e614738943d3f68c628ae3dbce7c3daffb196665f82f8c8ea6b65de73c79429"
+checksum = "103db485efc3e41214fe4fda9f3dbeae2eb9082f48fd236e6095627a9422066e"
 dependencies = [
  "brotli",
  "flate2",
@@ -424,9 +424,9 @@ dependencies = [
 
 [[package]]
 name = "cc"
-version = "1.1.28"
+version = "1.1.30"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2e80e3b6a3ab07840e1cae9b0666a63970dc28e8ed5ffbcdacbfc760c281bfc1"
+checksum = "b16803a61b81d9eabb7eae2588776c4c1e584b738ede45fdbb4c972cec1e9945"
 dependencies = [
  "jobserver",
  "libc",
@@ -486,9 +486,9 @@ dependencies = [
 
 [[package]]
 name = "clap"
-version = "4.5.19"
+version = "4.5.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7be5744db7978a28d9df86a214130d106a89ce49644cbc4e3f0c22c3fba30615"
+checksum = "b97f376d85a664d5837dbae44bf546e6477a679ff6610010f17276f686d867e8"
 dependencies = [
  "clap_builder",
  "clap_derive",
@@ -496,9 +496,9 @@ dependencies = [
 
 [[package]]
 name = "clap_builder"
-version = "4.5.19"
+version = "4.5.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a5fbc17d3ef8278f55b282b2a2e75ae6f6c7d4bb70ed3d0382375104bfafdb4b"
+checksum = "19bc80abd44e4bed93ca373a0704ccbd1b710dc5749406201bb018272808dc54"
 dependencies = [
  "anstream",
  "anstyle",
@@ -1206,9 +1206,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9"
 
 [[package]]
 name = "hyper"
-version = "0.14.30"
+version = "0.14.31"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a152ddd61dfaec7273fe8419ab357f33aee0d914c5f4efbf0d96fa749eea5ec9"
+checksum = "8c08302e8fa335b151b788c775ff56e7a03ae64ff85c548ee820fecb70356e85"
 dependencies = [
  "bytes",
  "futures-channel",
@@ -1230,9 +1230,9 @@ dependencies = [
 
 [[package]]
 name = "hyper"
-version = "1.4.1"
+version = "1.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "50dfd22e0e76d0f662d429a5f80fcaf3855009297eab6a0a9f8543834744ba05"
+checksum = "bbbff0a806a4728c99295b254c8838933b5b082d75e3cb70c8dab21fdfbcfa9a"
 dependencies = [
  "bytes",
  "futures-channel",
@@ -1257,7 +1257,7 @@ checksum = "ec3efd23720e2049821a693cbc7e65ea87c72f1c58ff2f9522ff332b1491e590"
 dependencies = [
  "futures-util",
  "http 0.2.12",
- "hyper 0.14.30",
+ "hyper 0.14.31",
  "rustls 0.21.12",
  "tokio",
  "tokio-rustls 0.24.1",
@@ -1270,7 +1270,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7a343d17fe7885302ed7252767dc7bb83609a874b6ff581142241ec4b73957ad"
 dependencies = [
  "http-body-util",
- "hyper 1.4.1",
+ "hyper 1.5.0",
  "hyper-util",
  "pin-project-lite",
  "tokio",
@@ -1289,7 +1289,7 @@ dependencies = [
  "futures-util",
  "http 1.1.0",
  "http-body 1.0.1",
- "hyper 1.4.1",
+ "hyper 1.5.0",
  "pin-project-lite",
  "socket2",
  "tokio",
@@ -1398,9 +1398,9 @@ dependencies = [
 
 [[package]]
 name = "js-sys"
-version = "0.3.70"
+version = "0.3.72"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1868808506b929d7b0cfa8f75951347aa71bb21144b7791bae35d9bccfcfe37a"
+checksum = "6a88f1bda2bd75b0452a14784937d796722fdebfe50df998aeb3f0b7603019a9"
 dependencies = [
  "wasm-bindgen",
 ]
@@ -1422,9 +1422,9 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
 
 [[package]]
 name = "libc"
-version = "0.2.159"
+version = "0.2.160"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "561d97a539a36e26a9a5fad1ea11a3039a67714694aaa379433e580854bc3dc5"
+checksum = "f0b21006cd1874ae9e650973c565615676dc4a274c965bb0a73796dac838ce4f"
 
 [[package]]
 name = "libloading"
@@ -1612,7 +1612,7 @@ dependencies = [
  "futures-util",
  "http 1.1.0",
  "http-body-util",
- "hyper 1.4.1",
+ "hyper 1.5.0",
  "hyper-tungstenite",
  "hyper-util",
  "lune-std-serde",
@@ -2013,9 +2013,9 @@ checksum = "17359afc20d7ab31fdb42bb844c8b3bb1dabd7dcf7e68428492da7f16966fcef"
 
 [[package]]
 name = "pathdiff"
-version = "0.2.1"
+version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8835116a5c179084a830efb3adc117ab007512b535bc1a21c991d3b32a6b44dd"
+checksum = "d61c5ce1153ab5b689d0c074c4e7fc613e942dfb7dd9eea5ab202d2ad91fe361"
 
 [[package]]
 name = "pbkdf2"
@@ -2142,27 +2142,27 @@ checksum = "dc375e1527247fe1a97d8b7156678dfe7c1af2fc075c9a4db3690ecd2a148068"
 
 [[package]]
 name = "proc-macro2"
-version = "1.0.87"
+version = "1.0.88"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b3e4daa0dcf6feba26f985457cdf104d4b4256fc5a09547140f3631bb076b19a"
+checksum = "7c3a7fc5db1e57d5a779a352c8cdb57b29aa4c40cc69c3a68a7fedc815fbf2f9"
 dependencies = [
  "unicode-ident",
 ]
 
 [[package]]
 name = "profiling"
-version = "1.0.15"
+version = "1.0.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "43d84d1d7a6ac92673717f9f6d1518374ef257669c24ebc5ac25d5033828be58"
+checksum = "afbdc74edc00b6f6a218ca6a5364d6226a259d4b8ea1af4a0ea063f27e179f4d"
 dependencies = [
  "profiling-procmacros",
 ]
 
 [[package]]
 name = "profiling-procmacros"
-version = "1.0.15"
+version = "1.0.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8021cf59c8ec9c432cfc2526ac6b8aa508ecaf29cd415f271b8406c1b851c3fd"
+checksum = "a65f2e60fbf1063868558d69c6beacf412dc755f9fc020f514b7955fc914fe30"
 dependencies = [
  "quote",
  "syn 2.0.79",
@@ -2413,7 +2413,7 @@ dependencies = [
  "h2 0.3.26",
  "http 0.2.12",
  "http-body 0.4.6",
- "hyper 0.14.30",
+ "hyper 0.14.31",
  "hyper-rustls",
  "ipnet",
  "js-sys",
@@ -2575,9 +2575,9 @@ dependencies = [
 
 [[package]]
 name = "rustls-pki-types"
-version = "1.9.0"
+version = "1.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0e696e35370c65c9c541198af4543ccd580cf17fc25d8e05c5a242b202488c55"
+checksum = "16f1201b3c9a7ee8039bcadc17b7e605e2945b27eee7631788c1bd2b0643674b"
 
 [[package]]
 name = "rustls-webpki"
@@ -2602,9 +2602,9 @@ dependencies = [
 
 [[package]]
 name = "rustversion"
-version = "1.0.17"
+version = "1.0.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "955d28af4278de8121b7ebeb796b6a45735dc01436d898801014aced2773a3d6"
+checksum = "0e819f2bc632f285be6d7cd36e25940d45b2391dd6d9b939e79de557f7014248"
 
 [[package]]
 name = "rustyline"
@@ -3499,9 +3499,9 @@ checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
 
 [[package]]
 name = "wasm-bindgen"
-version = "0.2.93"
+version = "0.2.95"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a82edfc16a6c469f5f44dc7b571814045d60404b55a0ee849f9bcfa2e63dd9b5"
+checksum = "128d1e363af62632b8eb57219c8fd7877144af57558fb2ef0368d0087bddeb2e"
 dependencies = [
  "cfg-if",
  "once_cell",
@@ -3510,9 +3510,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-backend"
-version = "0.2.93"
+version = "0.2.95"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9de396da306523044d3302746f1208fa71d7532227f15e347e2d93e4145dd77b"
+checksum = "cb6dd4d3ca0ddffd1dd1c9c04f94b868c37ff5fac97c30b97cff2d74fce3a358"
 dependencies = [
  "bumpalo",
  "log",
@@ -3525,9 +3525,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-futures"
-version = "0.4.43"
+version = "0.4.45"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "61e9300f63a621e96ed275155c108eb6f843b6a26d053f122ab69724559dc8ed"
+checksum = "cc7ec4f8827a71586374db3e87abdb5a2bb3a15afed140221307c3ec06b1f63b"
 dependencies = [
  "cfg-if",
  "js-sys",
@@ -3537,9 +3537,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro"
-version = "0.2.93"
+version = "0.2.95"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "585c4c91a46b072c92e908d99cb1dcdf95c5218eeb6f3bf1efa991ee7a68cccf"
+checksum = "e79384be7f8f5a9dd5d7167216f022090cf1f9ec128e6e6a482a2cb5c5422c56"
 dependencies = [
  "quote",
  "wasm-bindgen-macro-support",
@@ -3547,9 +3547,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro-support"
-version = "0.2.93"
+version = "0.2.95"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "afc340c74d9005395cf9dd098506f7f44e38f2b4a21c6aaacf9a105ea5e1e836"
+checksum = "26c6ab57572f7a24a4985830b120de1594465e5d500f24afe89e16b4e833ef68"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -3560,15 +3560,15 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-shared"
-version = "0.2.93"
+version = "0.2.95"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c62a0a307cb4a311d3a07867860911ca130c3494e8c2719593806c08bc5d0484"
+checksum = "65fc09f10666a9f147042251e0dda9c18f166ff7de300607007e96bdebc1068d"
 
 [[package]]
 name = "web-sys"
-version = "0.3.70"
+version = "0.3.72"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "26fdeaafd9bd129f65e7c031593c24d62186301e0c72c8978fa1678be7d532c0"
+checksum = "f6488b90108c040df0fe62fa815cbdee25124641df01814dd7282749234c6112"
 dependencies = [
  "js-sys",
  "wasm-bindgen",
diff --git a/crates/lune-std-luau/src/lib.rs b/crates/lune-std-luau/src/lib.rs
index 21eb9122..f3a61bdc 100644
--- a/crates/lune-std-luau/src/lib.rs
+++ b/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)
 }
diff --git a/crates/lune-utils/src/jit.rs b/crates/lune-utils/src/jit.rs
new file mode 100644
index 00000000..6f55f71a
--- /dev/null
+++ b/crates/lune-utils/src/jit.rs
@@ -0,0 +1,28 @@
+#[derive(Debug, Clone, Copy, Default)]
+pub struct JitStatus(bool);
+
+impl JitStatus {
+    pub fn new(enabled: bool) -> Self {
+        Self(enabled)
+    }
+
+    pub fn set_status(&mut self, enabled: bool) {
+        self.0 = enabled;
+    }
+
+    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)
+    }
+}
diff --git a/crates/lune-utils/src/lib.rs b/crates/lune-utils/src/lib.rs
index 52743a38..828426fb 100644
--- a/crates/lune-utils/src/lib.rs
+++ b/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;
diff --git a/crates/lune/src/rt/runtime.rs b/crates/lune/src/rt/runtime.rs
index 4e0bbed7..3f7cce6f 100644
--- a/crates/lune/src/rt/runtime.rs
+++ b/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(),
         }
     }
 
@@ -134,8 +137,8 @@ impl Runtime {
         Enables or disables JIT compilation.
     */
     #[must_use]
-    pub fn with_jit(self, enabled: bool) -> Self {
-        self.inner.lua().enable_jit(enabled);
+    pub fn with_jit(mut self, jit_status: impl Into<JitStatus>) -> Self {
+        self.jit_status = jit_status.into();
         self
     }
 
@@ -164,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())

From 4d25f0a66f41f4688d76c03685ed9d50102d693b Mon Sep 17 00:00:00 2001
From: Erica Marigold <hi@devcomp.xyz>
Date: Thu, 17 Oct 2024 12:41:48 +0100
Subject: [PATCH 4/5] fix: apply `must_use` attribute to required methods in
 `JitStatus`

---
 crates/lune-utils/src/jit.rs | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/crates/lune-utils/src/jit.rs b/crates/lune-utils/src/jit.rs
index 6f55f71a..e2ee4751 100644
--- a/crates/lune-utils/src/jit.rs
+++ b/crates/lune-utils/src/jit.rs
@@ -2,6 +2,7 @@
 pub struct JitStatus(bool);
 
 impl JitStatus {
+    #[must_use]
     pub fn new(enabled: bool) -> Self {
         Self(enabled)
     }
@@ -10,6 +11,7 @@ impl JitStatus {
         self.0 = enabled;
     }
 
+    #[must_use]
     pub fn enabled(self) -> bool {
         self.0
     }

From f6c39bc2a941f0361f697cf50132fd21466510af Mon Sep 17 00:00:00 2001
From: Erica Marigold <hi@devcomp.xyz>
Date: Thu, 17 Oct 2024 13:57:49 +0100
Subject: [PATCH 5/5] feat: fix JIT being disabled by default in tests

---
 crates/lune/src/tests.rs | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/crates/lune/src/tests.rs b/crates/lune/src/tests.rs
index d5f56404..ecf01b24 100644
--- a/crates/lune/src/tests.rs
+++ b/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")