From 4f21176fbebc3e5854660e997274c1b98bb4ada4 Mon Sep 17 00:00:00 2001
From: Joseph Guhlin <joseph.guhlin@gmail.com>
Date: Wed, 8 Jan 2025 11:04:54 +1300
Subject: [PATCH] Switch to jemalloc as default?

---
 minimappers2/Cargo.toml | 14 +++++++++++---
 minimappers2/build.rs   |  7 +++++++
 minimappers2/src/lib.rs | 41 +++++++++++++++++++++++++++++++++++++++--
 3 files changed, 57 insertions(+), 5 deletions(-)
 create mode 100644 minimappers2/build.rs

diff --git a/minimappers2/Cargo.toml b/minimappers2/Cargo.toml
index 5c9c6bc..a44cfa5 100644
--- a/minimappers2/Cargo.toml
+++ b/minimappers2/Cargo.toml
@@ -10,12 +10,20 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 minimap2 = { version = "0.1.23", features = ["simde"], path = ".." }
 crossbeam = "0.8.4"
-mimalloc = {version = "0.1", default-features = false }
-
-pyo3 = { version = "0.22" }
+pyo3 = { version = "0.22", features = ["abi3-py39", "chrono", "extension-module"] }
 polars = "0.45"
 pyo3-polars = "0.19"
 
+[target.'cfg(all(any(not(target_family = "unix"), target_os = "emscripten", allocator = "mimalloc"), not(allocator = "default")))'.dependencies]
+mimalloc = { version = "0.1", default-features = false }
+
+# Feature background_threads is unsupported on MacOS (https://github.com/jemalloc/jemalloc/issues/843).
+[target.'cfg(all(target_family = "unix", not(target_os = "macos"), not(target_os = "emscripten"), not(allocator = "mimalloc"), not(allocator = "default")))'.dependencies]
+jemallocator = { version = "0.5", features = ["disable_initial_exec_tls", "background_threads"] }
+
+[target.'cfg(all(target_family = "unix", target_os = "macos", not(allocator = "mimalloc"), not(allocator = "default")))'.dependencies]
+jemallocator = { version = "0.5", features = ["disable_initial_exec_tls"] }
+
 [profile.release]
 opt-level = 3
 lto = "fat"
diff --git a/minimappers2/build.rs b/minimappers2/build.rs
new file mode 100644
index 0000000..c54b7c4
--- /dev/null
+++ b/minimappers2/build.rs
@@ -0,0 +1,7 @@
+fn main() {
+    println!("cargo::rustc-check-cfg=cfg(allocator, values(\"default\", \"mimalloc\"))");
+    println!(
+        "cargo:rustc-env=TARGET={}",
+        std::env::var("TARGET").unwrap()
+    );
+}
\ No newline at end of file
diff --git a/minimappers2/src/lib.rs b/minimappers2/src/lib.rs
index b4a3f8a..d7c8385 100644
--- a/minimappers2/src/lib.rs
+++ b/minimappers2/src/lib.rs
@@ -2,15 +2,52 @@ use std::num::NonZeroI32;
 use std::sync::{Arc, Mutex};
 
 use crossbeam::queue::ArrayQueue;
-use mimalloc::MiMalloc;
 use minimap2::*;
 
 use polars::{df, prelude::*};
 use pyo3::prelude::*;
 use pyo3_polars::{error::PyPolarsErr, PyDataFrame};
 
+#[cfg(all(
+    target_family = "unix",
+    not(target_os = "emscripten"),
+    not(allocator = "default"),
+    not(allocator = "mimalloc"),
+))]
+use jemallocator::Jemalloc;
+#[cfg(all(
+    not(debug_assertions),
+    not(allocator = "default"),
+    any(
+        not(target_family = "unix"),
+        target_os = "emscripten",
+        allocator = "mimalloc"
+    ),
+))]
+use mimalloc::MiMalloc;
+
 #[global_allocator]
-static GLOBAL: MiMalloc = MiMalloc;
+#[cfg(all(
+    not(debug_assertions),
+    not(allocator = "mimalloc"),
+    not(allocator = "default"),
+    target_family = "unix",
+    not(target_os = "emscripten"),
+))]
+static ALLOC: Jemalloc = Jemalloc;
+
+#[global_allocator]
+#[cfg(all(
+    not(debug_assertions),
+    not(allocator = "default"),
+    any(
+        not(target_family = "unix"),
+        target_os = "emscripten",
+        allocator = "mimalloc"
+    ),
+))]
+static ALLOC: MiMalloc = MiMalloc;
+
 
 mod multithreading;