hotshot_example_types/membership/
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(
125                prev_rng,
126                calc_num_slots(count, round.is_multiple_of(2)),
127            ),
128            this_rng: NonRepeatValueIterator::new(this_rng, calc_num_slots(count, round % 2 == 1)),
129            round,
130            count,
131            overlap,
132            index: 0,
133        }
134    }
135}
136
137impl Iterator for StableQuorumIterator {
138    type Item = u64;
139
140    fn next(&mut self) -> Option<Self::Item> {
141        if self.index >= (self.count / 2) {
142            // Always return exactly half of the possible values. If we have OVERLAP>0 then
143            // we need to return (COUNT/2)-OVERLAP of the current set, even if there are additional
144            // even (or odd) numbers that we can return.
145            None
146        } else if self.index < self.overlap {
147            // Generate enough values for the previous round. If the current round is odd, then
148            // we want to pick even values that were selected from the previous round to create OVERLAP
149            // even values.
150            let v = self.prev_rng.next().unwrap();
151            self.index += 1;
152            Some(v * 2 + (1 - self.round % 2))
153        } else {
154            // Generate new values. If our current round is odd, we'll be creating (COUNT/2)-OVERLAP
155            // odd values here.
156            let v = self.this_rng.next().unwrap();
157            self.index += 1;
158            Some(v * 2 + self.round % 2)
159        }
160    }
161}
162
163#[must_use]
164/// Helper function to convert the arguments to a StableQuorumIterator into an ordered set of values.
165///
166/// # Panics
167///
168/// panics if the arguments are invalid for StableQuorumIterator::new
169pub fn stable_quorum_filter(seed: u64, round: u64, count: usize, overlap: u64) -> BTreeSet<usize> {
170    StableQuorumIterator::new(seed, round, count as u64, overlap)
171        // We should never have more than u32_max members in a test
172        .map(|x| usize::try_from(x).unwrap())
173        .collect()
174}
175
176/// Constructs a quorum with a random number of members and overlaps. Functions similar to StableQuorumIterator,
177/// except that the number of MEMBERS and OVERLAP are also (deterministically) random, to allow additional variance
178/// in testing.
179pub struct RandomOverlapQuorumIterator {
180    /// PRNG from the previous round
181    prev_rng: NonRepeatValueIterator,
182
183    /// PRNG for the current round
184    this_rng: NonRepeatValueIterator,
185
186    /// Current ROUND
187    round: u64,
188
189    /// Number of members to emit for the current round
190    members: u64,
191
192    /// OVERLAP of nodes to be carried over from the previous round
193    overlap: u64,
194
195    /// The next call to next() will emit the value with this index. Starts at 0 and is incremented for each
196    /// call to next()
197    index: u64,
198}
199
200impl RandomOverlapQuorumIterator {
201    #[must_use]
202    /// Create a new RandomOverlapQuorumIterator
203    ///
204    /// # Panics
205    ///
206    /// panics if overlap and members can produce invalid results or if ranges are invalid
207    pub fn new(
208        seed: u64,
209        round: u64,
210        count: u64,
211        members_min: u64,
212        members_max: u64,
213        overlap_min: u64,
214        overlap_max: u64,
215    ) -> Self {
216        assert!(
217            members_min <= members_max,
218            "Members_min cannot be greater than members_max"
219        );
220        assert!(
221            overlap_min <= overlap_max,
222            "Overlap_min cannot be greater than overlap_max"
223        );
224        assert!(
225            overlap_max < members_min,
226            "Overlap_max must be less than members_min"
227        );
228        assert!(
229            count / 2 > overlap_max,
230            "Overlap cannot be greater than the entire set size"
231        );
232        assert!(
233            count / 2 >= members_max - overlap_min,
234            "members_max must be greater or equal to half of the count plus overlap_min"
235        );
236
237        let (mut prev_rng, mut this_rng) = make_rngs(seed, round);
238
239        // Consume two values from prev_rng to advance it to the same state it was at the beginning of the previous round
240        let _prev_members = prev_rng.gen_range(members_min..=members_max);
241        let _prev_overlap = prev_rng.gen_range(overlap_min..=overlap_max);
242        let this_members = this_rng.gen_range(members_min..=members_max);
243        let this_overlap = this_rng.gen_range(overlap_min..=overlap_max);
244
245        Self {
246            prev_rng: NonRepeatValueIterator::new(
247                prev_rng,
248                calc_num_slots(count, round.is_multiple_of(2)),
249            ),
250            this_rng: NonRepeatValueIterator::new(this_rng, calc_num_slots(count, round % 2 == 1)),
251            round,
252            members: this_members,
253            overlap: this_overlap,
254            index: 0,
255        }
256    }
257}
258
259impl Iterator for RandomOverlapQuorumIterator {
260    type Item = u64;
261
262    fn next(&mut self) -> Option<Self::Item> {
263        if self.index >= self.members {
264            None
265        } else if self.index < self.overlap {
266            // Generate enough values for the previous round
267            let v = self.prev_rng.next().unwrap();
268            self.index += 1;
269            Some(v * 2 + (1 - self.round % 2))
270        } else {
271            // Generate new values
272            let v = self.this_rng.next().unwrap();
273            self.index += 1;
274            Some(v * 2 + self.round % 2)
275        }
276    }
277}
278
279#[must_use]
280/// Helper function to convert the arguments to a StableQuorumIterator into an ordered set of values.
281///
282/// # Panics
283///
284/// panics if the arguments are invalid for RandomOverlapQuorumIterator::new
285pub fn random_overlap_quorum_filter(
286    seed: u64,
287    round: u64,
288    count: usize,
289    members_min: u64,
290    members_max: u64,
291    overlap_min: u64,
292    overlap_max: u64,
293) -> BTreeSet<usize> {
294    RandomOverlapQuorumIterator::new(
295        seed,
296        round,
297        count as u64,
298        members_min,
299        members_max,
300        overlap_min,
301        overlap_max,
302    )
303    // We should never have more than u32_max members in a test
304    .map(|x| usize::try_from(x).unwrap())
305    .collect()
306}
307
308/// Trait wrapping a config for quorum filters. This allows selection between either the StableQuorumIterator or the
309/// RandomOverlapQuorumIterator functionality from above
310pub trait QuorumFilterConfig:
311    Copy
312    + Clone
313    + std::fmt::Debug
314    + Default
315    + Send
316    + Sync
317    + Ord
318    + PartialOrd
319    + Eq
320    + PartialEq
321    + Hash
322    + 'static
323{
324    /// Called to run the filter and return a set of indices
325    fn execute(epoch: u64, count: usize) -> BTreeSet<usize>;
326}
327
328#[derive(Debug, Copy, Clone, Default, Eq, PartialEq, Hash, Ord, PartialOrd)]
329/// Provides parameters to use the StableQuorumIterator
330pub struct StableQuorumFilterConfig<const SEED: u64, const OVERLAP: u64> {}
331
332impl<const SEED: u64, const OVERLAP: u64> QuorumFilterConfig
333    for StableQuorumFilterConfig<SEED, OVERLAP>
334{
335    fn execute(epoch: u64, count: usize) -> BTreeSet<usize> {
336        stable_quorum_filter(SEED, epoch, count, OVERLAP)
337    }
338}
339
340#[derive(Debug, Copy, Clone, Default, Eq, PartialEq, Hash, Ord, PartialOrd)]
341/// Provides parameters to use the RandomOverlapQuorumIterator
342pub struct RandomOverlapQuorumFilterConfig<
343    const SEED: u64,
344    const MEMBERS_MIN: u64,
345    const MEMBERS_MAX: u64,
346    const OVERLAP_MIN: u64,
347    const OVERLAP_MAX: u64,
348> {}
349
350impl<
351        const SEED: u64,
352        const MEMBERS_MIN: u64,
353        const MEMBERS_MAX: u64,
354        const OVERLAP_MIN: u64,
355        const OVERLAP_MAX: u64,
356    > QuorumFilterConfig
357    for RandomOverlapQuorumFilterConfig<SEED, MEMBERS_MIN, MEMBERS_MAX, OVERLAP_MIN, OVERLAP_MAX>
358{
359    fn execute(epoch: u64, count: usize) -> BTreeSet<usize> {
360        random_overlap_quorum_filter(
361            SEED,
362            epoch,
363            count,
364            MEMBERS_MIN,
365            MEMBERS_MAX,
366            OVERLAP_MIN,
367            OVERLAP_MAX,
368        )
369    }
370}
371
372#[cfg(test)]
373mod tests {
374    use super::*;
375
376    #[test]
377    fn test_stable() {
378        for _ in 0..100 {
379            let seed = rand::random::<u64>();
380            let prev_set: Vec<u64> = StableQuorumIterator::new(seed, 1, 10, 2).collect();
381            let this_set: Vec<u64> = StableQuorumIterator::new(seed, 2, 10, 2).collect();
382
383            // The first two elements from prev_set are from its previous round. But its 2nd and 3rd elements
384            // are new, and should be carried over to become the first two elements from this_set.
385            assert_eq!(
386                prev_set[2..4],
387                this_set[0..2],
388                "prev_set={prev_set:?}, this_set={this_set:?}"
389            );
390        }
391    }
392
393    #[test]
394    fn test_random_overlap() {
395        for _ in 0..100 {
396            let seed = rand::random::<u64>();
397            let prev_set: Vec<u64> =
398                RandomOverlapQuorumIterator::new(seed, 1, 20, 5, 10, 2, 3).collect();
399            let this_set: Vec<u64> =
400                RandomOverlapQuorumIterator::new(seed, 2, 20, 5, 10, 2, 3).collect();
401
402            // Similar to the overlap before, but there are 4 possible cases: the previous set might have had
403            // either 2 or 3 overlaps, meaning we should start with index 2 or 3, and the overlap size might
404            // be either 2 or 3. We'll just check for 2 overlaps, meaning we have two possible overlap cases
405            // to verify.
406            let matched = (prev_set[2..4] == this_set[0..2]) || (prev_set[3..5] == this_set[0..2]);
407            assert!(matched, "prev_set={prev_set:?}, this_set={this_set:?}");
408        }
409    }
410
411    #[test]
412    fn test_odd_even() {
413        for _ in 0..100 {
414            let seed = rand::random::<u64>();
415
416            let odd_set: Vec<u64> = StableQuorumIterator::new(seed, 1, 10, 2).collect();
417            let even_set: Vec<u64> = StableQuorumIterator::new(seed, 2, 10, 2).collect();
418
419            assert!(
420                odd_set[2] % 2 == 1,
421                "odd set non-overlap value should be odd (stable)"
422            );
423            assert!(
424                even_set[2].is_multiple_of(2),
425                "even set non-overlap value should be even (stable)"
426            );
427
428            let odd_set: Vec<u64> =
429                RandomOverlapQuorumIterator::new(seed, 1, 20, 5, 10, 2, 3).collect();
430            let even_set: Vec<u64> =
431                RandomOverlapQuorumIterator::new(seed, 2, 20, 5, 10, 2, 3).collect();
432
433            assert!(
434                odd_set[3] % 2 == 1,
435                "odd set non-overlap value should be odd (random overlap)"
436            );
437            assert!(
438                even_set[3].is_multiple_of(2),
439                "even set non-overlap value should be even (random overlap)"
440            );
441        }
442    }
443
444    #[test]
445    fn calc_num_slots_test() {
446        assert_eq!(calc_num_slots(5, true), 2);
447        assert_eq!(calc_num_slots(5, false), 3);
448
449        assert_eq!(calc_num_slots(6, true), 3);
450        assert_eq!(calc_num_slots(6, false), 3);
451    }
452}
hotshot_example_types/membership/helpers.rs

hotshot_example_types/membership/
helpers.rs
