agave_random/
range.rs

1use {
2    rand::Rng,
3    std::{
4        num::NonZero,
5        ops::{Bound, RangeBounds},
6    },
7};
8
9/// Compatibility uniform range `[0, limit)` sampler for `u64` numbers
10///
11/// Sampler uses provided `rand::Rng` reference to generate random `u64` numbers and
12/// map them to desired range maintaining uniform distribution of generated numbers.
13///
14/// This utility exists only to provide compatibility of sampling algorithm with `rand`
15/// library at versions <=0.8.5, since parts of the system rely on reproducible sequence
16/// of numbers given stable seeded random number generator.
17///
18/// Two sampling algorithms are supported (they initialize internal `zone` value in different ways)
19/// to provide compatibility with two ways `rand` sampling can be performed:
20/// - `new_like_instance_sample`: reproduces values obtained from sampler instance created with
21///   `rand::distributions::uniform::UniformSampler::new` and then used by calling `sample`
22/// - `new_like_trait_sample`: reproduces values obtained from trait function calls
23///   `rand::distributions::uniform::UniformSampler::sample_single`
24#[derive(Debug)]
25pub struct UniformU64Sampler {
26    range_end: NonZero<u64>,
27    zone: u64,
28}
29
30impl UniformU64Sampler {
31    /// Create sampler reproducing `sample` calls on `UniformInt` instance
32    ///
33    /// The `zone` internal threshold is obtained by calculating modulo of `u64::MAX`'s difference
34    /// with provided range's end to itself. See:
35    /// https://github.com/rust-random/rand/blob/937320cbfeebd4352a23086d9c6e68f067f74644/src/distributions/uniform.rs#L458-L504
36    #[allow(clippy::arithmetic_side_effects)]
37    pub fn new_like_instance_sample(range_end: NonZero<u64>) -> Self {
38        let ints_to_reject = (u64::MAX - range_end.get() + 1) % range_end.get();
39        let zone = u64::MAX - ints_to_reject;
40        Self { range_end, zone }
41    }
42
43    /// Create sampler reproducing direct `sample_single` calls on `UniformInt` trait
44    ///
45    /// The `zone` internal threshold is obtained by calculating 2^{number of leading zeros in provided range}. See
46    /// https://github.com/rust-random/rand/blob/937320cbfeebd4352a23086d9c6e68f067f74644/src/distributions/uniform.rs#L534-L553
47    pub fn new_like_trait_sample(range_end: NonZero<u64>) -> Self {
48        let zone = (range_end.get() << range_end.leading_zeros()).wrapping_sub(1);
49        Self { range_end, zone }
50    }
51
52    /// Obtain random number from `rng` and map it to the initialized range of this sampler
53    pub fn sample(&self, rng: &mut impl Rng) -> u64 {
54        loop {
55            let (hi, lo) = Self::wmul(rng.random(), self.range_end);
56            if lo <= self.zone {
57                return hi;
58            }
59        }
60    }
61
62    #[allow(clippy::arithmetic_side_effects)]
63    fn wmul(x: u64, y: NonZero<u64>) -> (u64, u64) {
64        let tmp = (x as u128) * (y.get() as u128);
65        ((tmp >> 64) as u64, tmp as u64)
66    }
67}
68
69/// Sample a random number in `range` using `rng` as generator for random `u64` numbers
70///
71/// This utility exists only to provide compatibility of sampling algorithm with `rand`
72/// library at versions <=0.8.5, it is equivalent to its `rng.gen_range(range)`.
73///
74/// Panics: when `range` is empty.
75pub fn random_u64_range(rng: &mut impl Rng, range: impl RangeBounds<u64>) -> u64 {
76    let start = match range.start_bound() {
77        Bound::Unbounded => 0,
78        Bound::Included(start) => *start,
79        Bound::Excluded(&u64::MAX) => panic!("Cannot generate number in empty range (max..)"),
80        Bound::Excluded(start) => start.wrapping_add(1),
81    };
82    let last = match range.end_bound() {
83        Bound::Unbounded | Bound::Included(&u64::MAX) if start == 0 => return rng.random(),
84        Bound::Unbounded => u64::MAX,
85        Bound::Included(last) => *last,
86        Bound::Excluded(0) => panic!("Cannot generate number in empty range (..0)"),
87        Bound::Excluded(end) => end.wrapping_sub(1),
88    };
89    let zero_range_end = last
90        .checked_sub(start)
91        .expect("Range must not be empty")
92        .wrapping_add(1);
93    // last - start != u64::MAX after check calculating last above, so +1 won't overflow
94    let zero_range_end = NonZero::new(zero_range_end).unwrap();
95    let sampler = UniformU64Sampler::new_like_trait_sample(zero_range_end);
96    sampler.sample(rng).wrapping_add(start)
97}
98
99#[cfg(test)]
100mod tests {
101    use {
102        super::*,
103        rand::SeedableRng as _,
104        rand_chacha::ChaChaRng,
105        solana_sha256_hasher::Hasher,
106        std::{array, ops::Range},
107        test_case::test_case,
108    };
109
110    const CHACHA_SEED: [u8; 32] = [16; 32];
111
112    #[test]
113    fn test_uniform_sample_like_instance_sample_example() {
114        let mut rng_compat = ChaChaRng::from_seed(CHACHA_SEED);
115        let sampler_compat =
116            UniformU64Sampler::new_like_instance_sample(NonZero::new(294_533).unwrap());
117        let values: [u64; 10] = array::from_fn(|_| sampler_compat.sample(&mut rng_compat));
118        assert_eq!(
119            values,
120            [
121                280405, 7507, 84194, 272634, 52124, 190984, 8676, 230277, 223574, 126007
122            ]
123        );
124    }
125
126    #[test_case(10, "5p3DrG89DLrJotbMuZJUDHMrqcfoQiDQpa3FbNDbnND6")]
127    #[test_case(2_729, "J65XfpzkppuxWEyvgVdDmGBsJSex6BuevYuDSRtaDNAK")]
128    #[test_case(4_098, "EaRABrcPBw5LLNBmjn5ay5YN5yTPgv3GnQJLbpUarRkJ")]
129    #[test_case(504_302_479, "HniJPVe7zir8XxHmtUuxzhPVvWjMcJxVhhj8QSs1PZTC")]
130    #[test_case(1_000_346_000_000, "BghMy7yLe6BVzMbc4zNvAB4sz3ZSPYKJS5oLGvXYVzw2")]
131    fn test_uniform_sampler_like_instance_sample_compat(range_end: u64, expected_hash: &str) {
132        let mut rng_compat = ChaChaRng::from_seed(CHACHA_SEED);
133        let sampler_compat =
134            UniformU64Sampler::new_like_instance_sample(NonZero::new(range_end).unwrap());
135
136        let mut hash = Hasher::default();
137        (0..600_000).for_each(|_| {
138            let compat = sampler_compat.sample(&mut rng_compat);
139            hash.hash(&compat.to_le_bytes());
140        });
141        assert_eq!(hash.result().to_string(), expected_hash);
142    }
143
144    #[test]
145    fn test_uniform_sample_like_trait_sample_example() {
146        let mut rng_compat = ChaChaRng::from_seed(CHACHA_SEED);
147        let sampler_compat =
148            UniformU64Sampler::new_like_trait_sample(NonZero::new(294_533).unwrap());
149        let values: [u64; 10] = array::from_fn(|_| sampler_compat.sample(&mut rng_compat));
150        assert_eq!(
151            values,
152            [
153                272634, 52124, 8676, 230277, 223574, 137788, 212533, 213080, 187008, 209168
154            ]
155        );
156    }
157
158    #[test_case(10, "2tsA6RuXekKvMMvAqjkMMipjYySpw8mrrwwoBeoKohD1")]
159    #[test_case(1_000, "58SG5gnD5wR6ngrjxuTv7m1SEXiXejZvuJPUCsyrmqWx")]
160    #[test_case(4098, "7w4TY1oaeEeqHmPw3iobD8WVq1BsL5eo7ZuNdvDkoSQo")]
161    #[test_case(10_000_000_000, "ENe5A82Wq2nYsH17q9WkKAudXdLE9FVGRbySuuCpaR5k")]
162    fn test_uniform_sampler_like_trait_sample_single(range_end: u64, expected_hash: &str) {
163        let mut rng_compat = ChaChaRng::from_seed(CHACHA_SEED);
164        let sampler_compat =
165            UniformU64Sampler::new_like_trait_sample(NonZero::new(range_end).unwrap());
166
167        let mut hash = Hasher::default();
168        (0..1_000).for_each(|_| {
169            let compat = sampler_compat.sample(&mut rng_compat);
170            hash.hash(&compat.to_le_bytes());
171        });
172        assert_eq!(hash.result().to_string(), expected_hash);
173    }
174
175    #[test_case(0..100, &[95, 2, 28, 92, 17, 78, 72, 72, 21, 65])]
176    #[test_case(4..=100, &[96, 6, 31, 21, 77, 45, 49, 74, 24, 62])]
177    #[test_case(4..101, &[96, 6, 31, 21, 77, 45, 49, 74, 24, 62])]
178    #[test_case(77..=77, &[77, 77, 77, 77, 77, 77, 77, 77, 77, 77])]
179    #[test_case((u64::MAX - 1)..,
180                &[18446744073709551614, 18446744073709551614, 18446744073709551615, 18446744073709551614,
181                  18446744073709551615, 18446744073709551615, 18446744073709551614, 18446744073709551615,
182                  18446744073709551615, 18446744073709551615])]
183    #[test_case((u64::MAX - 1)..=u64::MAX,
184                &[18446744073709551614, 18446744073709551614, 18446744073709551615, 18446744073709551614,
185                  18446744073709551615, 18446744073709551615, 18446744073709551614, 18446744073709551615,
186                  18446744073709551615, 18446744073709551615])]
187    #[test_case(1_000..3_463, &[1704, 2597, 1072, 2152, 2777, 1526, 2619, 3318, 1824, 2749])]
188    #[test_case(.., &[17561962876765024395, 470213581521912695, 5273148429961811548, 17075204616871931816,
189                      3264552321098222474, 11961477208003522799, 543394791199312464, 14422393090876740590,
190                      14002568482199435789, 7891933133259177431])]
191    #[test_case(0..=u64::MAX,
192                &[17561962876765024395, 470213581521912695, 5273148429961811548, 17075204616871931816,
193                  3264552321098222474, 11961477208003522799, 543394791199312464, 14422393090876740590,
194                  14002568482199435789, 7891933133259177431])]
195    fn test_random_range_example(range: impl RangeBounds<u64> + Clone, expected: &[u64]) {
196        let mut rng_compat = ChaChaRng::from_seed(CHACHA_SEED);
197        let result: Vec<u64> = (0..10)
198            .map(|_| random_u64_range(&mut rng_compat, range.clone()))
199            .collect();
200        assert_eq!(result, expected);
201    }
202
203    #[test_case(0..100, "9qrc9Atiwy7f5vYs6UK4fk4WMtRqjLF8Zd8WsG7MFsYX")]
204    #[test_case(1_000..9_999, "6qRfSYk55bPZty3qc3jLaSxASXgDRLaw8AWnXKEhiW6M")]
205    #[test_case(4..4098, "7rKu65HwouY8sr3o7jZrdRPCtGwVXjZDTXsjcKg3VVr")]
206    #[test_case(1_000..20_000_000_000, "J5vwi9DrgXrTNinMQmDw1zSq6ogeUencayoqEn4r64gH")]
207    fn test_random_range_reproducibility(range: Range<u64>, expected_hash: &str) {
208        let mut rng_compat = ChaChaRng::from_seed(CHACHA_SEED);
209
210        let mut hash = Hasher::default();
211        (0..10_000).for_each(|_| {
212            let compat = random_u64_range(&mut rng_compat, range.clone());
213            hash.hash(&compat.to_le_bytes());
214        });
215        assert_eq!(hash.result().to_string(), expected_hash);
216    }
217
218    #[test]
219    #[should_panic]
220    fn test_random_range_panic_empty() {
221        let mut rng = ChaChaRng::from_seed(CHACHA_SEED);
222        let _ = random_u64_range(&mut rng, 100..100);
223    }
224}
agave_random/range.rs

agave_random/
range.rs
