randomize/
bounded_rand.rs

1//! Types for bounded randomization.
2
3/// Allows sampling a `u32` number in `0 .. N`.
4#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
5pub struct BoundedRandU32 {
6  /// number of possible outputs. outputs will be in `0 .. count`
7  count: u32,
8  /// Multiplication threshold thing.
9  ///
10  /// <https://arxiv.org/abs/1805.10941>
11  threshold: u32,
12}
13impl BoundedRandU32 {
14  #[allow(missing_docs)]
15  pub const _4: Self = Self::new(4);
16  #[allow(missing_docs)]
17  pub const _6: Self = Self::new(6);
18  #[allow(missing_docs)]
19  pub const _8: Self = Self::new(8);
20  #[allow(missing_docs)]
21  pub const _10: Self = Self::new(10);
22  #[allow(missing_docs)]
23  pub const _12: Self = Self::new(12);
24  #[allow(missing_docs)]
25  pub const _20: Self = Self::new(20);
26
27  /// Constructs a new value.
28  ///
29  /// ## Panics
30  /// If the count is 0.
31  #[inline]
32  pub const fn new(count: u32) -> Self {
33    let threshold = count.wrapping_neg() % count;
34    Self { count, threshold }
35  }
36
37  /// Constructs a new value, or `None` on failure.
38  ///
39  /// ## Failure
40  /// If the count is 0.
41  #[inline]
42  pub const fn try_new(count: u32) -> Option<Self> {
43    if count > 0 {
44      Some(Self::new(count))
45    } else {
46      None
47    }
48  }
49
50  /// The number of possible outputs.
51  #[inline]
52  pub const fn count(self) -> u32 {
53    self.count
54  }
55
56  /// Given a `u32`, try to place it into this bounded range.
57  ///
58  /// ## Failure
59  /// * If the value is such that it doesn't fit evenly it is rejected.
60  #[inline]
61  pub const fn place_in_range(self, val: u32) -> Option<u32> {
62    let mul: u64 = (val as u64).wrapping_mul(self.count as u64);
63    let low_part: u32 = mul as u32;
64    if low_part < self.threshold {
65      None
66    } else {
67      debug_assert!(((mul >> 32) as u32) < self.count());
68      Some((mul >> 32) as u32)
69    }
70  }
71
72  /// Given a generator function, call it until
73  /// [`place_in_range`](Self::place_in_range) succeeds.
74  #[inline]
75  pub fn sample<F: FnMut() -> u32>(self, mut f: F) -> u32 {
76    loop {
77      if let Some(output) = self.place_in_range(f()) {
78        return output;
79      }
80    }
81  }
82}
83
84/// Allows sampling a `u16` number in `0 .. N`.
85///
86/// The primary advantage of this type over [BoundedRandU32] is that this type
87/// samples using only 32-bit multiplications rather than 64-bit
88/// multiplications, so on a 32-bit CPU this will perform much faster.
89#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
90pub struct BoundedRandU16 {
91  /// number of possible outputs. outputs will be in `0 .. count`
92  count: u16,
93  /// Multiplication threshold thing.
94  ///
95  /// <https://arxiv.org/abs/1805.10941>
96  threshold: u16,
97}
98impl BoundedRandU16 {
99  #[allow(missing_docs)]
100  pub const _4: Self = Self::new(4);
101  #[allow(missing_docs)]
102  pub const _6: Self = Self::new(6);
103  #[allow(missing_docs)]
104  pub const _8: Self = Self::new(8);
105  #[allow(missing_docs)]
106  pub const _10: Self = Self::new(10);
107  #[allow(missing_docs)]
108  pub const _12: Self = Self::new(12);
109  #[allow(missing_docs)]
110  pub const _20: Self = Self::new(20);
111
112  /// Constructs a new value.
113  ///
114  /// ## Panics
115  /// If the count is 0.
116  #[inline]
117  pub const fn new(count: u16) -> Self {
118    let threshold = count.wrapping_neg() % count;
119    Self { count, threshold }
120  }
121
122  /// Constructs a new value, or `None` on failure.
123  ///
124  /// ## Failure
125  /// If the count is 0.
126  #[inline]
127  pub const fn try_new(count: u16) -> Option<Self> {
128    if count > 0 {
129      Some(Self::new(count))
130    } else {
131      None
132    }
133  }
134
135  /// The number of possible outputs.
136  #[inline]
137  pub const fn count(self) -> u16 {
138    self.count
139  }
140
141  /// Given a `u16`, try to place it into this bounded range.
142  ///
143  /// ## Failure
144  /// * If the value is such that it doesn't fit evenly it is rejected.
145  #[inline]
146  pub const fn place_in_range(self, val: u16) -> Option<u16> {
147    let mul: u32 = (val as u32).wrapping_mul(self.count as u32);
148    let low_part: u16 = mul as u16;
149    if low_part < self.threshold {
150      None
151    } else {
152      debug_assert!(((mul >> 16) as u16) < self.count());
153      Some((mul >> 16) as u16)
154    }
155  }
156
157  /// Given a generator function, call it until
158  /// [`place_in_range`](Self::place_in_range) succeeds.
159  #[inline]
160  pub fn sample<F: FnMut() -> u16>(self, mut f: F) -> u16 {
161    loop {
162      if let Some(output) = self.place_in_range(f()) {
163        return output;
164      }
165    }
166  }
167}