Add a Rayon thread pool

Author: Zoxc

src/librustc/session/mod.rs

@@ -869,10 +869,16 @@ impl Session {
         ret
     }
 
+    /// Returns the number of query threads that should be used for this
+    /// compilation
+    pub fn query_threads_from_opts(opts: &config::Options) -> usize {
+        opts.debugging_opts.query_threads.unwrap_or(1)
+    }
+
     /// Returns the number of query threads that should be used for this
     /// compilation
     pub fn query_threads(&self) -> usize {
-        self.opts.debugging_opts.query_threads.unwrap_or(1)
+        Self::query_threads_from_opts(&self.opts)
     }
 
     /// Returns the number of codegen units that should be used for this

src/librustc/ty/context.rs

@@ -1800,9 +1800,11 @@ pub mod tls {
     /// in librustc otherwise. It is used to when diagnostic messages are
     /// emitted and stores them in the current query, if there is one.
     fn track_diagnostic(diagnostic: &Diagnostic) {
-        with_context(|context| {
-            if let Some(ref query) = context.query {
-                query.diagnostics.lock().push(diagnostic.clone());
+        with_context_opt(|icx| {
+            if let Some(icx) = icx {
+                if let Some(ref query) = icx.query {
+                    query.diagnostics.lock().push(diagnostic.clone());
+                }
             }
         })
     }

src/librustc_driver/Cargo.toml

@@ -13,6 +13,8 @@ arena = { path = "../libarena" }
 graphviz = { path = "../libgraphviz" }
 log = "0.4"
 env_logger = { version = "0.5", default-features = false }
+rustc-rayon = "0.1.0"
+scoped-tls = { version = "0.1.1", features = ["nightly"] }
 rustc = { path = "../librustc" }
 rustc_allocator = { path = "../librustc_allocator" }
 rustc_target = { path = "../librustc_target" }

src/librustc_driver/driver.rs

@@ -49,7 +49,7 @@ use std::fs;
 use std::io::{self, Write};
 use std::iter;
 use std::path::{Path, PathBuf};
-use rustc_data_structures::sync::Lrc;
+use rustc_data_structures::sync::{self, Lrc};
 use std::sync::mpsc;
 use syntax::{self, ast, attr, diagnostics, visit};
 use syntax::ext::base::ExtCtxt;
@@ -64,6 +64,51 @@ use pretty::ReplaceBodyWithLoop;
 
 use profile;
 
+#[cfg(not(parallel_queries))]
+pub fn spawn_thread_pool<F: FnOnce(config::Options) -> R + sync::Send, R: sync::Send>(
+    opts: config::Options,
+    f: F
+) -> R {
+    f(opts)
+}
+
+#[cfg(parallel_queries)]
+pub fn spawn_thread_pool<F: FnOnce(config::Options) -> R + sync::Send, R: sync::Send>(
+    opts: config::Options,
+    f: F
+) -> R {
+    use syntax;
+    use syntax_pos;
+    use rayon::{ThreadPoolBuilder, ThreadPool};
+
+    let config = ThreadPoolBuilder::new().num_threads(Session::query_threads_from_opts(&opts))
+                                         .stack_size(16 * 1024 * 1024);
+
+    let with_pool = move |pool: &ThreadPool| {
+        pool.install(move || f(opts))
+    };
+
+    syntax::GLOBALS.with(|syntax_globals| {
+        syntax_pos::GLOBALS.with(|syntax_pos_globals| {
+            // The main handler run for each Rayon worker thread and sets up
+            // the thread local rustc uses. syntax_globals and syntax_pos_globals are
+            // captured and set on the new threads. ty::tls::with_thread_locals sets up
+            // thread local callbacks from libsyntax
+            let main_handler = move |worker: &mut FnMut()| {
+                syntax::GLOBALS.set(syntax_globals, || {
+                    syntax_pos::GLOBALS.set(syntax_pos_globals, || {
+                        ty::tls::with_thread_locals(|| {
+                            worker()
+                        })
+                    })
+                })
+            };
+
+            ThreadPool::scoped_pool(config, main_handler, with_pool).unwrap()
+        })
+    })
+}
+
 pub fn compile_input(
     trans: Box<TransCrate>,
     sess: &Session,

src/librustc_driver/lib.rs

@@ -35,6 +35,7 @@ extern crate graphviz;
 extern crate env_logger;
 #[cfg(unix)]
 extern crate libc;
+extern crate rustc_rayon as rayon;
 extern crate rustc;
 extern crate rustc_allocator;
 extern crate rustc_target;
@@ -53,6 +54,7 @@ extern crate rustc_save_analysis;
 extern crate rustc_traits;
 extern crate rustc_trans_utils;
 extern crate rustc_typeck;
+extern crate scoped_tls;
 extern crate serialize;
 #[macro_use]
 extern crate log;
@@ -66,7 +68,7 @@ use pretty::{PpMode, UserIdentifiedItem};
 use rustc_resolve as resolve;
 use rustc_save_analysis as save;
 use rustc_save_analysis::DumpHandler;
-use rustc_data_structures::sync::Lrc;
+use rustc_data_structures::sync::{self, Lrc};
 use rustc_data_structures::OnDrop;
 use rustc::session::{self, config, Session, build_session, CompileResult};
 use rustc::session::CompileIncomplete;
@@ -450,36 +452,40 @@ fn get_trans_sysroot(backend_name: &str) -> fn() -> Box<TransCrate> {
 // See comments on CompilerCalls below for details about the callbacks argument.
 // The FileLoader provides a way to load files from sources other than the file system.
 pub fn run_compiler<'a>(args: &[String],
-                        callbacks: &mut CompilerCalls<'a>,
+                        callbacks: &mut (CompilerCalls<'a> + sync::Send),
                         file_loader: Option<Box<FileLoader + Send + Sync + 'static>>,
                         emitter_dest: Option<Box<Write + Send>>)
                         -> (CompileResult, Option<Session>)
 {
     syntax::with_globals(|| {
-        run_compiler_impl(args, callbacks, file_loader, emitter_dest)
+        let matches = match handle_options(args) {
+            Some(matches) => matches,
+            None => return (Ok(()), None),
+        };
+
+        let (sopts, cfg) = config::build_session_options_and_crate_config(&matches);
+
+        driver::spawn_thread_pool(sopts, |sopts| {
+            run_compiler_with_pool(matches, sopts, cfg, callbacks, file_loader, emitter_dest)
+        })
     })
 }
 
-fn run_compiler_impl<'a>(args: &[String],
-                         callbacks: &mut CompilerCalls<'a>,
-                         file_loader: Option<Box<FileLoader + Send + Sync + 'static>>,
-                         emitter_dest: Option<Box<Write + Send>>)
-                         -> (CompileResult, Option<Session>)
-{
+fn run_compiler_with_pool<'a>(
+    matches: getopts::Matches,
+    sopts: config::Options,
+    cfg: ast::CrateConfig,
+    callbacks: &mut (CompilerCalls<'a> + sync::Send),
+    file_loader: Option<Box<FileLoader + Send + Sync + 'static>>,
+    emitter_dest: Option<Box<Write + Send>>
+) -> (CompileResult, Option<Session>) {
     macro_rules! do_or_return {($expr: expr, $sess: expr) => {
         match $expr {
             Compilation::Stop => return (Ok(()), $sess),
             Compilation::Continue => {}
         }
     }}
 
-    let matches = match handle_options(args) {
-        Some(matches) => matches,
-        None => return (Ok(()), None),
-    };
-
-    let (sopts, cfg) = config::build_session_options_and_crate_config(&matches);
-
     let descriptions = diagnostics_registry();
 
     do_or_return!(callbacks.early_callback(&matches,

src/librustc_driver/test.rs

@@ -99,20 +99,25 @@ fn test_env<F>(source_string: &str,
     where F: FnOnce(Env)
 {
     syntax::with_globals(|| {
-        test_env_impl(source_string, args, body)
+        let mut options = config::basic_options();
+        options.debugging_opts.verbose = true;
+        options.unstable_features = UnstableFeatures::Allow;
+
+        driver::spawn_thread_pool(options, |options| {
+            test_env_with_pool(options, source_string, args, body)
+        })
     });
 }
 
-fn test_env_impl<F>(source_string: &str,
-                    (emitter, expected_err_count): (Box<Emitter + sync::Send>, usize),
-                    body: F)
+fn test_env_with_pool<F>(
+    options: config::Options,
+    source_string: &str,
+    (emitter, expected_err_count): (Box<Emitter + sync::Send>, usize),
+    body: F
+)
     where F: FnOnce(Env)
 {
-    let mut options = config::basic_options();
-    options.debugging_opts.verbose = true;
-    options.unstable_features = UnstableFeatures::Allow;
     let diagnostic_handler = errors::Handler::with_emitter(true, false, emitter);
-
     let sess = session::build_session_(options,
                                        None,
                                        diagnostic_handler,

src/librustdoc/clean/mod.rs

@@ -697,7 +697,7 @@ impl<'a> FromIterator<&'a DocFragment> for String {
 pub struct Attributes {
     pub doc_strings: Vec<DocFragment>,
     pub other_attrs: Vec<ast::Attribute>,
-    pub cfg: Option<Rc<Cfg>>,
+    pub cfg: Option<Arc<Cfg>>,
     pub span: Option<syntax_pos::Span>,
     /// map from Rust paths to resolved defs and potential URL fragments
     pub links: Vec<(String, Option<DefId>, Option<String>)>,
@@ -848,7 +848,7 @@ impl Attributes {
         Attributes {
             doc_strings,
             other_attrs,
-            cfg: if cfg == Cfg::True { None } else { Some(Rc::new(cfg)) },
+            cfg: if cfg == Cfg::True { None } else { Some(Arc::new(cfg)) },
             span: sp,
             links: vec![],
         }