Skip to content

Commit

Permalink
Add DistString
Browse files Browse the repository at this point in the history
  • Loading branch information
dhardy committed Jun 10, 2021
1 parent 1947c89 commit b4c1d66
Show file tree
Hide file tree
Showing 3 changed files with 64 additions and 1 deletion.
38 changes: 37 additions & 1 deletion src/distributions/distribution.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@

use crate::Rng;
use core::iter;
#[cfg(feature = "alloc")]
use alloc::string::String;

/// Types (distributions) that can be used to create a random instance of `T`.
///
Expand Down Expand Up @@ -187,9 +189,27 @@ where
}
}

/// `String` sampler
///
/// Sampling a `String` of random characters is not quite the same as collecting
/// a sequence of chars. This trait contains some helpers.
#[cfg(feature = "alloc")]
pub trait DistString {
/// Append `len` random chars to `string`
fn append_string<R: Rng + ?Sized>(&self, rng: &mut R, string: &mut String, len: usize);

/// Generate a `String` of `len` random chars
#[inline]
fn sample_string<R: Rng + ?Sized>(&self, rng: &mut R, len: usize) -> String {
let mut s = String::new();
self.append_string(rng, &mut s, len);
s
}
}

#[cfg(test)]
mod tests {
use crate::distributions::{Distribution, Uniform};
use crate::distributions::{Alphanumeric, Distribution, Standard, Uniform};
use crate::Rng;

#[test]
Expand Down Expand Up @@ -233,4 +253,20 @@ mod tests {
}
assert_eq!(count, 10);
}

#[test]
#[cfg(feature = "alloc")]
fn test_dist_string() {
use core::str;
use crate::distributions::DistString;
let mut rng = crate::test::rng(213);

let s1 = Alphanumeric.sample_string(&mut rng, 20);
assert_eq!(s1.len(), 20);
assert_eq!(str::from_utf8(s1.as_bytes()), Ok(s1.as_str()));

let s2 = Standard.sample_string(&mut rng, 20);
assert_eq!(s2.chars().count(), 20);
assert_eq!(str::from_utf8(s2.as_bytes()), Ok(s2.as_str()));
}
}
2 changes: 2 additions & 0 deletions src/distributions/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,8 @@ pub mod weighted;

pub use self::bernoulli::{Bernoulli, BernoulliError};
pub use self::distribution::{Distribution, DistIter, DistMap};
#[cfg(feature = "alloc")]
pub use self::distribution::DistString;
pub use self::float::{Open01, OpenClosed01};
pub use self::other::Alphanumeric;
pub use self::slice::Slice;
Expand Down
25 changes: 25 additions & 0 deletions src/distributions/other.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,12 @@

use core::char;
use core::num::Wrapping;
#[cfg(feature = "alloc")]
use alloc::string::String;

use crate::distributions::{Distribution, Standard, Uniform};
#[cfg(feature = "alloc")]
use crate::distributions::DistString;
use crate::Rng;

#[cfg(feature = "serde1")]
Expand Down Expand Up @@ -85,6 +89,16 @@ impl Distribution<char> for Standard {
}
}

/// Note: the `String` is potentially left with excess capacity; optionally the
/// user may call `string.shrink_to_fit()` afterwards.
#[cfg(feature = "alloc")]
impl DistString for Standard {
fn append_string<R: Rng + ?Sized>(&self, rng: &mut R, s: &mut String, len: usize) {
s.reserve(s.len() + 4 * len);
s.extend(Distribution::<char>::sample_iter(self, rng).take(len));
}
}

impl Distribution<u8> for Alphanumeric {
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u8 {
const RANGE: u32 = 26 + 26 + 10;
Expand All @@ -104,6 +118,17 @@ impl Distribution<u8> for Alphanumeric {
}
}

#[cfg(feature = "alloc")]
impl DistString for Alphanumeric {
fn append_string<R: Rng + ?Sized>(&self, rng: &mut R, string: &mut String, len: usize) {
unsafe {
let v = string.as_mut_vec();
v.reserve(v.len() + len);
v.extend(self.sample_iter(rng).take(len));
}
}
}

impl Distribution<bool> for Standard {
#[inline]
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> bool {
Expand Down

0 comments on commit b4c1d66

Please sign in to comment.