rand/distr/
slice.rs

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

rand/distr/
slice.rs
