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}
109