1// Copyright 2018 Developers of the Rand project.
2// Copyright 2013-2017 The Rust Project Developers.
3//
4// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
5// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
6// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
7// option. This file may not be copied, modified, or distributed
8// except according to those terms.
9
10//! Distribution trait and associates
11
12use crate::Rng;
13use core::iter;
14#[cfg(feature = "alloc")]
15use alloc::string::String;
16
17/// Types (distributions) that can be used to create a random instance of `T`.
18///
19/// It is possible to sample from a distribution through both the
20/// `Distribution` and [`Rng`] traits, via `distr.sample(&mut rng)` and
21/// `rng.sample(distr)`. They also both offer the [`sample_iter`] method, which
22/// produces an iterator that samples from the distribution.
23///
24/// All implementations are expected to be immutable; this has the significant
25/// advantage of not needing to consider thread safety, and for most
26/// distributions efficient state-less sampling algorithms are available.
27///
28/// Implementations are typically expected to be portable with reproducible
29/// results when used with a PRNG with fixed seed; see the
30/// [portability chapter](https://rust-random.github.io/book/portability.html)
31/// of The Rust Rand Book. In some cases this does not apply, e.g. the `usize`
32/// type requires different sampling on 32-bit and 64-bit machines.
33///
34/// [`sample_iter`]: Distribution::sample_iter
35pub trait Distribution<T> {
36 /// Generate a random value of `T`, using `rng` as the source of randomness.
37 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> T;
38
39 /// Create an iterator that generates random values of `T`, using `rng` as
40 /// the source of randomness.
41 ///
42 /// Note that this function takes `self` by value. This works since
43 /// `Distribution<T>` is impl'd for `&D` where `D: Distribution<T>`,
44 /// however borrowing is not automatic hence `distr.sample_iter(...)` may
45 /// need to be replaced with `(&distr).sample_iter(...)` to borrow or
46 /// `(&*distr).sample_iter(...)` to reborrow an existing reference.
47 ///
48 /// # Example
49 ///
50 /// ```
51 /// use rand::thread_rng;
52 /// use rand::distributions::{Distribution, Alphanumeric, Uniform, Standard};
53 ///
54 /// let mut rng = thread_rng();
55 ///
56 /// // Vec of 16 x f32:
57 /// let v: Vec<f32> = Standard.sample_iter(&mut rng).take(16).collect();
58 ///
59 /// // String:
60 /// let s: String = Alphanumeric
61 /// .sample_iter(&mut rng)
62 /// .take(7)
63 /// .map(char::from)
64 /// .collect();
65 ///
66 /// // Dice-rolling:
67 /// let die_range = Uniform::new_inclusive(1, 6);
68 /// let mut roll_die = die_range.sample_iter(&mut rng);
69 /// while roll_die.next().unwrap() != 6 {
70 /// println!("Not a 6; rolling again!");
71 /// }
72 /// ```
73 fn sample_iter<R>(self, rng: R) -> DistIter<Self, R, T>
74 where
75 R: Rng,
76 Self: Sized,
77 {
78 DistIter {
79 distr: self,
80 rng,
81 phantom: ::core::marker::PhantomData,
82 }
83 }
84
85 /// Create a distribution of values of 'S' by mapping the output of `Self`
86 /// through the closure `F`
87 ///
88 /// # Example
89 ///
90 /// ```
91 /// use rand::thread_rng;
92 /// use rand::distributions::{Distribution, Uniform};
93 ///
94 /// let mut rng = thread_rng();
95 ///
96 /// let die = Uniform::new_inclusive(1, 6);
97 /// let even_number = die.map(|num| num % 2 == 0);
98 /// while !even_number.sample(&mut rng) {
99 /// println!("Still odd; rolling again!");
100 /// }
101 /// ```
102 fn map<F, S>(self, func: F) -> DistMap<Self, F, T, S>
103 where
104 F: Fn(T) -> S,
105 Self: Sized,
106 {
107 DistMap {
108 distr: self,
109 func,
110 phantom: ::core::marker::PhantomData,
111 }
112 }
113}
114
115impl<'a, T, D: Distribution<T>> Distribution<T> for &'a D {
116 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> T {
117 (*self).sample(rng)
118 }
119}
120
121/// An iterator that generates random values of `T` with distribution `D`,
122/// using `R` as the source of randomness.
123///
124/// This `struct` is created by the [`sample_iter`] method on [`Distribution`].
125/// See its documentation for more.
126///
127/// [`sample_iter`]: Distribution::sample_iter
128#[derive(Debug)]
129pub struct DistIter<D, R, T> {
130 distr: D,
131 rng: R,
132 phantom: ::core::marker::PhantomData<T>,
133}
134
135impl<D, R, T> Iterator for DistIter<D, R, T>
136where
137 D: Distribution<T>,
138 R: Rng,
139{
140 type Item = T;
141
142 #[inline(always)]
143 fn next(&mut self) -> Option<T> {
144 // Here, self.rng may be a reference, but we must take &mut anyway.
145 // Even if sample could take an R: Rng by value, we would need to do this
146 // since Rng is not copyable and we cannot enforce that this is "reborrowable".
147 Some(self.distr.sample(&mut self.rng))
148 }
149
150 fn size_hint(&self) -> (usize, Option<usize>) {
151 (usize::max_value(), None)
152 }
153}
154
155impl<D, R, T> iter::FusedIterator for DistIter<D, R, T>
156where
157 D: Distribution<T>,
158 R: Rng,
159{
160}
161
162#[cfg(features = "nightly")]
163impl<D, R, T> iter::TrustedLen for DistIter<D, R, T>
164where
165 D: Distribution<T>,
166 R: Rng,
167{
168}
169
170/// A distribution of values of type `S` derived from the distribution `D`
171/// by mapping its output of type `T` through the closure `F`.
172///
173/// This `struct` is created by the [`Distribution::map`] method.
174/// See its documentation for more.
175#[derive(Debug)]
176pub struct DistMap<D, F, T, S> {
177 distr: D,
178 func: F,
179 phantom: ::core::marker::PhantomData<fn(T) -> S>,
180}
181
182impl<D, F, T, S> Distribution<S> for DistMap<D, F, T, S>
183where
184 D: Distribution<T>,
185 F: Fn(T) -> S,
186{
187 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> S {
188 (self.func)(self.distr.sample(rng))
189 }
190}
191
192/// `String` sampler
193///
194/// Sampling a `String` of random characters is not quite the same as collecting
195/// a sequence of chars. This trait contains some helpers.
196#[cfg(feature = "alloc")]
197pub trait DistString {
198 /// Append `len` random chars to `string`
199 fn append_string<R: Rng + ?Sized>(&self, rng: &mut R, string: &mut String, len: usize);
200
201 /// Generate a `String` of `len` random chars
202 #[inline]
203 fn sample_string<R: Rng + ?Sized>(&self, rng: &mut R, len: usize) -> String {
204 let mut s: String = String::new();
205 self.append_string(rng, &mut s, len);
206 s
207 }
208}
209
210#[cfg(test)]
211mod tests {
212 use crate::distributions::{Distribution, Uniform};
213 use crate::Rng;
214
215 #[test]
216 fn test_distributions_iter() {
217 use crate::distributions::Open01;
218 let mut rng = crate::test::rng(210);
219 let distr = Open01;
220 let mut iter = Distribution::<f32>::sample_iter(distr, &mut rng);
221 let mut sum: f32 = 0.;
222 for _ in 0..100 {
223 sum += iter.next().unwrap();
224 }
225 assert!(0. < sum && sum < 100.);
226 }
227
228 #[test]
229 fn test_distributions_map() {
230 let dist = Uniform::new_inclusive(0, 5).map(|val| val + 15);
231
232 let mut rng = crate::test::rng(212);
233 let val = dist.sample(&mut rng);
234 assert!((15..=20).contains(&val));
235 }
236
237 #[test]
238 fn test_make_an_iter() {
239 fn ten_dice_rolls_other_than_five<R: Rng>(
240 rng: &mut R,
241 ) -> impl Iterator<Item = i32> + '_ {
242 Uniform::new_inclusive(1, 6)
243 .sample_iter(rng)
244 .filter(|x| *x != 5)
245 .take(10)
246 }
247
248 let mut rng = crate::test::rng(211);
249 let mut count = 0;
250 for val in ten_dice_rolls_other_than_five(&mut rng) {
251 assert!((1..=6).contains(&val) && val != 5);
252 count += 1;
253 }
254 assert_eq!(count, 10);
255 }
256
257 #[test]
258 #[cfg(feature = "alloc")]
259 fn test_dist_string() {
260 use core::str;
261 use crate::distributions::{Alphanumeric, DistString, Standard};
262 let mut rng = crate::test::rng(213);
263
264 let s1 = Alphanumeric.sample_string(&mut rng, 20);
265 assert_eq!(s1.len(), 20);
266 assert_eq!(str::from_utf8(s1.as_bytes()), Ok(s1.as_str()));
267
268 let s2 = Standard.sample_string(&mut rng, 20);
269 assert_eq!(s2.chars().count(), 20);
270 assert_eq!(str::from_utf8(s2.as_bytes()), Ok(s2.as_str()));
271 }
272}
273