Add ThreadLocal

Author: Zoxc

src/librustc/lib.rs

@@ -41,6 +41,7 @@
 #![deny(warnings)]
 
 #![feature(asm)]
+#![feature(attr_literals)]
 #![feature(box_patterns)]
 #![feature(box_syntax)]
 #![feature(conservative_impl_trait)]
@@ -64,6 +65,7 @@
 #![feature(nonzero)]
 #![feature(quote)]
 #![feature(refcell_replace_swap)]
+#![cfg_attr(stage0, feature(repr_align))]
 #![feature(rustc_diagnostic_macros)]
 #![feature(set_stdio)]
 #![feature(slice_patterns)]

src/librustc/util/common.rs

@@ -20,11 +20,53 @@ use std::hash::{Hash, BuildHasher};
 use std::iter::repeat;
 use std::path::Path;
 use std::time::{Duration, Instant};
+use std::cell::UnsafeCell;
 
 use std::sync::mpsc::{Sender};
 use syntax_pos::{SpanData};
 use ty::maps::{QueryMsg};
 use dep_graph::{DepNode};
+use rayon_core::registry::Registry;
+
+scoped_thread_local!(pub static THREAD_INDEX: usize);
+
+#[repr(align(64))]
+struct CacheAligned<T>(T);
+
+// FIXME: Find a way to ensure this isn't transferred between multiple thread pools
+// Thread pools should be the only thing that has a valid THREAD_INDEX.
+// Make it contain a Arc<rayon::Registry> and get the index based on the current worker?
+pub struct ThreadLocal<T>(Vec<CacheAligned<UnsafeCell<T>>>);
+
+unsafe impl<T> Send for ThreadLocal<T> {}
+unsafe impl<T> Sync for ThreadLocal<T> {}
+
+impl<T> ThreadLocal<T> {
+    pub fn new<F>(f: F) -> ThreadLocal<T>
+        where F: Fn() -> T,
+    {
+        let n = Registry::current_num_threads();
+        ThreadLocal((0..(1 + n)).map(|_| CacheAligned(UnsafeCell::new(f()))).collect())
+    }
+
+    pub fn into_inner(self) -> Vec<T> {
+        self.0.into_iter().map(|c| c.0.into_inner()).collect()
+    }
+
+    pub fn current(&self) -> &mut T {
+        use std::ops::Index;
+
+        unsafe {
+            &mut *(self.0.index(THREAD_INDEX.with(|t| *t)).0.get())
+        }
+    }
+}
+
+impl<T> ThreadLocal<Vec<T>> {
+    pub fn collect(self) -> Vec<T> {
+        self.into_inner().into_iter().flat_map(|v| v).collect()
+    }
+}
 
 // The name of the associated type for `Fn` return types
 pub const FN_OUTPUT_NAME: &'static str = "Output";

src/librustc_driver/Cargo.toml

@@ -14,6 +14,7 @@ graphviz = { path = "../libgraphviz" }
 log = "0.4"
 env_logger = { version = "0.4", default-features = false }
 rayon = { git = "https://github.com/Zoxc/rayon.git", branch = "fiber" }
+rayon-core = { git = "https://github.com/Zoxc/rayon.git", branch = "fiber" }
 scoped-tls = { git = "https://github.com/Zoxc/scoped-tls.git", features=["nightly"] }
 rustc = { path = "../librustc" }
 rustc_allocator = { path = "../librustc_allocator" }

src/librustc_driver/driver.rs

@@ -68,7 +68,11 @@ use profile;
 
 #[cfg(not(parallel_queries))]
 pub fn spawn_thread_pool<F: FnOnce(Arc<Mutex<usize>>) -> R, R>(_: &Session, f: F) -> R {
-    f(Arc::new(Mutex::new(0)))
+    use rustc::util::common::THREAD_INDEX;
+
+    THREAD_INDEX.set(&0, || {
+        f(Arc::new(Mutex::new(0)))
+    })
 }
 
 #[cfg(parallel_queries)]
@@ -78,6 +82,8 @@ pub fn spawn_thread_pool<F: FnOnce(Arc<Mutex<usize>>) -> R, R>(sess: &Session, f
     use syntax_pos;
     use scoped_tls::ScopedKey;
     use rayon::{Configuration, ThreadPool};
+    use rayon_core;
+    use rustc::util::common::THREAD_INDEX;
 
     let gcx_ptr = Arc::new(Mutex::new(0));
 
@@ -90,7 +96,9 @@ pub fn spawn_thread_pool<F: FnOnce(Arc<Mutex<usize>>) -> R, R>(sess: &Session, f
 
     let with_pool = move |pool: &ThreadPool| {
         pool.with_global_registry(|| {
-            f(arg_gcx_ptr)
+            THREAD_INDEX.set(&0, || {
+                f(arg_gcx_ptr)
+            })
         })
     };
 
@@ -119,11 +127,16 @@ pub fn spawn_thread_pool<F: FnOnce(Arc<Mutex<usize>>) -> R, R>(sess: &Session, f
         try_with(&syntax::GLOBALS, |syntax_globals| {
             try_with(&syntax_pos::GLOBALS, |syntax_pos_globals| {
                 let main_handler = move |worker: &mut FnMut()| {
-                    maybe_set(&PROFQ_CHAN, prof_chan, || {
-                        maybe_set(&syntax::GLOBALS, syntax_globals, || {
-                            maybe_set(&syntax_pos::GLOBALS, syntax_pos_globals, || {
-                                ty::tls::with_thread_locals(|| {
-                                    worker()
+                    let idx = unsafe {
+                        1 + (*rayon_core::registry::WorkerThread::current()).index()
+                    };
+                    THREAD_INDEX.set(&idx, || {
+                        maybe_set(&PROFQ_CHAN, prof_chan, || {
+                            maybe_set(&syntax::GLOBALS, syntax_globals, || {
+                                maybe_set(&syntax_pos::GLOBALS, syntax_pos_globals, || {
+                                    ty::tls::with_thread_locals(|| {
+                                        worker()
+                                    })
                                 })
                             })
                         })

src/librustc_driver/lib.rs

@@ -34,6 +34,7 @@ extern crate env_logger;
 #[cfg(unix)]
 extern crate libc;
 extern crate rayon;
+extern crate rayon_core;
 extern crate rustc;
 extern crate rustc_allocator;
 extern crate rustc_back;