1// Copyright 2021 Developers of the Rand project.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9//! Distributions over slices
10
11use core::num::NonZeroUsize;
12
13use crate::distr::uniform::{UniformSampler, UniformUsize};
14use crate::distr::Distribution;
15#[cfg(feature = "alloc")]
16use alloc::string::String;
17
18/// A distribution to uniformly sample elements of a slice
19///
20/// Like [`IndexedRandom::choose`], this uniformly samples elements of a slice
21/// without modification of the slice (so called "sampling with replacement").
22/// This distribution object may be a little faster for repeated sampling (but
23/// slower for small numbers of samples).
24///
25/// ## Examples
26///
27/// Since this is a distribution, [`Rng::sample_iter`] and
28/// [`Distribution::sample_iter`] may be used, for example:
29/// ```
30/// use rand::distr::{Distribution, slice::Choose};
31///
32/// let vowels = ['a', 'e', 'i', 'o', 'u'];
33/// let vowels_dist = Choose::new(&vowels).unwrap();
34///
35/// // build a string of 10 vowels
36/// let vowel_string: String = vowels_dist
37/// .sample_iter(&mut rand::rng())
38/// .take(10)
39/// .collect();
40///
41/// println!("{}", vowel_string);
42/// assert_eq!(vowel_string.len(), 10);
43/// assert!(vowel_string.chars().all(|c| vowels.contains(&c)));
44/// ```
45///
46/// For a single sample, [`IndexedRandom::choose`] may be preferred:
47/// ```
48/// use rand::seq::IndexedRandom;
49///
50/// let vowels = ['a', 'e', 'i', 'o', 'u'];
51/// let mut rng = rand::rng();
52///
53/// println!("{}", vowels.choose(&mut rng).unwrap());
54/// ```
55///
56/// [`IndexedRandom::choose`]: crate::seq::IndexedRandom::choose
57/// [`Rng::sample_iter`]: crate::Rng::sample_iter
58#[derive(Debug, Clone, Copy)]
59pub struct Choose<'a, T> {
60 slice: &'a [T],
61 range: UniformUsize,
62 num_choices: NonZeroUsize,
63}
64
65impl<'a, T> Choose<'a, T> {
66 /// Create a new `Choose` instance which samples uniformly from the slice.
67 ///
68 /// Returns error [`Empty`] if the slice is empty.
69 pub fn new(slice: &'a [T]) -> Result<Self, Empty> {
70 let num_choices: ! = NonZeroUsize::new(slice.len()).ok_or(Empty)?;
71
72 Ok(Self {
73 slice,
74 range: UniformUsize::new(low:0, high:num_choices.get()).unwrap(),
75 num_choices,
76 })
77 }
78
79 /// Returns the count of choices in this distribution
80 pub fn num_choices(&self) -> NonZeroUsize {
81 self.num_choices
82 }
83}
84
85impl<'a, T> Distribution<&'a T> for Choose<'a, T> {
86 fn sample<R: crate::Rng + ?Sized>(&self, rng: &mut R) -> &'a T {
87 let idx: usize = self.range.sample(rng);
88
89 debug_assert!(
90 idx < self.slice.len(),
91 "Uniform::new(0, {}) somehow returned {}",
92 self.slice.len(),
93 idx
94 );
95
96 // Safety: at construction time, it was ensured that the slice was
97 // non-empty, and that the `Uniform` range produces values in range
98 // for the slice
99 unsafe { self.slice.get_unchecked(idx) }
100 }
101}
102
103/// Error: empty slice
104///
105/// This error is returned when [`Choose::new`] is given an empty slice.
106#[derive(Debug, Clone, Copy)]
107pub struct Empty;
108
109impl core::fmt::Display for Empty {
110 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
111 write!(
112 f,
113 "Tried to create a `rand::distr::slice::Choose` with an empty slice"
114 )
115 }
116}
117
118#[cfg(feature = "std")]
119impl std::error::Error for Empty {}
120
121#[cfg(feature = "alloc")]
122impl super::SampleString for Choose<'_, char> {
123 fn append_string<R: crate::Rng + ?Sized>(&self, rng: &mut R, string: &mut String, len: usize) {
124 // Get the max char length to minimize extra space.
125 // Limit this check to avoid searching for long slice.
126 let max_char_len = if self.slice.len() < 200 {
127 self.slice
128 .iter()
129 .try_fold(1, |max_len, char| {
130 // When the current max_len is 4, the result max_char_len will be 4.
131 Some(max_len.max(char.len_utf8())).filter(|len| *len < 4)
132 })
133 .unwrap_or(4)
134 } else {
135 4
136 };
137
138 // Split the extension of string to reuse the unused capacities.
139 // Skip the split for small length or only ascii slice.
140 let mut extend_len = if max_char_len == 1 || len < 100 {
141 len
142 } else {
143 len / 4
144 };
145 let mut remain_len = len;
146 while extend_len > 0 {
147 string.reserve(max_char_len * extend_len);
148 string.extend(self.sample_iter(&mut *rng).take(extend_len));
149 remain_len -= extend_len;
150 extend_len = extend_len.min(remain_len);
151 }
152 }
153}
154
155#[cfg(test)]
156mod test {
157 use super::*;
158 use core::iter;
159
160 #[test]
161 fn value_stability() {
162 let rng = crate::test::rng(651);
163 let slice = Choose::new(b"escaped emus explore extensively").unwrap();
164 let expected = b"eaxee";
165 assert!(iter::zip(slice.sample_iter(rng), expected).all(|(a, b)| a == b));
166 }
167}
168