Add a WorkerLocal abstraction

3 files changed, 31 insertions, 0 deletions

diff --git a/src/librustc_data_structures/Cargo.toml b/src/librustc_data_structures/Cargo.toml
index 15956829ff9..17ee771e529 100644
--- a/src/librustc_data_structures/Cargo.toml
+++ b/src/librustc_data_structures/Cargo.toml
@@ -17,6 +17,7 @@ cfg-if = "0.1.2"
 stable_deref_trait = "1.0.0"
 parking_lot_core = "0.2.8"
 rustc-rayon = "0.1.0"
+rustc-rayon-core = "0.1.0"
 rustc-hash = "1.0.1"
 
 [dependencies.parking_lot]
diff --git a/src/librustc_data_structures/lib.rs b/src/librustc_data_structures/lib.rs
index 23a920739b9..7046a2a2a49 100644
--- a/src/librustc_data_structures/lib.rs
+++ b/src/librustc_data_structures/lib.rs
@@ -44,6 +44,7 @@ extern crate parking_lot;
 extern crate cfg_if;
 extern crate stable_deref_trait;
 extern crate rustc_rayon as rayon;
+extern crate rustc_rayon_core as rayon_core;
 extern crate rustc_hash;
 
 // See librustc_cratesio_shim/Cargo.toml for a comment explaining this.
diff --git a/src/librustc_data_structures/sync.rs b/src/librustc_data_structures/sync.rs
index 9cceee65bf9..6f7d9e1b54b 100644
--- a/src/librustc_data_structures/sync.rs
+++ b/src/librustc_data_structures/sync.rs
@@ -100,6 +100,33 @@ cfg_if! {
         use std::cell::Cell;
 
         #[derive(Debug)]
+        pub struct WorkerLocal<T>(OneThread<T>);
+
+        impl<T> WorkerLocal<T> {
+            /// Creates a new worker local where the `initial` closure computes the
+            /// value this worker local should take for each thread in the thread pool.
+            #[inline]
+            pub fn new<F: FnMut(usize) -> T>(mut f: F) -> WorkerLocal<T> {
+                WorkerLocal(OneThread::new(f(0)))
+            }
+
+            /// Returns the worker-local value for each thread
+            #[inline]
+            pub fn into_inner(self) -> Vec<T> {
+                vec![OneThread::into_inner(self.0)]
+            }
+        }
+
+        impl<T> Deref for WorkerLocal<T> {
+            type Target = T;
+
+            #[inline(always)]
+            fn deref(&self) -> &T {
+                &*self.0
+            }
+        }
+
+        #[derive(Debug)]
         pub struct MTLock<T>(T);
 
         impl<T> MTLock<T> {
@@ -203,6 +230,8 @@ cfg_if! {
         use std::thread;
         pub use rayon::{join, scope};
 
+        pub use rayon_core::WorkerLocal;
+
         pub use rayon::iter::ParallelIterator;
         use rayon::iter::IntoParallelIterator;