1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169

use alloc::vec::Vec;
use core::{
    mem::transmute,
    sync::atomic::{AtomicU64, Ordering},
};
use num_traits::{Float, Zero};
use rand::{
    distributions::{
        uniform::{SampleRange, SampleUniform},
        Distribution,
    },
    seq::SliceRandom as _,
    Rng as _, SeedableRng,
};
use rand_chacha::ChaCha8Rng;

/// The seed type of the ChaCha algorithm.
pub type Seed = <ChaCha8Rng as SeedableRng>::Seed;

struct AtomicU128 {
    s1: AtomicU64,
    s2: AtomicU64,
}

impl AtomicU128 {
    fn new(v: u128) -> Self {
        let [a, b] = unsafe { transmute::<_, [_; 2]>(v) };
        Self { s1: AtomicU64::new(a), s2: AtomicU64::new(b) }
    }

    fn load(&self, order: Ordering) -> u128 {
        unsafe { transmute([self.s1.load(order), self.s2.load(order)]) }
    }

    fn store(&self, v: u128, order: Ordering) {
        let [a, b] = unsafe { transmute::<_, [_; 2]>(v) };
        self.s1.store(a, order);
        self.s2.store(b, order);
    }
}

/// An uniformed random number generator.
///
/// This generator doesn't require mutability,
/// because the state is saved as atomic values.
pub struct Rng {
    seed: Seed,
    stream: AtomicU64,
    word_pos: AtomicU128,
}

impl Rng {
    /// Create generator by a given seed.
    /// If none, create the seed from CPU random function.
    pub fn new(seed: Option<Seed>) -> Self {
        let rng = seed
            .map(ChaCha8Rng::from_seed)
            .unwrap_or_else(ChaCha8Rng::from_entropy);
        Self {
            seed: rng.get_seed(),
            stream: AtomicU64::new(rng.get_stream()),
            word_pos: AtomicU128::new(rng.get_word_pos()),
        }
    }

    /// Seed of this generator.
    pub fn seed(&self) -> Seed {
        self.seed
    }

    /// Low-level access to the RNG type.
    ///
    /// Please import necessary traits first.
    pub fn gen<R>(&self, f: impl FnOnce(&mut ChaCha8Rng) -> R) -> R {
        let mut rng = ChaCha8Rng::from_seed(self.seed);
        rng.set_stream(self.stream.load(Ordering::SeqCst));
        rng.set_word_pos(self.word_pos.load(Ordering::SeqCst));
        let r = f(&mut rng);
        self.stream.store(rng.get_stream(), Ordering::SeqCst);
        self.word_pos.store(rng.get_word_pos(), Ordering::SeqCst);
        r
    }

    /// Generate a classic random value between `0..1` (exclusive range).
    #[inline]
    pub fn rand(&self) -> f64 {
        self.ub(1.)
    }

    /// Generate a random boolean by positive (`true`) factor.
    #[inline]
    pub fn maybe(&self, p: f64) -> bool {
        self.gen(|r| r.gen_bool(p))
    }

    /// Generate a random value by range.
    #[inline]
    pub fn range<T, R>(&self, range: R) -> T
    where
        T: SampleUniform,
        R: SampleRange<T>,
    {
        self.gen(|r| r.gen_range(range))
    }

    /// Sample from a distribution.
    #[inline]
    pub fn sample<T, D>(&self, distr: D) -> T
    where
        D: Distribution<T>,
    {
        self.gen(|r| r.sample(distr))
    }

    /// Generate a random value by upper bound (exclusive range).
    ///
    /// The lower bound is zero.
    #[inline]
    pub fn ub<U>(&self, ub: U) -> U
    where
        U: Zero + SampleUniform,
        core::ops::Range<U>: SampleRange<U>,
    {
        self.range(U::zero()..ub)
    }

    /// Sample with Gaussian distribution.
    #[inline]
    pub fn normal<F>(&self, mean: F, std: F) -> F
    where
        F: Float,
        rand_distr::StandardNormal: Distribution<F>,
    {
        self.sample(rand_distr::Normal::new(mean, std).unwrap())
    }

    /// Shuffle a slice.
    pub fn shuffle<A>(&self, s: &mut [A]) {
        self.gen(|r| s.shuffle(r));
    }

    /// Generate a random array with no-repeat values.
    pub fn array<A, C, const N: usize>(&self, candi: C) -> [A; N]
    where
        A: Zero + Copy + PartialEq + SampleUniform,
        C: IntoIterator<Item = A>,
    {
        self.array_by([A::zero(); N], 0, candi)
    }

    /// Fill a mutable slice with no-repeat values.
    ///
    /// The start position of the vector can be set.
    pub fn array_by<A, V, C>(&self, mut v: V, start: usize, candi: C) -> V
    where
        A: PartialEq + SampleUniform,
        V: AsMut<[A]>,
        C: IntoIterator<Item = A>,
    {
        let (pre, curr) = v.as_mut().split_at_mut(start);
        let mut candi = candi
            .into_iter()
            .filter(|e| !pre.contains(e))
            .collect::<Vec<_>>();
        self.shuffle(&mut candi);
        curr.iter_mut().zip(candi).for_each(|(a, b)| *a = b);
        v
    }
}