From 2bf7b35268963881754edf1da09cab0b2ba5062c Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Mon, 13 Feb 2023 21:37:22 -0500
Subject: [PATCH] Re-enable custom allocators (#2876)

---
 Cargo.lock                  | 50 +++++++++++++++++++++++++++++++++----
 crates/ruff_cli/Cargo.toml  |  6 +++++
 crates/ruff_cli/src/main.rs | 15 +++++++++++
 3 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 95d1681c31..c3c5a88a01 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1539,6 +1539,16 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "libmimalloc-sys"
+version = "0.1.30"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dd8c7cbf8b89019683667e347572e6d55a7df7ea36b0c4ce69961b0cde67b174"
+dependencies = [
+ "cc",
+ "libc",
+]
+
 [[package]]
 name = "link-cplusplus"
 version = "1.0.8"
@@ -1645,6 +1655,15 @@ dependencies = [
  "autocfg",
 ]
 
+[[package]]
+name = "mimalloc"
+version = "0.1.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9dcb174b18635f7561a0c6c9fc2ce57218ac7523cf72c50af80e2d79ab8f3ba1"
+dependencies = [
+ "libmimalloc-sys",
+]
+
 [[package]]
 name = "mime"
 version = "0.3.16"
@@ -2497,6 +2516,7 @@ dependencies = [
  "itertools",
  "log",
  "mdcat",
+ "mimalloc",
  "notify",
  "path-absolutize",
  "pulldown-cmark",
@@ -2511,6 +2531,7 @@ dependencies = [
  "strum",
  "syntect",
  "textwrap",
+ "tikv-jemallocator",
  "ureq",
  "walkdir",
 ]
@@ -3128,11 +3149,10 @@ dependencies = [
 
 [[package]]
 name = "thread_local"
-version = "1.1.7"
+version = "1.1.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3fdd6f064ccff2d6567adcb3873ca630700f00b5ad3f060c25b5dcfd9a4ce152"
+checksum = "5516c27b78311c50bf42c071425c560ac799b11c30b31f87e3081965fe5e0180"
 dependencies = [
- "cfg-if",
  "once_cell",
 ]
 
@@ -3156,6 +3176,26 @@ dependencies = [
  "weezl",
 ]
 
+[[package]]
+name = "tikv-jemalloc-sys"
+version = "0.5.3+5.3.0-patched"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a678df20055b43e57ef8cddde41cdfda9a3c1a060b67f4c5836dfb1d78543ba8"
+dependencies = [
+ "cc",
+ "libc",
+]
+
+[[package]]
+name = "tikv-jemallocator"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "20612db8a13a6c06d57ec83953694185a367e16945f66565e8028d2c0bd76979"
+dependencies = [
+ "libc",
+ "tikv-jemalloc-sys",
+]
+
 [[package]]
 name = "time"
 version = "0.1.45"
@@ -3231,9 +3271,9 @@ dependencies = [
 
 [[package]]
 name = "tokio-util"
-version = "0.7.7"
+version = "0.7.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5427d89453009325de0d8f342c9490009f76e999cb7672d77e46267448f7e6b2"
+checksum = "bc6a3b08b64e6dfad376fa2432c7b1f01522e37a623c3050bc95db2d3ff21583"
 dependencies = [
  "bytes",
  "futures-core",
diff --git a/crates/ruff_cli/Cargo.toml b/crates/ruff_cli/Cargo.toml
index 5d30cc9674..85a5e03f79 100644
--- a/crates/ruff_cli/Cargo.toml
+++ b/crates/ruff_cli/Cargo.toml
@@ -65,3 +65,9 @@ ureq = { version = "2.5.0", features = [] }
 
 [package.metadata.maturin]
 name = "ruff"
+
+[target.'cfg(windows)'.dependencies]
+mimalloc = "0.1.29"
+
+[target.'cfg(all(not(windows), any(target_arch = "x86_64", target_arch = "aarch64", target_arch = "powerpc64")))'.dependencies]
+tikv-jemallocator = "0.5.0"
diff --git a/crates/ruff_cli/src/main.rs b/crates/ruff_cli/src/main.rs
index e4b92e3e32..cf806cde16 100644
--- a/crates/ruff_cli/src/main.rs
+++ b/crates/ruff_cli/src/main.rs
@@ -24,6 +24,21 @@ mod iterators;
 mod printer;
 mod resolve;
 
+#[cfg(target_os = "windows")]
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
+#[cfg(all(
+    not(target_os = "windows"),
+    any(
+        target_arch = "x86_64",
+        target_arch = "aarch64",
+        target_arch = "powerpc64"
+    )
+))]
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 enum ExitStatus {
     /// Linting was successful and there were no linting errors.
     Success,