diff --git a/Cargo.lock b/Cargo.lock
index 90a3106..7bd61af 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1165,6 +1165,7 @@ dependencies = [
"serde_json",
"serde_yaml",
"thiserror",
+ "tikv-jemallocator",
"tokio",
"tokio-tungstenite",
"toml",
@@ -2349,6 +2350,26 @@ dependencies = [
"once_cell",
]
+[[package]]
+name = "tikv-jemalloc-sys"
+version = "0.5.4+5.3.0-patched"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9402443cb8fd499b6f327e40565234ff34dbda27460c5b47db0db77443dd85d1"
+dependencies = [
+ "cc",
+ "libc",
+]
+
+[[package]]
+name = "tikv-jemallocator"
+version = "0.5.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "965fe0c26be5c56c94e38ba547249074803efd52adfb66de62107d95aab3eaca"
+dependencies = [
+ "libc",
+ "tikv-jemalloc-sys",
+]
+
[[package]]
name = "time"
version = "0.2.27"
diff --git a/Cargo.toml b/Cargo.toml
index 8338685..b0bb3bd 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -83,6 +83,8 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
mlua = { version = "0.9.1", features = ["luau", "luau-jit", "serialize"] }
tokio = { version = "1.24", features = ["full", "tracing"] }
os_str_bytes = { version = "6.4", features = ["conversions"] }
+[target.'cfg(not(target_env = "msvc"))'.dependencies]
+tikv-jemallocator = "0.5"
### SERDE
diff --git a/src/cli/mod.rs b/src/cli/mod.rs
index 065c326..0edb7e6 100644
--- a/src/cli/mod.rs
+++ b/src/cli/mod.rs
@@ -1,9 +1,13 @@
-use std::{env, fmt::Write as _, path::PathBuf, process::ExitCode};
+use std::{env, fmt::Write as _, ops::ControlFlow, path::PathBuf, process::ExitCode, sync::Mutex};
use anyhow::{Context, Result};
use clap::Parser;
use lune::Lune;
+use rayon::{
+ iter::{IndexedParallelIterator, ParallelIterator},
+ slice::ParallelSlice,
+};
use tokio::{
fs::read as read_to_vec,
io::{stdin, AsyncReadExt},
@@ -93,8 +97,8 @@ impl Cli {
let is_standalone = bin[bin.len() - signature.len()..bin.len()] == signature;
if is_standalone {
- let mut bytecode_offset = 0;
- let mut bytecode_size = 0;
+ let bytecode_offset = Mutex::new(0);
+ let bytecode_size = Mutex::new(0);
// standalone binary structure (reversed, 8 bytes per field)
// [0] => signature
@@ -112,31 +116,45 @@ impl Cli {
// The rchunks will have unequally sized sections in the beginning
// but that doesn't matter to us because we don't need anything past the
// middle chunks where the bytecode is stored
- for (idx, chunk) in bin.rchunks(signature.len()).enumerate() {
- if idx == 0 && chunk != signature {
- // Binary is guaranteed to be standalone, we've confirmed this before
- unreachable!()
- }
+ bin.par_rchunks(signature.len())
+ .enumerate()
+ .try_for_each(|(idx, chunk)| {
+ let mut bytecode_offset = bytecode_offset.lock().unwrap();
+ let mut bytecode_size = bytecode_size.lock().unwrap();
- if idx == 3 {
- bytecode_offset = u64::from_ne_bytes(chunk.try_into()?);
- }
+ if *bytecode_offset != 0 && *bytecode_size != 0 {
+ return ControlFlow::Break(());
+ }
- if idx == 2 {
- bytecode_size = u64::from_ne_bytes(chunk.try_into()?);
- }
- }
+ if idx == 0 && chunk != signature {
+ // Binary is guaranteed to be standalone, we've confirmed this before
+ unreachable!("expected proper signature for standalone binary")
+ }
+
+ if idx == 3 {
+ *bytecode_offset = u64::from_ne_bytes(chunk.try_into().unwrap());
+ }
+
+ if idx == 2 {
+ *bytecode_size = u64::from_ne_bytes(chunk.try_into().unwrap());
+ }
+
+ ControlFlow::Continue(())
+ });
+
+ let bytecode_offset_inner = bytecode_offset.into_inner().unwrap();
+ let bytecode_size_inner = bytecode_size.into_inner().unwrap();
// If we were able to retrieve the required metadata, we load
// and execute the bytecode
- if bytecode_offset != 0 && bytecode_size != 0 {
- // FIXME: Passing arguments does not work like it should, because the first
- // argument provided is treated as the script path. We should probably also not
- // allow any runner functionality within standalone binaries
+ // if bytecode_offset_inner != 0 && bytecode_size_inner != 0 {
+ // FIXME: Passing arguments does not work like it should, because the first
+ // argument provided is treated as the script path. We should probably also not
+ // allow any runner functionality within standalone binaries
- let mut reserved_args = Vec::new();
+ let mut reserved_args = Vec::new();
- macro_rules! include_reserved_args {
+ macro_rules! include_reserved_args {
($($arg_bool:expr=> $mapping:literal),*) => {
$(
if $arg_bool {
@@ -146,43 +164,41 @@ impl Cli {
};
}
- let mut real_args = Vec::new();
+ let mut real_args = Vec::new();
- if let Some(first_arg) = self.script_path {
- println!("{first_arg}");
-
- real_args.push(first_arg);
- }
-
- include_reserved_args! {
- self.setup => "--setup",
- self.generate_docs_file => "--generate-docs-file",
- self.generate_selene_types => "--generate-selene-types",
- self.generate_luau_types => "--generate-luau-types",
- self.list => "--list",
- self.build => "--build"
- }
-
- real_args.append(&mut reserved_args);
- real_args.append(&mut self.script_args.clone());
-
- let result = Lune::new()
- .with_args(real_args) // TODO: args should also include lune reserved ones
- .run(
- "STANDALONE",
- &bin[usize::try_from(bytecode_offset)?
- ..usize::try_from(bytecode_offset + bytecode_size)?],
- )
- .await;
-
- return Ok(match result {
- Err(err) => {
- eprintln!("{err}");
- ExitCode::FAILURE
- }
- Ok(code) => code,
- });
+ if let Some(first_arg) = self.script_path {
+ real_args.push(first_arg);
}
+
+ include_reserved_args! {
+ self.setup => "--setup",
+ self.generate_docs_file => "--generate-docs-file",
+ self.generate_selene_types => "--generate-selene-types",
+ self.generate_luau_types => "--generate-luau-types",
+ self.list => "--list",
+ self.build => "--build"
+ }
+
+ real_args.append(&mut reserved_args);
+ real_args.append(&mut self.script_args.clone());
+
+ let result = Lune::new()
+ .with_args(real_args) // TODO: args should also include lune reserved ones
+ .run(
+ "STANDALONE",
+ &bin[usize::try_from(bytecode_offset_inner)?
+ ..usize::try_from(bytecode_offset_inner + bytecode_size_inner)?],
+ )
+ .await;
+
+ return Ok(match result {
+ Err(err) => {
+ eprintln!("{err}");
+ ExitCode::FAILURE
+ }
+ Ok(code) => code,
+ });
+ // }
}
// List files in `lune` and `.lune` directories, if wanted
@@ -284,17 +300,13 @@ impl Cli {
output_path.to_string_lossy()
);
- return Ok(
- match build_standalone(output_path, strip_shebang(script_contents.clone()))
- .await
- {
- Ok(exitcode) => exitcode,
- Err(err) => {
- eprintln!("{err}");
- ExitCode::FAILURE
- }
- },
- );
+ return Ok(match build_standalone(output_path, script_contents).await {
+ Ok(exitcode) => exitcode,
+ Err(err) => {
+ eprintln!("{err}");
+ ExitCode::FAILURE
+ }
+ });
}
// Create a new lune object with all globals & run the script
diff --git a/src/main.rs b/src/main.rs
index cd13300..c2564b6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -16,6 +16,12 @@ pub(crate) mod cli;
use cli::Cli;
use console::style;
+#[cfg(not(target_env = "msvc"))]
+use tikv_jemallocator::Jemalloc;
+
+#[cfg(not(target_env = "msvc"))]
+#[global_allocator]
+static GLOBAL: Jemalloc = Jemalloc;
#[tokio::main(flavor = "multi_thread")]
async fn main() -> ExitCode {