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;
13#[cfg(feature = "alloc")]
14use alloc::string::String;
15use core::iter;
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::distr::{Distribution, Alphanumeric, Uniform, StandardUniform};
52 ///
53 /// let mut rng = rand::rng();
54 ///
55 /// // Vec of 16 x f32:
56 /// let v: Vec<f32> = StandardUniform.sample_iter(&mut rng).take(16).collect();
57 ///
58 /// // String:
59 /// let s: String = Alphanumeric
60 /// .sample_iter(&mut rng)
61 /// .take(7)
62 /// .map(char::from)
63 /// .collect();
64 ///
65 /// // Dice-rolling:
66 /// let die_range = Uniform::new_inclusive(1, 6).unwrap();
67 /// let mut roll_die = die_range.sample_iter(&mut rng);
68 /// while roll_die.next().unwrap() != 6 {
69 /// println!("Not a 6; rolling again!");
70 /// }
71 /// ```
72 fn sample_iter<R>(self, rng: R) -> Iter<Self, R, T>
73 where
74 R: Rng,
75 Self: Sized,
76 {
77 Iter {
78 distr: self,
79 rng,
80 phantom: core::marker::PhantomData,
81 }
82 }
83
84 /// Map sampled values to type `S`
85 ///
86 /// # Example
87 ///
88 /// ```
89 /// use rand::distr::{Distribution, Uniform};
90 ///
91 /// let die = Uniform::new_inclusive(1, 6).unwrap();
92 /// let even_number = die.map(|num| num % 2 == 0);
93 /// while !even_number.sample(&mut rand::rng()) {
94 /// println!("Still odd; rolling again!");
95 /// }
96 /// ```
97 fn map<F, S>(self, func: F) -> Map<Self, F, T, S>
98 where
99 F: Fn(T) -> S,
100 Self: Sized,
101 {
102 Map {
103 distr: self,
104 func,
105 phantom: core::marker::PhantomData,
106 }
107 }
108}
109
110impl<T, D: Distribution<T> + ?Sized> Distribution<T> for &D {
111 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> T {
112 (*self).sample(rng)
113 }
114}
115
116/// An iterator over a [`Distribution`]
117///
118/// This iterator yields random values of type `T` with distribution `D`
119/// from a random generator of type `R`.
120///
121/// Construct this `struct` using [`Distribution::sample_iter`] or
122/// [`Rng::sample_iter`]. It is also used by [`Rng::random_iter`] and
123/// [`crate::random_iter`].
124#[derive(Debug)]
125pub struct Iter<D, R, T> {
126 distr: D,
127 rng: R,
128 phantom: core::marker::PhantomData<T>,
129}
130
131impl<D, R, T> Iterator for Iter<D, R, T>
132where
133 D: Distribution<T>,
134 R: Rng,
135{
136 type Item = T;
137
138 #[inline(always)]
139 fn next(&mut self) -> Option<T> {
140 // Here, self.rng may be a reference, but we must take &mut anyway.
141 // Even if sample could take an R: Rng by value, we would need to do this
142 // since Rng is not copyable and we cannot enforce that this is "reborrowable".
143 Some(self.distr.sample(&mut self.rng))
144 }
145
146 fn size_hint(&self) -> (usize, Option<usize>) {
147 (usize::MAX, None)
148 }
149}
150
151impl<D, R, T> iter::FusedIterator for Iter<D, R, T>
152where
153 D: Distribution<T>,
154 R: Rng,
155{
156}
157
158/// A [`Distribution`] which maps sampled values to type `S`
159///
160/// This `struct` is created by the [`Distribution::map`] method.
161/// See its documentation for more.
162#[derive(Debug)]
163pub struct Map<D, F, T, S> {
164 distr: D,
165 func: F,
166 phantom: core::marker::PhantomData<fn(T) -> S>,
167}
168
169impl<D, F, T, S> Distribution<S> for Map<D, F, T, S>
170where
171 D: Distribution<T>,
172 F: Fn(T) -> S,
173{
174 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> S {
175 (self.func)(self.distr.sample(rng))
176 }
177}
178
179/// Sample or extend a [`String`]
180///
181/// Helper methods to extend a [`String`] or sample a new [`String`].
182#[cfg(feature = "alloc")]
183pub trait SampleString {
184 /// Append `len` random chars to `string`
185 ///
186 /// Note: implementations may leave `string` with excess capacity. If this
187 /// is undesirable, consider calling [`String::shrink_to_fit`] after this
188 /// method.
189 fn append_string<R: Rng + ?Sized>(&self, rng: &mut R, string: &mut String, len: usize);
190
191 /// Generate a [`String`] of `len` random chars
192 ///
193 /// Note: implementations may leave the string with excess capacity. If this
194 /// is undesirable, consider calling [`String::shrink_to_fit`] after this
195 /// method.
196 #[inline]
197 fn sample_string<R: Rng + ?Sized>(&self, rng: &mut R, len: usize) -> String {
198 let mut s = String::new();
199 self.append_string(rng, &mut s, len);
200 s
201 }
202}
203
204#[cfg(test)]
205mod tests {
206 use crate::distr::{Distribution, Uniform};
207 use crate::Rng;
208
209 #[test]
210 fn test_distributions_iter() {
211 use crate::distr::Open01;
212 let mut rng = crate::test::rng(210);
213 let distr = Open01;
214 let mut iter = Distribution::<f32>::sample_iter(distr, &mut rng);
215 let mut sum: f32 = 0.;
216 for _ in 0..100 {
217 sum += iter.next().unwrap();
218 }
219 assert!(0. < sum && sum < 100.);
220 }
221
222 #[test]
223 fn test_distributions_map() {
224 let dist = Uniform::new_inclusive(0, 5).unwrap().map(|val| val + 15);
225
226 let mut rng = crate::test::rng(212);
227 let val = dist.sample(&mut rng);
228 assert!((15..=20).contains(&val));
229 }
230
231 #[test]
232 fn test_make_an_iter() {
233 fn ten_dice_rolls_other_than_five<R: Rng>(rng: &mut R) -> impl Iterator<Item = i32> + '_ {
234 Uniform::new_inclusive(1, 6)
235 .unwrap()
236 .sample_iter(rng)
237 .filter(|x| *x != 5)
238 .take(10)
239 }
240
241 let mut rng = crate::test::rng(211);
242 let mut count = 0;
243 for val in ten_dice_rolls_other_than_five(&mut rng) {
244 assert!((1..=6).contains(&val) && val != 5);
245 count += 1;
246 }
247 assert_eq!(count, 10);
248 }
249
250 #[test]
251 #[cfg(feature = "alloc")]
252 fn test_dist_string() {
253 use crate::distr::{Alphanumeric, SampleString, StandardUniform};
254 use core::str;
255 let mut rng = crate::test::rng(213);
256
257 let s1 = Alphanumeric.sample_string(&mut rng, 20);
258 assert_eq!(s1.len(), 20);
259 assert_eq!(str::from_utf8(s1.as_bytes()), Ok(s1.as_str()));
260
261 let s2 = StandardUniform.sample_string(&mut rng, 20);
262 assert_eq!(s2.chars().count(), 20);
263 assert_eq!(str::from_utf8(s2.as_bytes()), Ok(s2.as_str()));
264 }
265}
266