rand/seq/
increasing_uniform.rs

1// Copyright 2018-2023 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
9use crate::{Rng, RngCore};
10
11/// Similar to a Uniform distribution,
12/// but after returning a number in the range [0,n], n is increased by 1.
13pub(crate) struct IncreasingUniform<R: RngCore> {
14    pub rng: R,
15    n: u32,
16    // Chunk is a random number in [0, (n + 1) * (n + 2) *..* (n + chunk_remaining) )
17    chunk: u32,
18    chunk_remaining: u8,
19}
20
21impl<R: RngCore> IncreasingUniform<R> {
22    /// Create a dice roller.
23    /// The next item returned will be a random number in the range [0,n]
24    pub fn new(rng: R, n: u32) -> Self {
25        // If n = 0, the first number returned will always be 0
26        // so we don't need to generate a random number
27        let chunk_remaining = if n == 0 { 1 } else { 0 };
28        Self {
29            rng,
30            n,
31            chunk: 0,
32            chunk_remaining,
33        }
34    }
35
36    /// Returns a number in [0,n] and increments n by 1.
37    /// Generates new random bits as needed
38    /// Panics if `n >= u32::MAX`
39    #[inline]
40    pub fn next_index(&mut self) -> usize {
41        let next_n = self.n + 1;
42
43        // There's room for further optimisation here:
44        // random_range uses rejection sampling (or other method; see #1196) to avoid bias.
45        // When the initial sample is biased for range 0..bound
46        // it may still be viable to use for a smaller bound
47        // (especially if small biases are considered acceptable).
48
49        let next_chunk_remaining = self.chunk_remaining.checked_sub(1).unwrap_or_else(|| {
50            // If the chunk is empty, generate a new chunk
51            let (bound, remaining) = calculate_bound_u32(next_n);
52            // bound = (n + 1) * (n + 2) *..* (n + remaining)
53            self.chunk = self.rng.random_range(..bound);
54            // Chunk is a random number in
55            // [0, (n + 1) * (n + 2) *..* (n + remaining) )
56
57            remaining - 1
58        });
59
60        let result = if next_chunk_remaining == 0 {
61            // `chunk` is a random number in the range [0..n+1)
62            // Because `chunk_remaining` is about to be set to zero
63            // we do not need to clear the chunk here
64            self.chunk as usize
65        } else {
66            // `chunk` is a random number in a range that is a multiple of n+1
67            // so r will be a random number in [0..n+1)
68            let r = self.chunk % next_n;
69            self.chunk /= next_n;
70            r as usize
71        };
72
73        self.chunk_remaining = next_chunk_remaining;
74        self.n = next_n;
75        result
76    }
77}
78
79#[inline]
80/// Calculates `bound`, `count` such that bound (m)*(m+1)*..*(m + remaining - 1)
81fn calculate_bound_u32(m: u32) -> (u32, u8) {
82    debug_assert!(m > 0);
83    #[inline]
84    const fn inner(m: u32) -> (u32, u8) {
85        let mut product = m;
86        let mut current = m + 1;
87
88        loop {
89            if let Some(p) = u32::checked_mul(product, current) {
90                product = p;
91                current += 1;
92            } else {
93                // Count has a maximum value of 13 for when min is 1 or 2
94                let count = (current - m) as u8;
95                return (product, count);
96            }
97        }
98    }
99
100    const RESULT2: (u32, u8) = inner(2);
101    if m == 2 {
102        // Making this value a constant instead of recalculating it
103        // gives a significant (~50%) performance boost for small shuffles
104        return RESULT2;
105    }
106
107    inner(m)
108}
rand/seq/increasing_uniform.rs

rand/seq/
increasing_uniform.rs
