1use std::cell::RefCell;
  2
  3use itertools::izip;
  4use num_traits::{FromPrimitive, PrimInt, Zero};
  5use rand::{distributions::Uniform, Rng, RngCore, SeedableRng};
  6use rand_chacha::ChaCha8Rng;
  7use rand_distr::{uniform::SampleUniform, Distribution};
  8
  9use crate::{backend::Modulus, utils::WithLocal};
 10
 11thread_local! {
 12    pub(crate) static DEFAULT_RNG: RefCell<DefaultSecureRng> = RefCell::new(DefaultSecureRng::new());
 13}
 14
 15pub trait NewWithSeed {
 16    type Seed;
 17    fn new_with_seed(seed: Self::Seed) -> Self;
 18}
 19
 20pub trait RandomElementInModulus<T, M> {
 21    /// Sample Random element of type T in range [0, modulus)
 22    fn random(&mut self, modulus: &M) -> T;
 23}
 24
 25pub trait RandomGaussianElementInModulus<T, M> {
 26    /// Sample Random gaussian element from \mu = 0.0 and \sigma = 3.19. Sampled
 27    /// element is converted to signed representation in modulus.
 28    fn random(&mut self, modulus: &M) -> T;
 29}
 30
 31pub trait RandomFill<M>
 32where
 33    M: ?Sized,
 34{
 35    /// Fill container with random elements of type of its elements
 36    fn random_fill(&mut self, container: &mut M);
 37}
 38
 39pub trait RandomFillGaussian<M>
 40where
 41    M: ?Sized,
 42{
 43    /// Fill container with random elements sampled from normal distribtuion
 44    /// with \mu = 0.0 and \sigma = 3.19.
 45    fn random_fill(&mut self, container: &mut M);
 46}
 47
 48pub trait RandomFillUniformInModulus<M, P>
 49where
 50    M: ?Sized,
 51{
 52    /// Fill container with random elements in range [0, modulus)
 53    fn random_fill(&mut self, modulus: &P, container: &mut M);
 54}
 55
 56pub trait RandomFillGaussianInModulus<M, P>
 57where
 58    M: ?Sized,
 59{
 60    /// Fill container with gaussian elements sampled from normal distribution
 61    /// with \mu = 0.0 and \sigma = 3.19. Elements are converted to signed
 62    /// represented in the modulus.
 63    fn random_fill(&mut self, modulus: &P, container: &mut M);
 64}
 65
 66pub struct DefaultSecureRng {
 67    rng: ChaCha8Rng,
 68}
 69
 70impl DefaultSecureRng {
 71    pub fn new_seeded(seed: <ChaCha8Rng as SeedableRng>::Seed) -> DefaultSecureRng {
 72        let rng = ChaCha8Rng::from_seed(seed);
 73        DefaultSecureRng { rng }
 74    }
 75
 76    pub fn new() -> DefaultSecureRng {
 77        let rng = ChaCha8Rng::from_entropy();
 78        DefaultSecureRng { rng }
 79    }
 80
 81    pub fn fill_bytes(&mut self, a: &mut [u8; 32]) {
 82        self.rng.fill_bytes(a);
 83    }
 84}
 85
 86impl NewWithSeed for DefaultSecureRng {
 87    type Seed = <ChaCha8Rng as SeedableRng>::Seed;
 88    fn new_with_seed(seed: Self::Seed) -> Self {
 89        DefaultSecureRng::new_seeded(seed)
 90    }
 91}
 92
 93impl<T, C> RandomFillUniformInModulus<[T], C> for DefaultSecureRng
 94where
 95    T: PrimInt + SampleUniform,
 96    C: Modulus<Element = T>,
 97{
 98    fn random_fill(&mut self, modulus: &C, container: &mut [T]) {
 99        izip!(
100            (&mut self.rng).sample_iter(Uniform::new_inclusive(
101                T::zero(),
102                modulus.largest_unsigned_value()
103            )),
104            container.iter_mut()
105        )
106        .for_each(|(from, to)| {
107            *to = from;
108        });
109    }
110}
111
112impl<T, C> RandomFillGaussianInModulus<[T], C> for DefaultSecureRng
113where
114    T: PrimInt,
115    C: Modulus<Element = T>,
116{
117    fn random_fill(&mut self, modulus: &C, container: &mut [T]) {
118        izip!(
119            rand_distr::Normal::new(0.0, 3.19f64)
120                .unwrap()
121                .sample_iter(&mut self.rng),
122            container.iter_mut()
123        )
124        .for_each(|(from, to)| {
125            *to = modulus.map_element_from_f64(from);
126        });
127    }
128}
129
130impl<T> RandomFill<[T]> for DefaultSecureRng
131where
132    T: PrimInt + SampleUniform,
133{
134    fn random_fill(&mut self, container: &mut [T]) {
135        izip!(
136            (&mut self.rng).sample_iter(Uniform::new_inclusive(T::zero(), T::max_value())),
137            container.iter_mut()
138        )
139        .for_each(|(from, to)| {
140            *to = from;
141        });
142    }
143}
144
145impl<T> RandomFillGaussian<[T]> for DefaultSecureRng
146where
147    T: FromPrimitive,
148{
149    fn random_fill(&mut self, container: &mut [T]) {
150        izip!(
151            rand_distr::Normal::new(0.0, 3.19f64)
152                .unwrap()
153                .sample_iter(&mut self.rng),
154            container.iter_mut()
155        )
156        .for_each(|(from, to)| {
157            *to = T::from_f64(from).unwrap();
158        });
159    }
160}
161
162impl<T> RandomFill<[T; 32]> for DefaultSecureRng
163where
164    T: PrimInt + SampleUniform,
165{
166    fn random_fill(&mut self, container: &mut [T; 32]) {
167        izip!(
168            (&mut self.rng).sample_iter(Uniform::new_inclusive(T::zero(), T::max_value())),
169            container.iter_mut()
170        )
171        .for_each(|(from, to)| {
172            *to = from;
173        });
174    }
175}
176
177impl<T> RandomElementInModulus<T, T> for DefaultSecureRng
178where
179    T: Zero + SampleUniform,
180{
181    fn random(&mut self, modulus: &T) -> T {
182        Uniform::new(T::zero(), modulus).sample(&mut self.rng)
183    }
184}
185
186impl<T, M: Modulus<Element = T>> RandomGaussianElementInModulus<T, M> for DefaultSecureRng {
187    fn random(&mut self, modulus: &M) -> T {
188        modulus.map_element_from_f64(
189            rand_distr::Normal::new(0.0, 3.19f64)
190                .unwrap()
191                .sample(&mut self.rng),
192        )
193    }
194}
195
196impl WithLocal for DefaultSecureRng {
197    fn with_local<F, R>(func: F) -> R
198    where
199        F: Fn(&Self) -> R,
200    {
201        DEFAULT_RNG.with_borrow(|r| func(r))
202    }
203
204    fn with_local_mut<F, R>(func: F) -> R
205    where
206        F: Fn(&mut Self) -> R,
207    {
208        DEFAULT_RNG.with_borrow_mut(|r| func(r))
209    }
210
211    fn with_local_mut_mut<F, R>(func: &mut F) -> R
212    where
213        F: FnMut(&mut Self) -> R,
214    {
215        DEFAULT_RNG.with_borrow_mut(|r| func(r))
216    }
217}