Add BlockRng64 wrapper

Author: pitdicker

rand-core/src/impls.rs

@@ -333,6 +333,168 @@ impl<R: BlockRngCore + SeedableRng> SeedableRng for BlockRng<R> {
     }
 }
 
+
+
+/// Wrapper around PRNGs that implement [`BlockRngCore`] to keep a results
+/// buffer and offer the methods from [`RngCore`].
+///
+/// This is similar to [`BlockRng`], but specialized for algorithms that operate
+/// on `u64` values (rare). Hopefully specialization will one day allow
+/// `BlockRng` to support both, so that this wrapper can be removed.
+///
+/// [`BlockRngCore`]: ../BlockRngCore.t.html
+/// [`RngCore`]: ../RngCore.t.html
+/// [`BlockRng`]: .BlockRng.s.html
+#[derive(Clone)]
+pub struct BlockRng64<R: BlockRngCore + ?Sized> {
+    pub results: R::Results,
+    pub index: usize,
+    pub half_used: bool, // true if only half of the previous result is used
+    pub core: R,
+}
+
+// Custom Debug implementation that does not expose the contents of `results`.
+impl<R: BlockRngCore + fmt::Debug> fmt::Debug for BlockRng64<R> {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        fmt.debug_struct("BlockRng64")
+           .field("core", &self.core)
+           .field("result_len", &self.results.as_ref().len())
+           .field("index", &self.index)
+           .field("half_used", &self.half_used)
+           .finish()
+    }
+}
+
+impl<R: BlockRngCore<Item=u64>> RngCore for BlockRng64<R>
+where <R as BlockRngCore>::Results: AsRef<[u64]>
+{
+    #[inline(always)]
+    fn next_u32(&mut self) -> u32 {
+        let mut index = self.index * 2 - self.half_used as usize;
+        if index >= self.results.as_ref().len() * 2 {
+            self.core.generate(&mut self.results);
+            self.index = 0;
+            // `self.half_used` is by definition `false`
+            self.half_used = false;
+            index = 0;
+        }
+
+        self.half_used = !self.half_used;
+        self.index += self.half_used as usize;
+
+        // Index as if this is a u32 slice.
+        unsafe {
+            let results =
+                &*(self.results.as_ref() as *const [u64] as *const [u32]);
+            if cfg!(target_endian = "little") {
+                *results.get_unchecked(index)
+            } else {
+                *results.get_unchecked(index ^ 1)
+            }
+        }
+    }
+
+    #[inline(always)]
+    fn next_u64(&mut self) -> u64 {
+        if self.index >= self.results.as_ref().len() {
+            self.core.generate(&mut self.results);
+            self.index = 0;
+        }
+
+        let value = self.results.as_ref()[self.index];
+        self.index += 1;
+        self.half_used = false;
+        value
+    }
+
+    // As an optimization we try to write directly into the output buffer.
+    // This is only enabled for little-endian platforms where unaligned writes
+    // are known to be safe and fast.
+    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+    fn fill_bytes(&mut self, dest: &mut [u8]) {
+        let mut filled = 0;
+
+        // Continue filling from the current set of results
+        if self.index < self.results.as_ref().len() {
+            let (consumed_u64, filled_u8) =
+                fill_via_u64_chunks(&self.results.as_ref()[self.index..],
+                                    dest);
+
+            self.index += consumed_u64;
+            filled += filled_u8;
+        }
+
+        let len_remainder =
+            (dest.len() - filled) % (self.results.as_ref().len() * 8);
+        let end_direct = dest.len() - len_remainder;
+
+        while filled < end_direct {
+            let dest_u64: &mut R::Results = unsafe {
+                ::core::mem::transmute(dest[filled..].as_mut_ptr())
+            };
+            self.core.generate(dest_u64);
+            filled += self.results.as_ref().len() * 8;
+        }
+        self.index = self.results.as_ref().len();
+
+        if len_remainder > 0 {
+            self.core.generate(&mut self.results);
+            let (consumed_u64, _) =
+                fill_via_u64_chunks(&mut self.results.as_ref(),
+                                    &mut dest[filled..]);
+
+            self.index = consumed_u64;
+        }
+    }
+
+    #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
+    fn fill_bytes(&mut self, dest: &mut [u8]) {
+        let mut read_len = 0;
+        while read_len < dest.len() {
+            if self.index as usize >= self.results.as_ref().len() {
+                self.core.generate(&mut self.results);
+                self.index = 0;
+                self.half_used = false;
+            }
+
+            let (consumed_u64, filled_u8) =
+                fill_via_u64_chunks(&self.results.as_ref()[self.index as usize..],
+                                    &mut dest[read_len..]);
+
+            self.index += consumed_u64;
+            read_len += filled_u8;
+        }
+    }
+
+    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
+        Ok(self.fill_bytes(dest))
+    }
+}
+
+impl<R: BlockRngCore + SeedableRng> SeedableRng for BlockRng64<R> {
+    type Seed = R::Seed;
+
+    fn from_seed(seed: Self::Seed) -> Self {
+        let results_empty = R::Results::default();
+        Self {
+            core: R::from_seed(seed),
+            index: results_empty.as_ref().len(), // generate on first use
+            half_used: false,
+            results: results_empty,
+        }
+    }
+
+    fn from_rng<RNG: RngCore>(rng: &mut RNG) -> Result<Self, Error> {
+        let results_empty = R::Results::default();
+        Ok(Self {
+            core: R::from_rng(rng)?,
+            index: results_empty.as_ref().len(), // generate on first use
+            half_used: false,
+            results: results_empty,
+        })
+    }
+}
+
 impl<R: BlockRngCore + CryptoRng> CryptoRng for BlockRng<R> {}
 
 // TODO: implement tests for the above