From d194c707c82198864cd1493a9e104f904d354e16 Mon Sep 17 00:00:00 2001
From: bjorn3 <bjorn3@users.noreply.github.com>
Date: Sat, 6 Mar 2021 18:59:01 +0100
Subject: [PATCH] Use jemalloc

---
 Cargo.toml         |  3 ++-
 src/bin/cg_clif.rs | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index 9861af1f8ea..79c9b57532d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -35,9 +35,10 @@ smallvec = "1.6.1"
 #gimli = { path = "../" }
 
 [features]
-default = ["jit", "inline_asm"]
+default = ["jit", "inline_asm", "jemalloc"]
 jit = ["cranelift-jit", "libloading"]
 inline_asm = []
+jemalloc = []
 
 [profile.dev]
 # By compiling dependencies with optimizations, performing tests gets much faster.
diff --git a/src/bin/cg_clif.rs b/src/bin/cg_clif.rs
index 983839d48d2..b8d6958b3b2 100644
--- a/src/bin/cg_clif.rs
+++ b/src/bin/cg_clif.rs
@@ -1,5 +1,7 @@
 #![feature(rustc_private)]
 
+#[cfg(feature = "jemalloc")]
+extern crate jemalloc_sys;
 extern crate rustc_data_structures;
 extern crate rustc_driver;
 extern crate rustc_interface;
@@ -33,6 +35,46 @@ fn config(&mut self, config: &mut interface::Config) {
 }
 
 fn main() {
+    // Pull in jemalloc when enabled.
+    //
+    // Note that we're pulling in a static copy of jemalloc which means that to
+    // pull it in we need to actually reference its symbols for it to get
+    // linked. The two crates we link to here, std and rustc_driver, are both
+    // dynamic libraries. That means to pull in jemalloc we need to actually
+    // reference allocation symbols one way or another (as this file is the only
+    // object code in the rustc executable).
+    #[cfg(feature = "jemalloc")]
+    {
+        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;
+        }
+    }
+
     let start_time = std::time::Instant::now();
     let start_rss = get_resident_set_size();
     rustc_driver::init_rustc_env_logger();