1use std::cell::RefCell;
  2
  3use itertools::izip;
  4use num_traits::{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 RandomFillUniformInModulus<M, P>
 40where
 41    M: ?Sized,
 42{
 43    /// Fill container with random elements in range [0, modulus)
 44    fn random_fill(&mut self, modulus: &P, container: &mut M);
 45}
 46
 47pub trait RandomFillGaussianInModulus<M, P>
 48where
 49    M: ?Sized,
 50{
 51    /// Fill container with gaussian elements sampled from normal distribution
 52    /// with \mu = 0.0 and \sigma = 3.19. Elements are converted to signed
 53    /// represented in the modulus.
 54    fn random_fill(&mut self, modulus: &P, container: &mut M);
 55}
 56
 57pub struct DefaultSecureRng {
 58    rng: ChaCha8Rng,
 59}
 60
 61impl DefaultSecureRng {
 62    pub fn new_seeded(seed: <ChaCha8Rng as SeedableRng>::Seed) -> DefaultSecureRng {
 63        let rng = ChaCha8Rng::from_seed(seed);
 64        DefaultSecureRng { rng }
 65    }
 66
 67    pub fn new() -> DefaultSecureRng {
 68        let rng = ChaCha8Rng::from_entropy();
 69        DefaultSecureRng { rng }
 70    }
 71
 72    pub fn fill_bytes(&mut self, a: &mut [u8; 32]) {
 73        self.rng.fill_bytes(a);
 74    }
 75}
 76
 77impl NewWithSeed for DefaultSecureRng {
 78    type Seed = <ChaCha8Rng as SeedableRng>::Seed;
 79    fn new_with_seed(seed: Self::Seed) -> Self {
 80        DefaultSecureRng::new_seeded(seed)
 81    }
 82}
 83
 84impl<T, C> RandomFillUniformInModulus<[T], C> for DefaultSecureRng
 85where
 86    T: PrimInt + SampleUniform,
 87    C: Modulus<Element = T>,
 88{
 89    fn random_fill(&mut self, modulus: &C, container: &mut [T]) {
 90        izip!(
 91            (&mut self.rng).sample_iter(Uniform::new_inclusive(
 92                T::zero(),
 93                modulus.largest_unsigned_value()
 94            )),
 95            container.iter_mut()
 96        )
 97        .for_each(|(from, to)| {
 98            *to = from;
 99        });
100    }
101}
102
103impl<T, C> RandomFillGaussianInModulus<[T], C> for DefaultSecureRng
104where
105    T: PrimInt,
106    C: Modulus<Element = T>,
107{
108    fn random_fill(&mut self, modulus: &C, container: &mut [T]) {
109        izip!(
110            rand_distr::Normal::new(0.0, 3.19f64)
111                .unwrap()
112                .sample_iter(&mut self.rng),
113            container.iter_mut()
114        )
115        .for_each(|(from, to)| {
116            *to = modulus.map_element_from_f64(from);
117        });
118    }
119}
120
121impl<T> RandomFill<[T]> for DefaultSecureRng
122where
123    T: PrimInt + SampleUniform,
124{
125    fn random_fill(&mut self, container: &mut [T]) {
126        izip!(
127            (&mut self.rng).sample_iter(Uniform::new_inclusive(T::zero(), T::max_value())),
128            container.iter_mut()
129        )
130        .for_each(|(from, to)| {
131            *to = from;
132        });
133    }
134}
135
136impl<T> RandomFill<[T; 32]> for DefaultSecureRng
137where
138    T: PrimInt + SampleUniform,
139{
140    fn random_fill(&mut self, container: &mut [T; 32]) {
141        izip!(
142            (&mut self.rng).sample_iter(Uniform::new_inclusive(T::zero(), T::max_value())),
143            container.iter_mut()
144        )
145        .for_each(|(from, to)| {
146            *to = from;
147        });
148    }
149}
150
151impl<T> RandomElementInModulus<T, T> for DefaultSecureRng
152where
153    T: Zero + SampleUniform,
154{
155    fn random(&mut self, modulus: &T) -> T {
156        Uniform::new(T::zero(), modulus).sample(&mut self.rng)
157    }
158}
159
160impl<T, M: Modulus<Element = T>> RandomGaussianElementInModulus<T, M> for DefaultSecureRng {
161    fn random(&mut self, modulus: &M) -> T {
162        modulus.map_element_from_f64(
163            rand_distr::Normal::new(0.0, 3.19f64)
164                .unwrap()
165                .sample(&mut self.rng),
166        )
167    }
168}
169
170impl WithLocal for DefaultSecureRng {
171    fn with_local<F, R>(func: F) -> R
172    where
173        F: Fn(&Self) -> R,
174    {
175        DEFAULT_RNG.with_borrow(|r| func(r))
176    }
177
178    fn with_local_mut<F, R>(func: F) -> R
179    where
180        F: Fn(&mut Self) -> R,
181    {
182        DEFAULT_RNG.with_borrow_mut(|r| func(r))
183    }
184
185    fn with_local_mut_mut<F, R>(func: &mut F) -> R
186    where
187        F: FnMut(&mut Self) -> R,
188    {
189        DEFAULT_RNG.with_borrow_mut(|r| func(r))
190    }
191}