use jemalloc as global allocator

3 files changed, 53 insertions, 2 deletions

diff --git a/src/tools/miri/Cargo.lock b/src/tools/miri/Cargo.lock
index 62370206956..f6d7776a7ba 100644
--- a/src/tools/miri/Cargo.lock
+++ b/src/tools/miri/Cargo.lock
@@ -418,6 +418,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "af150ab688ff2122fcef229be89cb50dd66af9e01a4ff320cc137eecc9bacc38"
 
 [[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 = "lazy_static"
 version = "1.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -533,6 +543,7 @@ dependencies = [
  "ctrlc",
  "env_logger",
  "getrandom",
+ "jemalloc-sys",
  "lazy_static",
  "libc",
  "libffi",
diff --git a/src/tools/miri/Cargo.toml b/src/tools/miri/Cargo.toml
index f8e507a11b0..d054d9bfac1 100644
--- a/src/tools/miri/Cargo.toml
+++ b/src/tools/miri/Cargo.toml
@@ -24,10 +24,14 @@ log = "0.4"
 rand = "0.8"
 smallvec = "1.7"
 aes = { version = "0.8.3", features = ["hazmat"] }
-
 measureme = "10.0.0"
 ctrlc = "3.2.5"
 
+# Copied from `compiler/rustc/Cargo.toml`.
+[dependencies.jemalloc-sys]
+version = "0.5.0"
+features = ['unprefixed_malloc_on_supported_platforms']
+
 [target.'cfg(unix)'.dependencies]
 libc = "0.2"
 
diff --git a/src/tools/miri/src/bin/miri.rs b/src/tools/miri/src/bin/miri.rs
index ad6a5a669f0..c774961d486 100644
--- a/src/tools/miri/src/bin/miri.rs
+++ b/src/tools/miri/src/bin/miri.rs
@@ -293,13 +293,49 @@ fn run_compiler(
 }
 
 /// Parses a comma separated list of `T` from the given string:
-///
 /// `<value1>,<value2>,<value3>,...`
 fn parse_comma_list<T: FromStr>(input: &str) -> Result<Vec<T>, T::Err> {
     input.split(',').map(str::parse::<T>).collect()
 }
 
+fn jemalloc_magic() {
+    // These magic runes are copied from
+    // <https://github.com/rust-lang/rust/blob/e89bd9428f621545c979c0ec686addc6563a394e/compiler/rustc/src/main.rs#L39>.
+    // See there for further comments.
+    use std::os::raw::{c_int, c_void};
+
+    #[used]
+    static _F1: unsafe extern "C" fn(usize, usize) -> *mut c_void = jemalloc_sys::calloc;
+    #[used]
+    static _F2: unsafe extern "C" fn(*mut *mut c_void, usize, usize) -> c_int =
+        jemalloc_sys::posix_memalign;
+    #[used]
+    static _F3: unsafe extern "C" fn(usize, usize) -> *mut c_void = jemalloc_sys::aligned_alloc;
+    #[used]
+    static _F4: unsafe extern "C" fn(usize) -> *mut c_void = jemalloc_sys::malloc;
+    #[used]
+    static _F5: unsafe extern "C" fn(*mut c_void, usize) -> *mut c_void = jemalloc_sys::realloc;
+    #[used]
+    static _F6: unsafe extern "C" fn(*mut c_void) = jemalloc_sys::free;
+
+    // On OSX, jemalloc doesn't directly override malloc/free, but instead
+    // registers itself with the allocator's zone APIs in a ctor. However,
+    // the linker doesn't seem to consider ctors as "used" when statically
+    // linking, so we need to explicitly depend on the function.
+    #[cfg(target_os = "macos")]
+    {
+        extern "C" {
+            fn _rjem_je_zone_register();
+        }
+
+        #[used]
+        static _F7: unsafe extern "C" fn() = _rjem_je_zone_register;
+    }
+}
+
 fn main() {
+    jemalloc_magic();
+
     let early_dcx = EarlyDiagCtxt::new(ErrorOutputType::default());
 
     // Snapshot a copy of the environment before `rustc` starts messing with it.