agave_random/
weighted.rs

1use {
2    crate::range::UniformU64Sampler,
3    rand::{Rng, distr::weighted::Error},
4    std::num::NonZero,
5};
6
7/// Compatibility weighted sampler for `u64` numbers
8///
9/// Sampler uses provided `rand::Rng` reference to generate random `u64` numbers and
10/// map them to the index of the weight from weights vector provided on initialization.
11///
12/// This utility exists only to provide compatibility of sampling algorithm with `rand`
13/// library at versions <=0.8.5, since parts of the system rely on reproducible sequence
14/// of numbers given stable seeded random number generator.
15///
16/// The algorithm reproduces indices returned by `rand::distributions::WeightedIndex`
17/// for the same weights vector and seeded random number generator.
18#[derive(Debug)]
19pub struct WeightedU64Index {
20    weights: Vec<u64>,
21    total_weight_sampler: UniformU64Sampler,
22}
23
24impl WeightedU64Index {
25    pub fn new(mut weights: Vec<u64>) -> Result<Self, Error> {
26        // Calculate prefix sum of weights such that binary search can find the index of the
27        // chosen weight.
28        let mut total_weight = 0u64;
29        for weight in weights.iter_mut() {
30            total_weight = total_weight.checked_add(*weight).ok_or(Error::Overflow)?;
31            *weight = total_weight;
32        }
33        if weights.pop().is_none() {
34            return Err(Error::InvalidInput);
35        }
36        let Some(total_weight) = NonZero::new(total_weight) else {
37            return Err(Error::InsufficientNonZero);
38        };
39
40        Ok(Self {
41            weights,
42            total_weight_sampler: UniformU64Sampler::new_like_instance_sample(total_weight),
43        })
44    }
45
46    pub fn sample(&self, rng: &mut impl Rng) -> usize {
47        let chosen_weight = self.total_weight_sampler.sample(rng);
48        // Find the first item which has a weight *higher* than the chosen weight.
49        self.weights.partition_point(|w| *w <= chosen_weight)
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use {
56        super::*, assert_matches::assert_matches, rand::SeedableRng as _, rand_chacha::ChaChaRng,
57        solana_sha256_hasher::Hasher, std::array, test_case::test_case,
58    };
59
60    const CHACHA_SEED: [u8; 32] = [16; 32];
61
62    #[test_case(100, 0, [95, 2, 28, 92, 17, 64, 2, 78, 75, 42])]
63    #[test_case(1_000, 0, [952, 25, 285, 925, 176, 648, 29, 781, 759, 427])]
64    #[test_case(1_000, 1, [975, 160, 534, 962, 420, 805, 172, 884, 871, 654])]
65    #[test_case(1_000, 2, [983, 294, 658, 974, 561, 865, 309, 921, 912, 753])]
66    #[test_case(10_000, 1, [9757, 1596, 5346, 9621, 4207, 8052, 1716, 8842, 8712, 6540])]
67    fn test_weighted_u64_index_example(num_weights: u64, pow: u32, expected_indices: [usize; 10]) {
68        let weights: Vec<_> = (0..num_weights).map(|i| i.pow(pow)).collect();
69
70        let mut rng_compat = ChaChaRng::from_seed(CHACHA_SEED);
71        let index_compat =
72            WeightedU64Index::new(weights.clone()).expect("non empty and non zero is ok");
73
74        let indices = array::from_fn(|_| index_compat.sample(&mut rng_compat));
75        assert_eq!(indices, expected_indices);
76    }
77
78    #[test_case(30_000, 0, 50_000, "23K23NXJpui3d9nrKLNfvwpHHRs4dxFfZ8saxH6ZPJyw")]
79    #[test_case(20_000, 1, 45_000, "9yvuyu8JDQtUo7cvWJKkc3cUmzWw5RzBJoEtoVR2N2r2")]
80    #[test_case(10_000, 2, 35_000, "842FcJe1kmnmrZXAA3rETtBakk1jFdz1dnzMbkf875gh")]
81    #[test_case(10_000, 3, 30_000, "5LNbaEBQrb5CzsoHdK79XNDENAJ9WJqW9LpWktqkRchf")]
82    fn test_weighted_u64_index_compat(num_weights: u64, pow: u32, len: usize, expected_hash: &str) {
83        let weights: Vec<_> = (0..num_weights).map(|i| i.pow(pow)).collect();
84
85        let mut rng_compat = ChaChaRng::from_seed(CHACHA_SEED);
86        let index_compat = WeightedU64Index::new(weights).expect("non empty and non zero is ok");
87
88        let mut hash = Hasher::default();
89        (0..len).for_each(|_| {
90            let compat = index_compat.sample(&mut rng_compat);
91            hash.hash(&compat.to_le_bytes());
92        });
93        assert_eq!(hash.result().to_string(), expected_hash);
94    }
95
96    #[test]
97    fn test_weighted_u64_index_error_on_new() {
98        assert_matches!(WeightedU64Index::new(vec![]), Err(Error::InvalidInput));
99        assert_matches!(
100            WeightedU64Index::new(vec![0, 0, 0, 0, 0]),
101            Err(Error::InsufficientNonZero)
102        );
103        assert_matches!(
104            WeightedU64Index::new(vec![u64::MAX / 3, u64::MAX / 2, 0, u64::MAX / 3]),
105            Err(Error::Overflow)
106        );
107    }
108}
agave_random/weighted.rs

agave_random/
weighted.rs
