/rust/registry/src/index.crates.io-1949cf8c6b5b557f/tower-0.5.2/src/util/rng.rs
Line | Count | Source |
1 | | //! [PRNG] utilities for tower middleware. |
2 | | //! |
3 | | //! This module provides a generic [`Rng`] trait and a [`HasherRng`] that |
4 | | //! implements the trait based on [`RandomState`] or any other [`Hasher`]. |
5 | | //! |
6 | | //! These utilities replace tower's internal usage of `rand` with these smaller, |
7 | | //! more lightweight methods. Most of the implementations are extracted from |
8 | | //! their corresponding `rand` implementations. |
9 | | //! |
10 | | //! [PRNG]: https://en.wikipedia.org/wiki/Pseudorandom_number_generator |
11 | | |
12 | | use std::{ |
13 | | collections::hash_map::RandomState, |
14 | | hash::{BuildHasher, Hasher}, |
15 | | ops::Range, |
16 | | }; |
17 | | |
18 | | /// A simple [PRNG] trait for use within tower middleware. |
19 | | /// |
20 | | /// [PRNG]: https://en.wikipedia.org/wiki/Pseudorandom_number_generator |
21 | | pub trait Rng { |
22 | | /// Generate a random [`u64`]. |
23 | | fn next_u64(&mut self) -> u64; |
24 | | |
25 | | /// Generate a random [`f64`] between `[0, 1)`. |
26 | 0 | fn next_f64(&mut self) -> f64 { |
27 | | // Borrowed from: |
28 | | // https://github.com/rust-random/rand/blob/master/src/distributions/float.rs#L106 |
29 | 0 | let float_size = std::mem::size_of::<f64>() as u32 * 8; |
30 | 0 | let precision = 52 + 1; |
31 | 0 | let scale = 1.0 / ((1u64 << precision) as f64); |
32 | | |
33 | 0 | let value = self.next_u64(); |
34 | 0 | let value = value >> (float_size - precision); |
35 | | |
36 | 0 | scale * value as f64 |
37 | 0 | } |
38 | | |
39 | | /// Randomly pick a value within the range. |
40 | | /// |
41 | | /// # Panic |
42 | | /// |
43 | | /// - If start < end this will panic in debug mode. |
44 | 0 | fn next_range(&mut self, range: Range<u64>) -> u64 { |
45 | 0 | debug_assert!( |
46 | 0 | range.start < range.end, |
47 | | "The range start must be smaller than the end" |
48 | | ); |
49 | 0 | let start = range.start; |
50 | 0 | let end = range.end; |
51 | | |
52 | 0 | let range = end - start; |
53 | | |
54 | 0 | let n = self.next_u64(); |
55 | | |
56 | 0 | (n % range) + start |
57 | 0 | } |
58 | | } |
59 | | |
60 | | impl<R: Rng + ?Sized> Rng for Box<R> { |
61 | 0 | fn next_u64(&mut self) -> u64 { |
62 | 0 | (**self).next_u64() |
63 | 0 | } |
64 | | } |
65 | | |
66 | | /// A [`Rng`] implementation that uses a [`Hasher`] to generate the random |
67 | | /// values. The implementation uses an internal counter to pass to the hasher |
68 | | /// for each iteration of [`Rng::next_u64`]. |
69 | | /// |
70 | | /// # Default |
71 | | /// |
72 | | /// This hasher has a default type of [`RandomState`] which just uses the |
73 | | /// libstd method of getting a random u64. |
74 | | #[derive(Clone, Debug)] |
75 | | pub struct HasherRng<H = RandomState> { |
76 | | hasher: H, |
77 | | counter: u64, |
78 | | } |
79 | | |
80 | | impl HasherRng { |
81 | | /// Create a new default [`HasherRng`]. |
82 | 0 | pub fn new() -> Self { |
83 | 0 | HasherRng::default() |
84 | 0 | } |
85 | | } |
86 | | |
87 | | impl Default for HasherRng { |
88 | 0 | fn default() -> Self { |
89 | 0 | HasherRng::with_hasher(RandomState::default()) |
90 | 0 | } |
91 | | } |
92 | | |
93 | | impl<H> HasherRng<H> { |
94 | | /// Create a new [`HasherRng`] with the provided hasher. |
95 | 0 | pub fn with_hasher(hasher: H) -> Self { |
96 | 0 | HasherRng { hasher, counter: 0 } |
97 | 0 | } |
98 | | } |
99 | | |
100 | | impl<H> Rng for HasherRng<H> |
101 | | where |
102 | | H: BuildHasher, |
103 | | { |
104 | 0 | fn next_u64(&mut self) -> u64 { |
105 | 0 | let mut hasher = self.hasher.build_hasher(); |
106 | 0 | hasher.write_u64(self.counter); |
107 | 0 | self.counter = self.counter.wrapping_add(1); |
108 | 0 | hasher.finish() |
109 | 0 | } |
110 | | } |
111 | | |
112 | | /// A sampler modified from the Rand implementation for use internally for the balance middleware. |
113 | | /// |
114 | | /// It's an implementation of Floyd's combination algorithm with amount fixed at 2. This uses no allocated |
115 | | /// memory and finishes in constant time (only 2 random calls). |
116 | | /// |
117 | | /// ref: This was borrowed and modified from the following Rand implementation |
118 | | /// https://github.com/rust-random/rand/blob/b73640705d6714509f8ceccc49e8df996fa19f51/src/seq/index.rs#L375-L411 |
119 | | #[cfg(feature = "balance")] |
120 | | pub(crate) fn sample_floyd2<R: Rng>(rng: &mut R, length: u64) -> [u64; 2] { |
121 | | debug_assert!(2 <= length); |
122 | | let aidx = rng.next_range(0..length - 1); |
123 | | let bidx = rng.next_range(0..length); |
124 | | let aidx = if aidx == bidx { length - 1 } else { aidx }; |
125 | | [aidx, bidx] |
126 | | } |
127 | | |
128 | | #[cfg(test)] |
129 | | mod tests { |
130 | | use super::*; |
131 | | use quickcheck::*; |
132 | | |
133 | | quickcheck! { |
134 | | fn next_f64(counter: u64) -> TestResult { |
135 | | let mut rng = HasherRng::default(); |
136 | | rng.counter = counter; |
137 | | let n = rng.next_f64(); |
138 | | |
139 | | TestResult::from_bool(n < 1.0 && n >= 0.0) |
140 | | } |
141 | | |
142 | | fn next_range(counter: u64, range: Range<u64>) -> TestResult { |
143 | | if range.start >= range.end{ |
144 | | return TestResult::discard(); |
145 | | } |
146 | | |
147 | | let mut rng = HasherRng::default(); |
148 | | rng.counter = counter; |
149 | | |
150 | | let n = rng.next_range(range.clone()); |
151 | | |
152 | | TestResult::from_bool(n >= range.start && (n < range.end || range.start == range.end)) |
153 | | } |
154 | | |
155 | | fn sample_floyd2(counter: u64, length: u64) -> TestResult { |
156 | | if length < 2 || length > 256 { |
157 | | return TestResult::discard(); |
158 | | } |
159 | | |
160 | | let mut rng = HasherRng::default(); |
161 | | rng.counter = counter; |
162 | | |
163 | | let [a, b] = super::sample_floyd2(&mut rng, length); |
164 | | |
165 | | if a >= length || b >= length || a == b { |
166 | | return TestResult::failed(); |
167 | | } |
168 | | |
169 | | TestResult::passed() |
170 | | } |
171 | | } |
172 | | |
173 | | #[test] |
174 | | fn sample_inplace_boundaries() { |
175 | | let mut r = HasherRng::default(); |
176 | | match super::sample_floyd2(&mut r, 2) { |
177 | | [0, 1] | [1, 0] => (), |
178 | | array => panic!("unexpected inplace boundaries: {:?}", array), |
179 | | } |
180 | | } |
181 | | } |