Skip to content

Commit cbd2cc7

Browse files
committed
Add fill and try_fill methods to Rng
1 parent 8ce7435 commit cbd2cc7

File tree

1 file changed

+140
-1
lines changed

1 file changed

+140
-1
lines changed

src/lib.rs

Lines changed: 140 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -265,7 +265,7 @@
265265
#[cfg(all(feature="std", not(feature = "log")))] macro_rules! error { ($($x:tt)*) => () }
266266

267267

268-
use core::{marker, mem};
268+
use core::{marker, mem, slice};
269269
#[cfg(feature="std")] use std::cell::RefCell;
270270
#[cfg(feature="std")] use std::rc::Rc;
271271
#[cfg(all(feature="alloc", not(feature="std")))] use alloc::boxed::Box;
@@ -508,6 +508,71 @@ pub trait RngCore {
508508
///
509509
/// [`RngCore`]: trait.RngCore.html
510510
pub trait Rng: RngCore + Sized {
511+
/// Fill `dest` entirely with random bytes, where `dest` is any type
512+
/// supporting [`AsByteSliceMut`], namely slices over primitive integer
513+
/// types (`i8`, `i16`, `u32`, etc.).
514+
///
515+
/// On big-endian platforms this performs byte-swapping to ensure
516+
/// portability of results from reproducible generators.
517+
///
518+
/// This uses [`fill_bytes`] internally which may handle some RNG errors
519+
/// implicitly (e.g. waiting if the OS generator is not ready), but panics
520+
/// on other errors. See also [`try_fill`] which returns errors.
521+
///
522+
/// # Example
523+
///
524+
/// ```rust
525+
/// use rand::{thread_rng, Rng};
526+
///
527+
/// let mut arr = [0i8; 20];
528+
/// thread_rng().try_fill(&mut arr[..]);
529+
/// ```
530+
///
531+
/// [`fill_bytes`]: trait.RngCore.html#method.fill_bytes
532+
/// [`try_fill`]: trait.Rng.html#method.try_fill
533+
/// [`AsByteSliceMut`]: trait.AsByteSliceMut.html
534+
fn fill<T: AsByteSliceMut + ?Sized>(&mut self, dest: &mut T) where Self: Sized {
535+
self.fill_bytes(dest.as_byte_slice_mut());
536+
dest.to_le();
537+
}
538+
539+
/// Fill `dest` entirely with random bytes, where `dest` is any type
540+
/// supporting [`AsByteSliceMut`], namely slices over primitive integer
541+
/// types (`i8`, `i16`, `u32`, etc.).
542+
///
543+
/// On big-endian platforms this performs byte-swapping to ensure
544+
/// portability of results from reproducible generators.
545+
///
546+
/// This uses [`try_fill_bytes`] internally and forwards all RNG errors. In
547+
/// some cases errors may be resolvable; see [`ErrorKind`] and
548+
/// documentation for the RNG in use. If you do not plan to handle these
549+
/// errors you may prefer to use [`fill`].
550+
///
551+
/// # Example
552+
///
553+
/// ```rust
554+
/// # use rand::Error;
555+
/// use rand::{thread_rng, Rng};
556+
///
557+
/// # fn try_inner() -> Result<(), Error> {
558+
/// let mut arr = [0u64; 4];
559+
/// thread_rng().try_fill(&mut arr[..])?;
560+
/// # Ok(())
561+
/// # }
562+
///
563+
/// # try_inner().unwrap()
564+
/// ```
565+
///
566+
/// [`ErrorKind`]: enum.ErrorKind.html
567+
/// [`try_fill_bytes`]: trait.RngCore.html#method.try_fill_bytes
568+
/// [`fill`]: trait.Rng.html#method.fill
569+
/// [`AsByteSliceMut`]: trait.AsByteSliceMut.html
570+
fn try_fill<T: AsByteSliceMut + ?Sized>(&mut self, dest: &mut T) -> Result<(), Error> where Self: Sized {
571+
self.try_fill_bytes(dest.as_byte_slice_mut())?;
572+
dest.to_le();
573+
Ok(())
574+
}
575+
511576
/// Sample a new value, using the given distribution.
512577
///
513578
/// ### Example
@@ -746,6 +811,62 @@ impl<R: RngCore + ?Sized> RngCore for Box<R> {
746811
}
747812
}
748813

814+
/// Trait for casting types to byte slices
815+
///
816+
/// This is used by the [`fill`] and [`try_fill`] methods.
817+
///
818+
/// [`fill`]: trait.Rng.html#method.fill
819+
/// [`try_fill`]: trait.Rng.html#method.try_fill
820+
pub trait AsByteSliceMut {
821+
/// Return a mutable reference to self as a byte slice
822+
fn as_byte_slice_mut<'a>(&'a mut self) -> &'a mut [u8];
823+
824+
/// Call `to_le` on each element (i.e. byte-swap on Big Endian platforms).
825+
fn to_le(&mut self);
826+
}
827+
828+
impl AsByteSliceMut for [u8] {
829+
fn as_byte_slice_mut<'a>(&'a mut self) -> &'a mut [u8] {
830+
self
831+
}
832+
833+
fn to_le(&mut self) {}
834+
}
835+
836+
macro_rules! impl_as_byte_slice {
837+
($t:ty) => {
838+
impl AsByteSliceMut for [$t] {
839+
fn as_byte_slice_mut<'a>(&'a mut self) -> &'a mut [u8] {
840+
unsafe {
841+
slice::from_raw_parts_mut(&mut self[0]
842+
as *mut $t
843+
as *mut u8,
844+
self.len() * mem::size_of::<$t>()
845+
)
846+
}
847+
}
848+
849+
fn to_le(&mut self) {
850+
for mut x in self {
851+
*x = x.to_le();
852+
}
853+
}
854+
}
855+
}
856+
}
857+
858+
impl_as_byte_slice!(u16);
859+
impl_as_byte_slice!(u32);
860+
impl_as_byte_slice!(u64);
861+
#[cfg(feature="i128_support")] impl_as_byte_slice!(u128);
862+
impl_as_byte_slice!(usize);
863+
impl_as_byte_slice!(i8);
864+
impl_as_byte_slice!(i16);
865+
impl_as_byte_slice!(i32);
866+
impl_as_byte_slice!(i64);
867+
#[cfg(feature="i128_support")] impl_as_byte_slice!(i128);
868+
impl_as_byte_slice!(isize);
869+
749870
/// Iterator which will generate a stream of random items.
750871
///
751872
/// This iterator is created via the [`gen_iter`] method on [`Rng`].
@@ -1317,6 +1438,24 @@ mod test {
13171438
}
13181439
}
13191440
}
1441+
1442+
#[test]
1443+
fn test_fill() {
1444+
let x = 9041086907909331047; // a random u64
1445+
let mut rng = ConstRng { i: x };
1446+
1447+
// Convert to byte sequence and back to u64; byte-swap twice if BE.
1448+
let mut array = [0u64; 2];
1449+
rng.fill(&mut array[..]);
1450+
assert_eq!(array, [x, x]);
1451+
assert_eq!(rng.next_u64(), x);
1452+
1453+
// Convert to bytes then u32 in LE order
1454+
let mut array = [0u32; 2];
1455+
rng.fill(&mut array[..]);
1456+
assert_eq!(array, [x as u32, (x >> 32) as u32]);
1457+
assert_eq!(rng.next_u32(), x as u32);
1458+
}
13201459

13211460
#[test]
13221461
fn test_gen_range() {

0 commit comments

Comments
 (0)