hotshot/traits/election/
helpers.rs

1// Copyright (c) 2021-2024 Espresso Systems (espressosys.com)
2// This file is part of the HotShot repository.
3
4// You should have received a copy of the MIT License
5// along with the HotShot repository. If not, see <https://mit-license.org/>.
6
7use std::{collections::BTreeSet, hash::Hash};
8
9use rand::{rngs::StdRng, Rng, SeedableRng};
10
11/// Helper which allows producing random numbers within a range and preventing duplicates
12/// If consumed as a regular iterator, will return a randomly ordered permutation of all
13/// values from 0..max
14struct NonRepeatValueIterator {
15    /// Random number generator to use
16    rng: StdRng,
17
18    /// Values which have already been emitted, to avoid duplicates
19    values: BTreeSet<u64>,
20
21    /// Maximum value, open-ended. Numbers returned will be 0..max
22    max: u64,
23}
24
25impl NonRepeatValueIterator {
26    /// Create a new NonRepeatValueIterator
27    pub fn new(rng: StdRng, max: u64) -> Self {
28        Self {
29            rng,
30            values: BTreeSet::new(),
31            max,
32        }
33    }
34}
35
36impl Iterator for NonRepeatValueIterator {
37    type Item = u64;
38
39    fn next(&mut self) -> Option<Self::Item> {
40        if self.values.len() as u64 >= self.max {
41            return None;
42        }
43
44        loop {
45            let v = self.rng.gen_range(0..self.max);
46            if !self.values.contains(&v) {
47                self.values.insert(v);
48                return Some(v);
49            }
50        }
51    }
52}
53
54/// Create a single u64 seed by merging two u64s. Done this way to allow easy seeding of the number generator
55/// from both a stable SOUND as well as a moving value ROUND (typically, epoch). Shift left by 8 to avoid
56/// scenarios where someone manually stepping seeds would pass over the same space of random numbers across
57/// sequential rounds. Doesn't have to be 8, but has to be large enough that it is unlikely that a given
58/// test run will collide; using 8 means that 256 rounds (epochs) would have to happen inside of a test before
59/// the test starts repeating values from SEED+1.
60fn make_seed(seed: u64, round: u64) -> u64 {
61    seed.wrapping_add(round.wrapping_shl(8))
62}
63
64/// Create a pair of PRNGs for the given SEED and ROUND. Prev_rng is the PRNG for the previous ROUND, used to
65/// deterministically replay random numbers generated for the previous ROUND.
66fn make_rngs(seed: u64, round: u64) -> (StdRng, StdRng) {
67    let prev_rng = SeedableRng::seed_from_u64(make_seed(seed, round.wrapping_sub(1)));
68    let this_rng = SeedableRng::seed_from_u64(make_seed(seed, round));
69
70    (prev_rng, this_rng)
71}
72
73/// Iterator which returns odd/even values for a given COUNT of nodes. For OVERLAP=0, this will return
74/// [0, 2, 4, 6, ...] for an even round, and [1, 3, 5, 7, ...] for an odd round. Setting OVERLAP>0 will
75/// randomly introduce OVERLAP elements from the previous round, so an even round with OVERLAP=2 will contain
76/// something like [1, 7, 2, 4, 0, ...]. Note that the total number of nodes will always be COUNT/2, so
77/// for OVERLAP>0 a random number of nodes which would have been in the round for OVERLAP=0 will be dropped.
78/// Ordering of nodes is random. Outputs is deterministic when prev_rng and this_rng are provided by make_rngs
79/// using the same values for SEED and ROUND.
80pub struct StableQuorumIterator {
81    /// PRNG from the previous round
82    prev_rng: NonRepeatValueIterator,
83
84    /// PRNG for the current round
85    this_rng: NonRepeatValueIterator,
86
87    /// Current ROUND
88    round: u64,
89
90    /// Count of nodes in the source quorum being filtered against
91    count: u64,
92
93    /// OVERLAP of nodes to be carried over from the previous round
94    overlap: u64,
95
96    /// The next call to next() will emit the value with this index. Starts at 0 and is incremented for each
97    /// call to next()
98    index: u64,
99}
100
101/// Determines how many possible values can be made for the given odd/even
102/// E.g. if count is 5, then possible values would be [0, 1, 2, 3, 4]
103/// if odd = true, slots = 2 (1 or 3), else slots = 3 (0, 2, 4)
104fn calc_num_slots(count: u64, odd: bool) -> u64 {
105    (count / 2) + if odd { 0 } else { count % 2 }
106}
107
108impl StableQuorumIterator {
109    #[must_use]
110    /// Create a new StableQuorumIterator
111    ///
112    /// # Panics
113    ///
114    /// panics if overlap is greater than half of count
115    pub fn new(seed: u64, round: u64, count: u64, overlap: u64) -> Self {
116        assert!(
117            count / 2 > overlap,
118            "Overlap cannot be greater than the entire set size"
119        );
120
121        let (prev_rng, this_rng) = make_rngs(seed, round);
122
123        Self {
124            prev_rng: NonRepeatValueIterator::new(prev_rng, calc_num_slots(count, round % 2 == 0)),
125            this_rng: NonRepeatValueIterator::new(this_rng, calc_num_slots(count, round % 2 == 1)),
126            round,
127            count,
128            overlap,
129            index: 0,
130        }
131    }
132}
133
134impl Iterator for StableQuorumIterator {
135    type Item = u64;
136
137    fn next(&mut self) -> Option<Self::Item> {
138        if self.index >= (self.count / 2) {
139            // Always return exactly half of the possible values. If we have OVERLAP>0 then
140            // we need to return (COUNT/2)-OVERLAP of the current set, even if there are additional
141            // even (or odd) numbers that we can return.
142            None
143        } else if self.index < self.overlap {
144            // Generate enough values for the previous round. If the current round is odd, then
145            // we want to pick even values that were selected from the previous round to create OVERLAP
146            // even values.
147            let v = self.prev_rng.next().unwrap();
148            self.index += 1;
149            Some(v * 2 + (1 - self.round % 2))
150        } else {
151            // Generate new values. If our current round is odd, we'll be creating (COUNT/2)-OVERLAP
152            // odd values here.
153            let v = self.this_rng.next().unwrap();
154            self.index += 1;
155            Some(v * 2 + self.round % 2)
156        }
157    }
158}
159
160#[must_use]
161/// Helper function to convert the arguments to a StableQuorumIterator into an ordered set of values.
162///
163/// # Panics
164///
165/// panics if the arguments are invalid for StableQuorumIterator::new
166pub fn stable_quorum_filter(seed: u64, round: u64, count: usize, overlap: u64) -> BTreeSet<usize> {
167    StableQuorumIterator::new(seed, round, count as u64, overlap)
168        // We should never have more than u32_max members in a test
169        .map(|x| usize::try_from(x).unwrap())
170        .collect()
171}
172
173/// Constructs a quorum with a random number of members and overlaps. Functions similar to StableQuorumIterator,
174/// except that the number of MEMBERS and OVERLAP are also (deterministically) random, to allow additional variance
175/// in testing.
176pub struct RandomOverlapQuorumIterator {
177    /// PRNG from the previous round
178    prev_rng: NonRepeatValueIterator,
179
180    /// PRNG for the current round
181    this_rng: NonRepeatValueIterator,
182
183    /// Current ROUND
184    round: u64,
185
186    /// Number of members to emit for the current round
187    members: u64,
188
189    /// OVERLAP of nodes to be carried over from the previous round
190    overlap: u64,
191
192    /// The next call to next() will emit the value with this index. Starts at 0 and is incremented for each
193    /// call to next()
194    index: u64,
195}
196
197impl RandomOverlapQuorumIterator {
198    #[must_use]
199    /// Create a new RandomOverlapQuorumIterator
200    ///
201    /// # Panics
202    ///
203    /// panics if overlap and members can produce invalid results or if ranges are invalid
204    pub fn new(
205        seed: u64,
206        round: u64,
207        count: u64,
208        members_min: u64,
209        members_max: u64,
210        overlap_min: u64,
211        overlap_max: u64,
212    ) -> Self {
213        assert!(
214            members_min <= members_max,
215            "Members_min cannot be greater than members_max"
216        );
217        assert!(
218            overlap_min <= overlap_max,
219            "Overlap_min cannot be greater than overlap_max"
220        );
221        assert!(
222            overlap_max < members_min,
223            "Overlap_max must be less than members_min"
224        );
225        assert!(
226            count / 2 > overlap_max,
227            "Overlap cannot be greater than the entire set size"
228        );
229        assert!(
230            count / 2 >= members_max - overlap_min,
231            "members_max must be greater or equal to half of the count plus overlap_min"
232        );
233
234        let (mut prev_rng, mut this_rng) = make_rngs(seed, round);
235
236        // Consume two values from prev_rng to advance it to the same state it was at the beginning of the previous round
237        let _prev_members = prev_rng.gen_range(members_min..=members_max);
238        let _prev_overlap = prev_rng.gen_range(overlap_min..=overlap_max);
239        let this_members = this_rng.gen_range(members_min..=members_max);
240        let this_overlap = this_rng.gen_range(overlap_min..=overlap_max);
241
242        Self {
243            prev_rng: NonRepeatValueIterator::new(prev_rng, calc_num_slots(count, round % 2 == 0)),
244            this_rng: NonRepeatValueIterator::new(this_rng, calc_num_slots(count, round % 2 == 1)),
245            round,
246            members: this_members,
247            overlap: this_overlap,
248            index: 0,
249        }
250    }
251}
252
253impl Iterator for RandomOverlapQuorumIterator {
254    type Item = u64;
255
256    fn next(&mut self) -> Option<Self::Item> {
257        if self.index >= self.members {
258            None
259        } else if self.index < self.overlap {
260            // Generate enough values for the previous round
261            let v = self.prev_rng.next().unwrap();
262            self.index += 1;
263            Some(v * 2 + (1 - self.round % 2))
264        } else {
265            // Generate new values
266            let v = self.this_rng.next().unwrap();
267            self.index += 1;
268            Some(v * 2 + self.round % 2)
269        }
270    }
271}
272
273#[must_use]
274/// Helper function to convert the arguments to a StableQuorumIterator into an ordered set of values.
275///
276/// # Panics
277///
278/// panics if the arguments are invalid for RandomOverlapQuorumIterator::new
279pub fn random_overlap_quorum_filter(
280    seed: u64,
281    round: u64,
282    count: usize,
283    members_min: u64,
284    members_max: u64,
285    overlap_min: u64,
286    overlap_max: u64,
287) -> BTreeSet<usize> {
288    RandomOverlapQuorumIterator::new(
289        seed,
290        round,
291        count as u64,
292        members_min,
293        members_max,
294        overlap_min,
295        overlap_max,
296    )
297    // We should never have more than u32_max members in a test
298    .map(|x| usize::try_from(x).unwrap())
299    .collect()
300}
301
302/// Trait wrapping a config for quorum filters. This allows selection between either the StableQuorumIterator or the
303/// RandomOverlapQuorumIterator functionality from above
304pub trait QuorumFilterConfig:
305    Copy
306    + Clone
307    + std::fmt::Debug
308    + Default
309    + Send
310    + Sync
311    + Ord
312    + PartialOrd
313    + Eq
314    + PartialEq
315    + Hash
316    + 'static
317{
318    /// Called to run the filter and return a set of indices
319    fn execute(epoch: u64, count: usize) -> BTreeSet<usize>;
320}
321
322#[derive(Debug, Copy, Clone, Default, Eq, PartialEq, Hash, Ord, PartialOrd)]
323/// Provides parameters to use the StableQuorumIterator
324pub struct StableQuorumFilterConfig<const SEED: u64, const OVERLAP: u64> {}
325
326impl<const SEED: u64, const OVERLAP: u64> QuorumFilterConfig
327    for StableQuorumFilterConfig<SEED, OVERLAP>
328{
329    fn execute(epoch: u64, count: usize) -> BTreeSet<usize> {
330        stable_quorum_filter(SEED, epoch, count, OVERLAP)
331    }
332}
333
334#[derive(Debug, Copy, Clone, Default, Eq, PartialEq, Hash, Ord, PartialOrd)]
335/// Provides parameters to use the RandomOverlapQuorumIterator
336pub struct RandomOverlapQuorumFilterConfig<
337    const SEED: u64,
338    const MEMBERS_MIN: u64,
339    const MEMBERS_MAX: u64,
340    const OVERLAP_MIN: u64,
341    const OVERLAP_MAX: u64,
342> {}
343
344impl<
345        const SEED: u64,
346        const MEMBERS_MIN: u64,
347        const MEMBERS_MAX: u64,
348        const OVERLAP_MIN: u64,
349        const OVERLAP_MAX: u64,
350    > QuorumFilterConfig
351    for RandomOverlapQuorumFilterConfig<SEED, MEMBERS_MIN, MEMBERS_MAX, OVERLAP_MIN, OVERLAP_MAX>
352{
353    fn execute(epoch: u64, count: usize) -> BTreeSet<usize> {
354        random_overlap_quorum_filter(
355            SEED,
356            epoch,
357            count,
358            MEMBERS_MIN,
359            MEMBERS_MAX,
360            OVERLAP_MIN,
361            OVERLAP_MAX,
362        )
363    }
364}
365
366#[cfg(test)]
367mod tests {
368    use super::*;
369
370    #[test]
371    fn test_stable() {
372        for _ in 0..100 {
373            let seed = rand::random::<u64>();
374            let prev_set: Vec<u64> = StableQuorumIterator::new(seed, 1, 10, 2).collect();
375            let this_set: Vec<u64> = StableQuorumIterator::new(seed, 2, 10, 2).collect();
376
377            // The first two elements from prev_set are from its previous round. But its 2nd and 3rd elements
378            // are new, and should be carried over to become the first two elements from this_set.
379            assert_eq!(
380                prev_set[2..4],
381                this_set[0..2],
382                "prev_set={prev_set:?}, this_set={this_set:?}"
383            );
384        }
385    }
386
387    #[test]
388    fn test_random_overlap() {
389        for _ in 0..100 {
390            let seed = rand::random::<u64>();
391            let prev_set: Vec<u64> =
392                RandomOverlapQuorumIterator::new(seed, 1, 20, 5, 10, 2, 3).collect();
393            let this_set: Vec<u64> =
394                RandomOverlapQuorumIterator::new(seed, 2, 20, 5, 10, 2, 3).collect();
395
396            // Similar to the overlap before, but there are 4 possible cases: the previous set might have had
397            // either 2 or 3 overlaps, meaning we should start with index 2 or 3, and the overlap size might
398            // be either 2 or 3. We'll just check for 2 overlaps, meaning we have two possible overlap cases
399            // to verify.
400            let matched = (prev_set[2..4] == this_set[0..2]) || (prev_set[3..5] == this_set[0..2]);
401            assert!(matched, "prev_set={prev_set:?}, this_set={this_set:?}");
402        }
403    }
404
405    #[test]
406    fn test_odd_even() {
407        for _ in 0..100 {
408            let seed = rand::random::<u64>();
409
410            let odd_set: Vec<u64> = StableQuorumIterator::new(seed, 1, 10, 2).collect();
411            let even_set: Vec<u64> = StableQuorumIterator::new(seed, 2, 10, 2).collect();
412
413            assert!(
414                odd_set[2] % 2 == 1,
415                "odd set non-overlap value should be odd (stable)"
416            );
417            assert!(
418                even_set[2] % 2 == 0,
419                "even set non-overlap value should be even (stable)"
420            );
421
422            let odd_set: Vec<u64> =
423                RandomOverlapQuorumIterator::new(seed, 1, 20, 5, 10, 2, 3).collect();
424            let even_set: Vec<u64> =
425                RandomOverlapQuorumIterator::new(seed, 2, 20, 5, 10, 2, 3).collect();
426
427            assert!(
428                odd_set[3] % 2 == 1,
429                "odd set non-overlap value should be odd (random overlap)"
430            );
431            assert!(
432                even_set[3] % 2 == 0,
433                "even set non-overlap value should be even (random overlap)"
434            );
435        }
436    }
437
438    #[test]
439    fn calc_num_slots_test() {
440        assert_eq!(calc_num_slots(5, true), 2);
441        assert_eq!(calc_num_slots(5, false), 3);
442
443        assert_eq!(calc_num_slots(6, true), 3);
444        assert_eq!(calc_num_slots(6, false), 3);
445    }
446}