Optimize and improve the proc_macro RPC interface for cross-thread execution

Cargo.lock

@@ -3840,6 +3840,7 @@ dependencies = [
 name = "rustc_expand"
 version = "0.0.0"
 dependencies = [
+ "crossbeam-channel",
  "rustc_ast",
  "rustc_ast_passes",
  "rustc_ast_pretty",
@@ -3856,6 +3857,7 @@ dependencies = [
  "rustc_span",
  "smallvec",
  "tracing",
+ "unicode-normalization",
 ]
 
 [[package]]

compiler/rustc_expand/Cargo.toml

@@ -25,3 +25,5 @@ rustc_parse = { path = "../rustc_parse" }
 rustc_session = { path = "../rustc_session" }
 smallvec = { version = "1.6.1", features = ["union", "may_dangle"] }
 rustc_ast = { path = "../rustc_ast" }
+crossbeam-channel = "0.5.0"
+unicode-normalization = "0.1.11"

compiler/rustc_expand/src/proc_macro.rs

@@ -12,7 +12,32 @@ use rustc_parse::parser::ForceCollect;
 use rustc_span::def_id::CrateNum;
 use rustc_span::{Span, DUMMY_SP};
 
-const EXEC_STRATEGY: pm::bridge::server::SameThread = pm::bridge::server::SameThread;
+struct CrossbeamMessagePipe<T> {
+    tx: crossbeam_channel::Sender<T>,
+    rx: crossbeam_channel::Receiver<T>,
+}
+
+impl<T> pm::bridge::server::MessagePipe<T> for CrossbeamMessagePipe<T> {
+    fn new() -> (Self, Self) {
+        let (tx1, rx1) = crossbeam_channel::unbounded();
+        let (tx2, rx2) = crossbeam_channel::unbounded();
+        (CrossbeamMessagePipe { tx: tx1, rx: rx2 }, CrossbeamMessagePipe { tx: tx2, rx: rx1 })
+    }
+
+    fn send(&mut self, value: T) {
+        self.tx.send(value).unwrap();
+    }
+
+    fn recv(&mut self) -> Option<T> {
+        self.rx.recv().ok()
+    }
+}
+
+fn exec_strategy(ecx: &ExtCtxt<'_>) -> impl pm::bridge::server::ExecutionStrategy {
+    <pm::bridge::server::MaybeCrossThread<CrossbeamMessagePipe<_>>>::new(
+        ecx.sess.opts.debugging_opts.proc_macro_cross_thread,
+    )
+}
 
 pub struct BangProcMacro {
     pub client: pm::bridge::client::Client<fn(pm::TokenStream) -> pm::TokenStream>,
@@ -27,14 +52,16 @@ impl base::ProcMacro for BangProcMacro {
         input: TokenStream,
     ) -> Result<TokenStream, ErrorReported> {
         let server = proc_macro_server::Rustc::new(ecx, self.krate);
-        self.client.run(&EXEC_STRATEGY, server, input, ecx.ecfg.proc_macro_backtrace).map_err(|e| {
-            let mut err = ecx.struct_span_err(span, "proc macro panicked");
-            if let Some(s) = e.as_str() {
-                err.help(&format!("message: {}", s));
-            }
-            err.emit();
-            ErrorReported
-        })
+        self.client.run(&exec_strategy(ecx), server, input, ecx.ecfg.proc_macro_backtrace).map_err(
+            |e| {
+                let mut err = ecx.struct_span_err(span, "proc macro panicked");
+                if let Some(s) = e.as_str() {
+                    err.help(&format!("message: {}", s));
+                }
+                err.emit();
+                ErrorReported
+            },
+        )
     }
 }
 
@@ -53,7 +80,7 @@ impl base::AttrProcMacro for AttrProcMacro {
     ) -> Result<TokenStream, ErrorReported> {
         let server = proc_macro_server::Rustc::new(ecx, self.krate);
         self.client
-            .run(&EXEC_STRATEGY, server, annotation, annotated, ecx.ecfg.proc_macro_backtrace)
+            .run(&exec_strategy(ecx), server, annotation, annotated, ecx.ecfg.proc_macro_backtrace)
             .map_err(|e| {
                 let mut err = ecx.struct_span_err(span, "custom attribute panicked");
                 if let Some(s) = e.as_str() {
@@ -102,18 +129,22 @@ impl MultiItemModifier for ProcMacroDerive {
         };
 
         let server = proc_macro_server::Rustc::new(ecx, self.krate);
-        let stream =
-            match self.client.run(&EXEC_STRATEGY, server, input, ecx.ecfg.proc_macro_backtrace) {
-                Ok(stream) => stream,
-                Err(e) => {
-                    let mut err = ecx.struct_span_err(span, "proc-macro derive panicked");
-                    if let Some(s) = e.as_str() {
-                        err.help(&format!("message: {}", s));
-                    }
-                    err.emit();
-                    return ExpandResult::Ready(vec![]);
+        let stream = match self.client.run(
+            &exec_strategy(ecx),
+            server,
+            input,
+            ecx.ecfg.proc_macro_backtrace,
+        ) {
+            Ok(stream) => stream,
+            Err(e) => {
+                let mut err = ecx.struct_span_err(span, "proc-macro derive panicked");
+                if let Some(s) = e.as_str() {
+                    err.help(&format!("message: {}", s));
                 }
-            };
+                err.emit();
+                return ExpandResult::Ready(vec![]);
+            }
+        };
 
         let error_count_before = ecx.sess.parse_sess.span_diagnostic.err_count();
         let mut parser =

compiler/rustc_session/src/options.rs

@@ -1207,6 +1207,8 @@ options! {
         "print layout information for each type encountered (default: no)"),
     proc_macro_backtrace: bool = (false, parse_bool, [UNTRACKED],
          "show backtraces for panics during proc-macro execution (default: no)"),
+    proc_macro_cross_thread: bool = (false, parse_bool, [UNTRACKED],
+        "run proc-macro code on a separate thread (default: no)"),
     profile: bool = (false, parse_bool, [TRACKED],
         "insert profiling code (default: no)"),
     profile_closures: bool = (false, parse_no_flag, [UNTRACKED],

library/proc_macro/src/bridge/buffer.rs

@@ -5,6 +5,35 @@ use std::mem;
 use std::ops::{Deref, DerefMut};
 use std::slice;
 
+#[repr(C)]
+pub struct Slice<'a, T> {
+    data: &'a [T; 0],
+    len: usize,
+}
+
+unsafe impl<'a, T: Sync> Sync for Slice<'a, T> {}
+unsafe impl<'a, T: Sync> Send for Slice<'a, T> {}
+
+impl<T> Copy for Slice<'a, T> {}
+impl<T> Clone for Slice<'a, T> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<T> From<&'a [T]> for Slice<'a, T> {
+    fn from(xs: &'a [T]) -> Self {
+        Slice { data: unsafe { &*(xs.as_ptr() as *const [T; 0]) }, len: xs.len() }
+    }
+}
+
+impl<T> Deref for Slice<'a, T> {
+    type Target = [T];
+    fn deref(&self) -> &[T] {
+        unsafe { slice::from_raw_parts(self.data.as_ptr(), self.len) }
+    }
+}
+
 #[repr(C)]
 pub struct Buffer<T: Copy> {
     data: *mut T,

library/proc_macro/src/bridge/handle.rs

@@ -1,9 +1,11 @@
 //! Server-side handles and storage for per-handle data.
 
+use std::borrow::Borrow;
 use std::collections::{BTreeMap, HashMap};
 use std::hash::Hash;
 use std::num::NonZeroU32;
 use std::ops::{Index, IndexMut};
+use std::rc::Rc;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
 pub(super) type Handle = NonZeroU32;
@@ -27,10 +29,14 @@ impl<T> OwnedStore<T> {
     pub(super) fn alloc(&mut self, x: T) -> Handle {
         let counter = self.counter.fetch_add(1, Ordering::SeqCst);
         let handle = Handle::new(counter as u32).expect("`proc_macro` handle counter overflowed");
-        assert!(self.data.insert(handle, x).is_none());
+        self.init(handle, x);
         handle
     }
 
+    pub(super) fn init(&mut self, h: Handle, x: T) {
+        assert!(self.data.insert(h, x).is_none());
+    }
+
     pub(super) fn take(&mut self, h: Handle) -> T {
         self.data.remove(&h).expect("use-after-free in `proc_macro` handle")
     }
@@ -49,22 +55,60 @@ impl<T> IndexMut<Handle> for OwnedStore<T> {
     }
 }
 
+pub(super) trait FromKey<Q: ?Sized> {
+    fn from_key(key: &Q) -> Self;
+}
+
+impl<T: Clone> FromKey<T> for T {
+    fn from_key(key: &T) -> T {
+        key.clone()
+    }
+}
+
+impl<T: ?Sized> FromKey<T> for Rc<T>
+where
+    Rc<T>: for<'a> From<&'a T>,
+{
+    fn from_key(key: &T) -> Rc<T> {
+        key.into()
+    }
+}
+
 pub(super) struct InternedStore<T: 'static> {
     owned: OwnedStore<T>,
     interner: HashMap<T, Handle>,
 }
 
-impl<T: Copy + Eq + Hash> InternedStore<T> {
+impl<T: Clone + Eq + Hash> InternedStore<T> {
     pub(super) fn new(counter: &'static AtomicUsize) -> Self {
         InternedStore { owned: OwnedStore::new(counter), interner: HashMap::new() }
     }
 
-    pub(super) fn alloc(&mut self, x: T) -> Handle {
+    pub(super) fn alloc<'a, Q: ?Sized>(&mut self, x: &'a Q) -> Handle
+    where
+        T: Borrow<Q> + FromKey<Q>,
+        Q: Hash + Eq,
+    {
         let owned = &mut self.owned;
-        *self.interner.entry(x).or_insert_with(|| owned.alloc(x))
+        *self
+            .interner
+            .raw_entry_mut()
+            .from_key(x)
+            .or_insert_with(|| {
+                let own = T::from_key(x);
+                (own.clone(), owned.alloc(own))
+            })
+            .1
     }
 
     pub(super) fn copy(&mut self, h: Handle) -> T {
-        self.owned[h]
+        self.owned[h].clone()
+    }
+}
+
+impl<T> Index<Handle> for InternedStore<T> {
+    type Output = T;
+    fn index(&self, h: Handle) -> &T {
+        self.owned.index(h)
     }
 }

library/proc_macro/src/bridge/rpc.rs

@@ -43,15 +43,17 @@ macro_rules! rpc_encode_decode {
             }
         }
     };
-    (struct $name:ident { $($field:ident),* $(,)? }) => {
-        impl<S> Encode<S> for $name {
+    (struct $name:ident $(<$($T:ident),+>)? { $($field:ident),* $(,)? }) => {
+        impl<S, $($($T: Encode<S>),+)?> Encode<S> for $name $(<$($T),+>)? {
             fn encode(self, w: &mut Writer, s: &mut S) {
                 $(self.$field.encode(w, s);)*
             }
         }
 
-        impl<S> DecodeMut<'_, '_, S> for $name {
-            fn decode(r: &mut Reader<'_>, s: &mut S) -> Self {
+        impl<S, $($($T: for<'s> DecodeMut<'a, 's, S>),+)?> DecodeMut<'a, '_, S>
+            for $name $(<$($T),+>)?
+        {
+            fn decode(r: &mut Reader<'a>, s: &mut S) -> Self {
                 $name {
                     $($field: DecodeMut::decode(r, s)),*
                 }
@@ -126,6 +128,7 @@ impl<S> DecodeMut<'_, '_, S> for u8 {
     }
 }
 
+rpc_encode_decode!(le u16);
 rpc_encode_decode!(le u32);
 rpc_encode_decode!(le usize);
 
@@ -246,6 +249,26 @@ impl<S> DecodeMut<'_, '_, S> for String {
     }
 }
 
+impl<S, T: Encode<S>> Encode<S> for Vec<T> {
+    fn encode(self, w: &mut Writer, s: &mut S) {
+        self.len().encode(w, s);
+        for x in self {
+            x.encode(w, s);
+        }
+    }
+}
+
+impl<S, T: for<'s> DecodeMut<'a, 's, S>> DecodeMut<'a, '_, S> for Vec<T> {
+    fn decode(r: &mut Reader<'a>, s: &mut S) -> Self {
+        let len = usize::decode(r, s);
+        let mut vec = Vec::with_capacity(len);
+        for _ in 0..len {
+            vec.push(T::decode(r, s));
+        }
+        vec
+    }
+}
+
 /// Simplified version of panic payloads, ignoring
 /// types other than `&'static str` and `String`.
 pub enum PanicMessage {

src/tools/tidy/src/deps.rs

@@ -97,6 +97,7 @@ const PERMITTED_DEPENDENCIES: &[&str] = &[
     "compiler_builtins",
     "cpuid-bool",
     "crc32fast",
+    "crossbeam-channel",
     "crossbeam-deque",
     "crossbeam-epoch",
     "crossbeam-queue",