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