randgen/
lib.rs

1#![allow(unused_imports)]
2
3use std::{collections::BTreeSet, ops::Range};
4
5use rand::{
6    Rng, SeedableRng,
7    distributions::{Distribution, Uniform},
8};
9use rand_chacha::ChaCha20Rng;
10
11pub trait Gen {
12    type Output;
13    fn generate<R: Rng>(&self, rng: &mut R) -> Self::Output;
14}
15
16pub struct StrictAsc<B> {
17    bound: B,
18    len: usize,
19}
20
21pub struct Asc<B> {
22    bound: B,
23    len: usize,
24}
25
26impl Gen for Range<i32> {
27    type Output = i32;
28    fn generate<R: Rng>(&self, rng: &mut R) -> Self::Output {
29        let between = Uniform::from(self.clone());
30        between.sample(rng)
31    }
32}
33
34impl Gen for StrictAsc<Range<i32>> {
35    type Output = Vec<i32>;
36    fn generate<R: Rng>(&self, rng: &mut R) -> Self::Output {
37        let Self { bound: Range { start, end }, len } = self;
38
39        // n = end - start, k = len
40        let dense = (2 * len) as i32 > (end - start) / 2;
41        let count = if dense { (end - start) as usize - len } else { *len };
42
43        let mut seen = BTreeSet::new();
44        while seen.len() < count {
45            seen.insert((*start..*end).generate(rng));
46        }
47
48        if dense {
49            (*start..*end).filter(|x| !seen.contains(x)).collect()
50        } else {
51            seen.into_iter().collect()
52        }
53    }
54}
55
56impl Gen for Asc<Range<i32>> {
57    type Output = Vec<i32>;
58    fn generate<R: Rng>(&self, rng: &mut R) -> Self::Output {
59        let Self { bound: Range { start, end }, len } = self;
60        let mut strict =
61            StrictAsc { bound: *start..end + *len as i32, len: *len }
62                .generate(rng);
63        for i in 0..*len {
64            strict[i] -= i as i32;
65        }
66        strict
67    }
68}
69
70#[test]
71fn uniformity() {
72    use std::collections::BTreeMap;
73
74    let mut rng = ChaCha20Rng::from_seed([0; 32]);
75    let n = 10_usize.pow(6);
76
77    let mut map = BTreeMap::new();
78    for _ in 0..n {
79        let tmp = StrictAsc { bound: 0..4, len: 3 }.generate(&mut rng);
80        *map.entry(tmp).or_insert(0) += 1;
81    }
82    let k = 4;
83    assert_eq!(map.len(), k);
84    for &v in map.values() {
85        assert!(v >= (n / k) * 99 / 100);
86        assert!(v <= (n / k) * 101 / 100);
87    }
88
89    let mut map = BTreeMap::new();
90    for _ in 0..n {
91        let tmp = Asc { bound: 0..4, len: 3 }.generate(&mut rng);
92        *map.entry(tmp).or_insert(0) += 1;
93    }
94    let k = 35;
95    assert_eq!(map.len(), k);
96    for &v in map.values() {
97        assert!(v >= (n / k) * 97 / 100);
98        assert!(v <= (n / k) * 103 / 100);
99    }
100}
101
102#[test]
103#[cfg(false)]
104fn macros() {
105    rand_gen! {
106        rng = _; // Default
107        n in 1_usize..10;
108        a in [0..10_i32.pow(9); n];
109        b in Asc { bound: 1..=5, len: 3 };
110        c in StrictAsc { bound: 1..=5, len: 3 };
111    }
112    let mut rng = ChaCha20Rng::from_seed([0; 32]);
113    rand_gen! {
114        rng = &mut rng;
115        n in 1_usize..10;
116        a in [0..10_i32.pow(9); n];
117        b in Asc { bound: 1..=5, len: 3 };
118        c in StrictAsc { bound: 1..=5, len: 3 };
119    }
120}
randgen/lib.rs

randgen/
lib.rs
