auto merge of #10588 : huonw/rust/un@mutilate-task_rng, r=alexcrichton

Replace with some unsafe code by storing a pointer into TLS-owned heap
data.

3 files changed, 43 insertions, 30 deletions

diff --git a/src/libstd/rand/distributions/mod.rs b/src/libstd/rand/distributions/mod.rs
index 9697fc22ccd..dfdb08b7550 100644
--- a/src/libstd/rand/distributions/mod.rs
+++ b/src/libstd/rand/distributions/mod.rs
@@ -444,17 +444,17 @@ mod tests {
     fn test_rand_sample() {
         let mut rand_sample = RandSample::<ConstRand>;
 
-        assert_eq!(*rand_sample.sample(task_rng()), 0);
-        assert_eq!(*rand_sample.ind_sample(task_rng()), 0);
+        assert_eq!(*rand_sample.sample(&mut task_rng()), 0);
+        assert_eq!(*rand_sample.ind_sample(&mut task_rng()), 0);
     }
 
     #[test]
     fn test_normal() {
         let mut norm = Normal::new(10.0, 10.0);
-        let rng = task_rng();
+        let mut rng = task_rng();
         for _ in range(0, 1000) {
-            norm.sample(rng);
-            norm.ind_sample(rng);
+            norm.sample(&mut rng);
+            norm.ind_sample(&mut rng);
         }
     }
     #[test]
@@ -466,10 +466,10 @@ mod tests {
     #[test]
     fn test_exp() {
         let mut exp = Exp::new(10.0);
-        let rng = task_rng();
+        let mut rng = task_rng();
         for _ in range(0, 1000) {
-            assert!(exp.sample(rng) >= 0.0);
-            assert!(exp.ind_sample(rng) >= 0.0);
+            assert!(exp.sample(&mut rng) >= 0.0);
+            assert!(exp.ind_sample(&mut rng) >= 0.0);
         }
     }
     #[test]
diff --git a/src/libstd/rand/distributions/range.rs b/src/libstd/rand/distributions/range.rs
index 1b805a0b8f7..db9cefa4d79 100644
--- a/src/libstd/rand/distributions/range.rs
+++ b/src/libstd/rand/distributions/range.rs
@@ -183,7 +183,7 @@ mod tests {
 
     #[test]
     fn test_integers() {
-        let rng = task_rng();
+        let mut rng = task_rng();
         macro_rules! t (
             ($($ty:ty),*) => {{
                 $(
@@ -193,9 +193,9 @@ mod tests {
                    for &(low, high) in v.iter() {
                         let mut sampler: Range<$ty> = Range::new(low, high);
                         for _ in range(0, 1000) {
-                            let v = sampler.sample(rng);
+                            let v = sampler.sample(&mut rng);
                             assert!(low <= v && v < high);
-                            let v = sampler.ind_sample(rng);
+                            let v = sampler.ind_sample(&mut rng);
                             assert!(low <= v && v < high);
                         }
                     }
@@ -208,7 +208,7 @@ mod tests {
 
     #[test]
     fn test_floats() {
-        let rng = task_rng();
+        let mut rng = task_rng();
         macro_rules! t (
             ($($ty:ty),*) => {{
                 $(
@@ -219,9 +219,9 @@ mod tests {
                    for &(low, high) in v.iter() {
                         let mut sampler: Range<$ty> = Range::new(low, high);
                         for _ in range(0, 1000) {
-                            let v = sampler.sample(rng);
+                            let v = sampler.sample(&mut rng);
                             assert!(low <= v && v < high);
-                            let v = sampler.ind_sample(rng);
+                            let v = sampler.ind_sample(&mut rng);
                             assert!(low <= v && v < high);
                         }
                     }
diff --git a/src/libstd/rand/mod.rs b/src/libstd/rand/mod.rs
index a2b555028fa..39b4cca2063 100644
--- a/src/libstd/rand/mod.rs
+++ b/src/libstd/rand/mod.rs
@@ -577,11 +577,24 @@ impl reseeding::Reseeder<StdRng> for TaskRngReseeder {
     }
 }
 static TASK_RNG_RESEED_THRESHOLD: uint = 32_768;
+type TaskRngInner = reseeding::ReseedingRng<StdRng, TaskRngReseeder>;
 /// The task-local RNG.
-pub type TaskRng = reseeding::ReseedingRng<StdRng, TaskRngReseeder>;
+#[no_send]
+pub struct TaskRng {
+    // This points into TLS (specifically, it points to the endpoint
+    // of a ~ stored in TLS, to make it robust against TLS moving
+    // things internally) and so this struct cannot be legally
+    // transferred between tasks *and* it's unsafe to deallocate the
+    // RNG other than when a task is finished.
+    //
+    // The use of unsafe code here is OK if the invariants above are
+    // satisfied; and it allows us to avoid (unnecessarily) using a
+    // GC'd or RC'd pointer.
+    priv rng: *mut TaskRngInner
+}
 
 // used to make space in TLS for a random number generator
-local_data_key!(TASK_RNG_KEY: @mut TaskRng)
+local_data_key!(TASK_RNG_KEY: ~TaskRngInner)
 
 /// Retrieve the lazily-initialized task-local random number
 /// generator, seeded by the system. Intended to be used in method
@@ -594,34 +607,34 @@ local_data_key!(TASK_RNG_KEY: @mut TaskRng)
 /// if the operating system random number generator is rigged to give
 /// the same sequence always. If absolute consistency is required,
 /// explicitly select an RNG, e.g. `IsaacRng` or `Isaac64Rng`.
-pub fn task_rng() -> @mut TaskRng {
-    let r = local_data::get(TASK_RNG_KEY, |k| k.map(|k| *k));
-    match r {
+pub fn task_rng() -> TaskRng {
+    local_data::get_mut(TASK_RNG_KEY, |rng| match rng {
         None => {
-            let rng = @mut reseeding::ReseedingRng::new(StdRng::new(),
+            let mut rng = ~reseeding::ReseedingRng::new(StdRng::new(),
                                                         TASK_RNG_RESEED_THRESHOLD,
                                                         TaskRngReseeder);
+            let ptr = &mut *rng as *mut TaskRngInner;
+
             local_data::set(TASK_RNG_KEY, rng);
-            rng
+
+            TaskRng { rng: ptr }
         }
-        Some(rng) => rng
-    }
+        Some(rng) => TaskRng { rng: &mut **rng }
+    })
 }
 
-// Allow direct chaining with `task_rng`
-impl<R: Rng> Rng for @mut R {
-    #[inline]
+impl Rng for TaskRng {
     fn next_u32(&mut self) -> u32 {
-        (**self).next_u32()
+        unsafe { (*self.rng).next_u32() }
     }
-    #[inline]
+
     fn next_u64(&mut self) -> u64 {
-        (**self).next_u64()
+        unsafe { (*self.rng).next_u64() }
     }
 
     #[inline]
     fn fill_bytes(&mut self, bytes: &mut [u8]) {
-        (**self).fill_bytes(bytes);
+        unsafe { (*self.rng).fill_bytes(bytes) }
     }
 }