From 863e39fe5fcdfb3d33f82d27dbf465201602841d Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@rome.tools>
Date: Sat, 11 Feb 2023 19:26:07 +0100
Subject: [PATCH] perf: Use custom allocator (#2768)

This PR replaces the system allocator with a custom allocator to improve performance:

* Windows: mimalloc
* Unix: tikv-jemallocator

## Performance:

* Linux
  * `cpython --no-cache`: 208.8ms -> 190.5ms
  * `cpython`: 32.8ms -> 31ms
* Mac:
  * `cpython --no-cache`: 436.3ms -> 380ms
  * `cpython`: 40.9ms -> 39.6ms
* Windows:
  * `cpython --no-cache`: 367ms -> 268ms
  * `cpython`: 92.5ms -> 92.3ms

## Size

* Linux: +5MB from 13MB -> 18MB (I need to double check this)
* Mac: +0.7MB from 8.3MB-> 9MB
* Windows: -0.16MB from 8.29MB -> 8.13MB (that's unexpected)
---
 Cargo.lock                  | 41 +++++++++++++++++++++++++++++++++++++
 crates/ruff_cli/Cargo.toml  |  6 ++++++
 crates/ruff_cli/src/main.rs |  8 ++++++++
 3 files changed, 55 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 2a08953234..ec4d9ce163 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1186,6 +1186,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"
@@ -1256,6 +1266,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 = "miniz_oxide"
 version = "0.6.2"
@@ -1970,6 +1989,7 @@ dependencies = [
  "ignore",
  "itertools",
  "log",
+ "mimalloc",
  "notify",
  "path-absolutize",
  "quick-junit",
@@ -1982,6 +2002,7 @@ dependencies = [
  "similar",
  "strum",
  "textwrap",
+ "tikv-jemallocator",
  "ureq",
  "walkdir",
 ]
@@ -2508,6 +2529,26 @@ dependencies = [
  "once_cell",
 ]
 
+[[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"
diff --git a/crates/ruff_cli/Cargo.toml b/crates/ruff_cli/Cargo.toml
index 6b1d98843d..1d1d947a13 100644
--- a/crates/ruff_cli/Cargo.toml
+++ b/crates/ruff_cli/Cargo.toml
@@ -61,3 +61,9 @@ ureq = { version = "2.5.0", features = [] }
 
 [package.metadata.maturin]
 name = "ruff"
+
+[target.'cfg(windows)'.dependencies]
+mimalloc = "0.1.29"
+
+[target.'cfg(not(windows))'.dependencies]
+tikv-jemallocator = "0.5.0"
diff --git a/crates/ruff_cli/src/main.rs b/crates/ruff_cli/src/main.rs
index 7b6de0b7ab..32392d95a3 100644
--- a/crates/ruff_cli/src/main.rs
+++ b/crates/ruff_cli/src/main.rs
@@ -24,6 +24,14 @@ mod iterators;
 mod printer;
 mod resolve;
 
+#[cfg(target_os = "windows")]
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
+#[cfg(not(target_os = "windows"))]
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 enum ExitStatus {
     /// Linting was successful and there were no linting errors.
     Success,