From 2945b96e587e162a52c66dfaa95aec8944b11797 Mon Sep 17 00:00:00 2001
From: Alan Egerton <eggyal@gmail.com>
Date: Wed, 28 Apr 2021 14:44:18 +0100
Subject: [PATCH] Multithreading support for lazy-jit

---
 src/driver/jit.rs | 81 +++++++++++++++++++++++++++++++++++++++++++----
 src/lib.rs        |  9 +++++-
 2 files changed, 82 insertions(+), 8 deletions(-)

diff --git a/src/driver/jit.rs b/src/driver/jit.rs
index 4a99cb727c83..3d8f837a66c8 100644
--- a/src/driver/jit.rs
+++ b/src/driver/jit.rs
@@ -3,7 +3,9 @@
 
 use std::cell::RefCell;
 use std::ffi::CString;
+use std::lazy::{Lazy, SyncOnceCell};
 use std::os::raw::{c_char, c_int};
+use std::sync::{mpsc, Mutex};
 
 use cranelift_codegen::binemit::{NullStackMapSink, NullTrapSink};
 use rustc_codegen_ssa::CrateInfo;
@@ -23,6 +25,39 @@ thread_local! {
     static LAZY_JIT_STATE: RefCell<Option<JitState>> = RefCell::new(None);
 }
 
+/// The Sender owned by the rustc thread
+static GLOBAL_MESSAGE_SENDER: SyncOnceCell<Mutex<mpsc::Sender<UnsafeMessage>>> = SyncOnceCell::new();
+
+/// A message that is sent from the jitted runtime to the rustc thread.
+/// Senders are responsible for upholding `Send` semantics.
+enum UnsafeMessage {
+    /// Request that the specified `Instance` be lazily jitted.
+    ///
+    /// Nothing accessible through `instance_ptr` may be moved or mutated by the sender after
+    /// this message is sent.
+    JitFn {
+        instance_ptr: *const Instance<'static>,
+        tx: mpsc::Sender<*const u8>,
+    },
+}
+unsafe impl Send for UnsafeMessage {}
+
+impl UnsafeMessage {
+    /// Send the message.
+    fn send(self) -> Result<(), mpsc::SendError<UnsafeMessage>> {
+        thread_local! {
+            /// The Sender owned by the local thread
+            static LOCAL_MESSAGE_SENDER: Lazy<mpsc::Sender<UnsafeMessage>> = Lazy::new(||
+                GLOBAL_MESSAGE_SENDER
+                    .get().unwrap()
+                    .lock().unwrap()
+                    .clone()
+            );
+        }
+        LOCAL_MESSAGE_SENDER.with(|sender| sender.send(self))
+    }
+}
+
 fn create_jit_module<'tcx>(
     tcx: TyCtxt<'tcx>,
     backend_config: &BackendConfig,
@@ -116,11 +151,6 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
         .chain(backend_config.jit_args.iter().map(|arg| &**arg))
         .map(|arg| CString::new(arg).unwrap())
         .collect::<Vec<_>>();
-    let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();
-
-    // Push a null pointer as a terminating argument. This is required by POSIX and
-    // useful as some dynamic linkers use it as a marker to jump over.
-    argv.push(std::ptr::null());
 
     let start_sig = Signature {
         params: vec![
@@ -141,12 +171,49 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
 
     let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
         unsafe { ::std::mem::transmute(finalized_start) };
-    let ret = f(args.len() as c_int, argv.as_ptr());
-    std::process::exit(ret);
+
+    let (tx, rx) = mpsc::channel();
+    GLOBAL_MESSAGE_SENDER.set(Mutex::new(tx)).unwrap();
+
+    // Spawn the jitted runtime in a new thread so that this rustc thread can handle messages
+    // (eg to lazily JIT further functions as required)
+    std::thread::spawn(move || {
+        let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();
+
+        // Push a null pointer as a terminating argument. This is required by POSIX and
+        // useful as some dynamic linkers use it as a marker to jump over.
+        argv.push(std::ptr::null());
+
+        let ret = f(args.len() as c_int, argv.as_ptr());
+        std::process::exit(ret);
+    });
+
+    // Handle messages
+    loop {
+        match rx.recv().unwrap() {
+            // lazy JIT compilation request - compile requested instance and return pointer to result
+            UnsafeMessage::JitFn { instance_ptr, tx } => {
+                tx.send(jit_fn(instance_ptr))
+                  .expect("jitted runtime hung up before response to lazy JIT request was sent");
+            }
+        }
+    }
 }
 
 #[no_mangle]
 extern "C" fn __clif_jit_fn(instance_ptr: *const Instance<'static>) -> *const u8 {
+    // send the JIT request to the rustc thread, with a channel for the response
+    let (tx, rx) = mpsc::channel();
+    UnsafeMessage::JitFn { instance_ptr, tx }
+        .send()
+        .expect("rustc thread hung up before lazy JIT request was sent");
+
+    // block on JIT compilation result
+    rx.recv()
+      .expect("rustc thread hung up before responding to sent lazy JIT request")
+}
+
+fn jit_fn(instance_ptr: *const Instance<'static>) -> *const u8 {
     rustc_middle::ty::tls::with(|tcx| {
         // lift is used to ensure the correct lifetime for instance.
         let instance = tcx.lift(unsafe { *instance_ptr }).unwrap();
diff --git a/src/lib.rs b/src/lib.rs
index cfc5902cbe3d..91ef62659387 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,11 @@
-#![feature(rustc_private, decl_macro, never_type, hash_drain_filter, vec_into_raw_parts)]
+#![feature(
+    rustc_private,
+    decl_macro,
+    never_type,
+    hash_drain_filter,
+    vec_into_raw_parts,
+    once_cell,
+)]
 #![warn(rust_2018_idioms)]
 #![warn(unused_lifetimes)]
 #![warn(unreachable_pub)]