From 538f2163847683505f9ac9a5d8a46143a082597a Mon Sep 17 00:00:00 2001
From: _Kerman <kermanx@qq.com>
Date: Fri, 11 Oct 2024 10:49:05 +0800
Subject: [PATCH] perf: enable allocator api

---
 Cargo.lock  | 57 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 Cargo.toml  | 11 +++++++++++
 src/main.rs |  8 ++++++++
 3 files changed, 76 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index a19bd1e9..9f04849a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -91,6 +91,15 @@ dependencies = [
  "rustversion",
 ]
 
+[[package]]
+name = "cc"
+version = "1.1.28"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2e80e3b6a3ab07840e1cae9b0666a63970dc28e8ed5ffbcdacbfc760c281bfc1"
+dependencies = [
+ "shlex",
+]
+
 [[package]]
 name = "cfg-if"
 version = "1.0.0"
@@ -437,6 +446,26 @@ version = "1.0.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b"
 
+[[package]]
+name = "jemalloc-sys"
+version = "0.5.4+5.3.0-patched"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ac6c1946e1cea1788cbfde01c993b52a10e2da07f4bac608228d1bed20bfebf2"
+dependencies = [
+ "cc",
+ "libc",
+]
+
+[[package]]
+name = "jemallocator"
+version = "0.5.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a0de374a9f8e63150e6f5e8a60cc14c668226d7a347d8aee1a45766e3c4dd3bc"
+dependencies = [
+ "jemalloc-sys",
+ "libc",
+]
+
 [[package]]
 name = "js-sys"
 version = "0.3.70"
@@ -488,6 +517,16 @@ dependencies = [
  "windows-targets 0.52.6",
 ]
 
+[[package]]
+name = "libmimalloc-sys"
+version = "0.1.39"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23aa6811d3bd4deb8a84dde645f943476d13b248d818edcf8ce0b2f37f036b44"
+dependencies = [
+ "cc",
+ "libc",
+]
+
 [[package]]
 name = "linked-hash-map"
 version = "0.5.6"
@@ -550,6 +589,15 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "mimalloc"
+version = "0.1.43"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "68914350ae34959d83f732418d51e2427a794055d0b9529f48259ac07af65633"
+dependencies = [
+ "libmimalloc-sys",
+]
+
 [[package]]
 name = "napi"
 version = "2.16.11"
@@ -1195,6 +1243,12 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "shlex"
+version = "1.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
+
 [[package]]
 name = "similar"
 version = "2.6.0"
@@ -1281,11 +1335,14 @@ dependencies = [
 name = "tree-shake"
 version = "0.1.0"
 dependencies = [
+ "allocator-api2",
  "bitflags",
  "codspeed-criterion-compat",
  "criterion",
  "dashmap",
  "insta",
+ "jemallocator",
+ "mimalloc",
  "oxc",
  "regex",
  "rustc-hash",
diff --git a/Cargo.toml b/Cargo.toml
index 24f449e0..effc6267 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,6 +9,7 @@ edition = "2021"
 [lib]
 
 [dependencies]
+allocator-api2 = "0.2.18"
 bitflags = "2.6.0"
 dashmap = "6.0.1"
 insta = { version = "1.39.0", features = ["glob"] }
@@ -16,6 +17,12 @@ oxc = { version = "0.30.0", features = ["codegen", "semantic", "minifier"] }
 regex = "1.10.6"
 rustc-hash = "2.0.0"
 
+[target.'cfg(all(not(target_env = "msvc"), not(target_os = "windows")))'.dependencies]
+jemallocator = { version = "0.5.4", optional = true }
+
+[target.'cfg(target_os = "windows")'.dependencies]
+mimalloc = { version = "0.1.43", optional = true }
+
 [dev-dependencies]
 codspeed-criterion-compat = "2.7.2"
 criterion = "0.5.1"
@@ -33,3 +40,7 @@ codegen-units = 1
 strip = "symbols" # Set to `false` for debug information
 debug = false # Set to `true` for debug information
 panic = "abort" # Let it crash and force ourselves to write safe Rust
+
+[features]
+default = ["allocator"]
+allocator = ["dep:jemallocator", "dep:mimalloc"]
diff --git a/src/main.rs b/src/main.rs
index 0d126808..b76e365d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,14 @@
 use std::{env, fs::File, io::Write, path::Path};
 use tree_shake::{tree_shake, TreeShakeConfig, TreeShakeOptions};
 
+#[cfg(all(feature = "allocator", not(miri), not(target_env = "msvc"), not(target_os = "windows")))]
+#[global_allocator]
+static GLOBAL: jemallocator::Jemalloc = jemallocator::Jemalloc;
+
+#[cfg(all(feature = "allocator", not(miri), target_os = "windows"))]
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
 fn main() {
   let args: Vec<String> = env::args().collect();