diff --git a/src/distributions/distribution.rs b/src/distributions/distribution.rs index 412c91f736e..e7f7677a85a 100644 --- a/src/distributions/distribution.rs +++ b/src/distributions/distribution.rs @@ -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`. /// @@ -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(&self, rng: &mut R, string: &mut String, len: usize); + + /// Generate a `String` of `len` random chars + #[inline] + fn sample_string(&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] @@ -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())); + } } diff --git a/src/distributions/mod.rs b/src/distributions/mod.rs index ef93fd8665a..e3086680b71 100644 --- a/src/distributions/mod.rs +++ b/src/distributions/mod.rs @@ -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; diff --git a/src/distributions/other.rs b/src/distributions/other.rs index d1f060f483d..2e932703a6d 100644 --- a/src/distributions/other.rs +++ b/src/distributions/other.rs @@ -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")] @@ -85,6 +89,16 @@ impl Distribution 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(&self, rng: &mut R, s: &mut String, len: usize) { + s.reserve(s.len() + 4 * len); + s.extend(Distribution::::sample_iter(self, rng).take(len)); + } +} + impl Distribution for Alphanumeric { fn sample(&self, rng: &mut R) -> u8 { const RANGE: u32 = 26 + 26 + 10; @@ -104,6 +118,17 @@ impl Distribution for Alphanumeric { } } +#[cfg(feature = "alloc")] +impl DistString for Alphanumeric { + fn append_string(&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 for Standard { #[inline] fn sample(&self, rng: &mut R) -> bool {