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 RandomGaussianDist<M>
 21where
 22    M: ?Sized,
 23{
 24    type Parameters: ?Sized;
 25    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut M);
 26}
 27
 28pub trait RandomUniformDist<M>
 29where
 30    M: ?Sized,
 31{
 32    type Parameters: ?Sized;
 33    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut M);
 34}
 35
 36pub trait RandomUniformDist1<M, P>
 37where
 38    M: ?Sized,
 39{
 40    fn random_fill(&mut self, modulus: &P, container: &mut M);
 41}
 42
 43pub(crate) struct DefaultSecureRng {
 44    rng: ChaCha8Rng,
 45}
 46
 47impl DefaultSecureRng {
 48    pub fn new_seeded(seed: <ChaCha8Rng as SeedableRng>::Seed) -> DefaultSecureRng {
 49        let rng = ChaCha8Rng::from_seed(seed);
 50        DefaultSecureRng { rng }
 51    }
 52
 53    pub fn new() -> DefaultSecureRng {
 54        let rng = ChaCha8Rng::from_entropy();
 55        DefaultSecureRng { rng }
 56    }
 57
 58    pub fn fill_bytes(&mut self, a: &mut [u8; 32]) {
 59        self.rng.fill_bytes(a);
 60    }
 61}
 62
 63impl NewWithSeed for DefaultSecureRng {
 64    type Seed = <ChaCha8Rng as SeedableRng>::Seed;
 65    fn new_with_seed(seed: Self::Seed) -> Self {
 66        DefaultSecureRng::new_seeded(seed)
 67    }
 68}
 69
 70impl RandomUniformDist<usize> for DefaultSecureRng {
 71    type Parameters = usize;
 72    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut usize) {
 73        *container = self.rng.gen_range(0..*parameters);
 74    }
 75}
 76
 77impl RandomUniformDist<[u8]> for DefaultSecureRng {
 78    type Parameters = u8;
 79    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut [u8]) {
 80        self.rng.fill_bytes(container);
 81    }
 82}
 83
 84impl RandomUniformDist<[u32]> for DefaultSecureRng {
 85    type Parameters = u32;
 86    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut [u32]) {
 87        izip!(
 88            (&mut self.rng).sample_iter(Uniform::new(0, parameters)),
 89            container.iter_mut()
 90        )
 91        .for_each(|(from, to)| {
 92            *to = from;
 93        });
 94    }
 95}
 96
 97impl<T, C> RandomUniformDist1<[T], C> for DefaultSecureRng
 98where
 99    T: PrimInt + SampleUniform,
100    C: Modulus<Element = T>,
101{
102    fn random_fill(&mut self, modulus: &C, container: &mut [T]) {
103        izip!(
104            (&mut self.rng).sample_iter(Uniform::new_inclusive(
105                T::zero(),
106                modulus.largest_unsigned_value()
107            )),
108            container.iter_mut()
109        )
110        .for_each(|(from, to)| {
111            *to = from;
112        });
113    }
114}
115
116impl RandomUniformDist<[u64]> for DefaultSecureRng {
117    type Parameters = u64;
118    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut [u64]) {
119        izip!(
120            (&mut self.rng).sample_iter(Uniform::new(0, parameters)),
121            container.iter_mut()
122        )
123        .for_each(|(from, to)| {
124            *to = from;
125        });
126    }
127}
128
129impl RandomGaussianDist<u64> for DefaultSecureRng {
130    type Parameters = u64;
131    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut u64) {
132        let o = rand_distr::Normal::new(0.0, 3.19f64)
133            .unwrap()
134            .sample(&mut self.rng)
135            .round();
136
137        // let o = 0.0f64;
138
139        let is_neg = o.is_sign_negative() && o != 0.0;
140        if is_neg {
141            *container = parameters - (o.abs() as u64);
142        } else {
143            *container = o as u64;
144        }
145    }
146}
147
148impl RandomGaussianDist<u32> for DefaultSecureRng {
149    type Parameters = u32;
150    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut u32) {
151        let o = rand_distr::Normal::new(0.0, 3.19f32)
152            .unwrap()
153            .sample(&mut self.rng)
154            .round();
155
156        // let o = 0.0f32;
157        let is_neg = o.is_sign_negative() && o != 0.0;
158
159        if is_neg {
160            *container = parameters - (o.abs() as u32);
161        } else {
162            *container = o as u32;
163        }
164    }
165}
166
167impl RandomGaussianDist<[u64]> for DefaultSecureRng {
168    type Parameters = u64;
169    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut [u64]) {
170        izip!(
171            rand_distr::Normal::new(0.0, 3.19f64)
172                .unwrap()
173                .sample_iter(&mut self.rng),
174            container.iter_mut()
175        )
176        .for_each(|(oi, v)| {
177            let oi = oi.round();
178            let is_neg = oi.is_sign_negative() && oi != 0.0;
179            if is_neg {
180                *v = parameters - (oi.abs() as u64);
181            } else {
182                *v = oi as u64;
183            }
184        });
185    }
186}
187
188impl RandomGaussianDist<[u32]> for DefaultSecureRng {
189    type Parameters = u32;
190    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut [u32]) {
191        izip!(
192            rand_distr::Normal::new(0.0, 3.19f32)
193                .unwrap()
194                .sample_iter(&mut self.rng),
195            container.iter_mut()
196        )
197        .for_each(|(oi, v)| {
198            let oi = oi.round();
199            let is_neg = oi.is_sign_negative() && oi != 0.0;
200            if is_neg {
201                *v = parameters - (oi.abs() as u32);
202            } else {
203                *v = oi as u32;
204            }
205        });
206    }
207}
208
209impl WithLocal for DefaultSecureRng {
210    fn with_local<F, R>(func: F) -> R
211    where
212        F: Fn(&Self) -> R,
213    {
214        DEFAULT_RNG.with_borrow(|r| func(r))
215    }
216
217    fn with_local_mut<F, R>(func: F) -> R
218    where
219        F: Fn(&mut Self) -> R,
220    {
221        DEFAULT_RNG.with_borrow_mut(|r| func(r))
222    }
223}