1use std::cell::RefCell;
  2
  3use itertools::izip;
  4use rand::{distributions::Uniform, thread_rng, CryptoRng, Rng, RngCore, SeedableRng};
  5use rand_chacha::ChaCha8Rng;
  6use rand_distr::Distribution;
  7
  8use crate::utils::WithLocal;
  9
 10thread_local! {
 11    pub(crate) static DEFAULT_RNG: RefCell<DefaultSecureRng> = RefCell::new(DefaultSecureRng::new());
 12}
 13
 14pub trait RandomGaussianDist<M>
 15where
 16    M: ?Sized,
 17{
 18    type Parameters: ?Sized;
 19    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut M);
 20}
 21
 22pub trait RandomUniformDist<M>
 23where
 24    M: ?Sized,
 25{
 26    type Parameters: ?Sized;
 27    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut M);
 28}
 29
 30pub(crate) struct DefaultSecureRng {
 31    rng: ChaCha8Rng,
 32}
 33
 34impl DefaultSecureRng {
 35    pub fn new_seeded(seed: <ChaCha8Rng as SeedableRng>::Seed) -> DefaultSecureRng {
 36        let rng = ChaCha8Rng::from_seed(seed);
 37        DefaultSecureRng { rng }
 38    }
 39
 40    pub fn new() -> DefaultSecureRng {
 41        let rng = ChaCha8Rng::from_entropy();
 42        DefaultSecureRng { rng }
 43    }
 44}
 45
 46impl RandomUniformDist<usize> for DefaultSecureRng {
 47    type Parameters = usize;
 48    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut usize) {
 49        *container = self.rng.gen_range(0..*parameters);
 50    }
 51}
 52
 53impl RandomUniformDist<[u8]> for DefaultSecureRng {
 54    type Parameters = u8;
 55    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut [u8]) {
 56        self.rng.fill_bytes(container);
 57    }
 58}
 59
 60impl RandomUniformDist<[u32]> for DefaultSecureRng {
 61    type Parameters = u32;
 62    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut [u32]) {
 63        izip!(
 64            (&mut self.rng).sample_iter(Uniform::new(0, parameters)),
 65            container.iter_mut()
 66        )
 67        .for_each(|(from, to)| {
 68            *to = from;
 69        });
 70    }
 71}
 72
 73impl RandomUniformDist<[u64]> for DefaultSecureRng {
 74    type Parameters = u64;
 75    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut [u64]) {
 76        izip!(
 77            (&mut self.rng).sample_iter(Uniform::new(0, parameters)),
 78            container.iter_mut()
 79        )
 80        .for_each(|(from, to)| {
 81            *to = from;
 82        });
 83    }
 84}
 85
 86impl RandomGaussianDist<u64> for DefaultSecureRng {
 87    type Parameters = u64;
 88    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut u64) {
 89        let o = rand_distr::Normal::new(0.0, 3.2f64)
 90            .unwrap()
 91            .sample(&mut self.rng)
 92            .round();
 93
 94        // let o = 0.0f64;
 95
 96        let is_neg = o.is_sign_negative() && o != 0.0;
 97        if is_neg {
 98            *container = parameters - (o.abs() as u64);
 99        } else {
100            *container = o as u64;
101        }
102    }
103}
104
105impl RandomGaussianDist<u32> for DefaultSecureRng {
106    type Parameters = u32;
107    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut u32) {
108        let o = rand_distr::Normal::new(0.0, 3.2f32)
109            .unwrap()
110            .sample(&mut self.rng)
111            .round();
112
113        // let o = 0.0f32;
114        let is_neg = o.is_sign_negative() && o != 0.0;
115
116        if is_neg {
117            *container = parameters - (o.abs() as u32);
118        } else {
119            *container = o as u32;
120        }
121    }
122}
123
124impl RandomGaussianDist<[u64]> for DefaultSecureRng {
125    type Parameters = u64;
126    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut [u64]) {
127        izip!(
128            rand_distr::Normal::new(0.0, 3.2f64)
129                .unwrap()
130                .sample_iter(&mut self.rng),
131            container.iter_mut()
132        )
133        .for_each(|(oi, v)| {
134            let oi = oi.round();
135            let is_neg = oi.is_sign_negative() && oi != 0.0;
136            if is_neg {
137                *v = parameters - (oi.abs() as u64);
138            } else {
139                *v = oi as u64;
140            }
141        });
142    }
143}
144
145impl RandomGaussianDist<[u32]> for DefaultSecureRng {
146    type Parameters = u32;
147    fn random_fill(&mut self, parameters: &Self::Parameters, container: &mut [u32]) {
148        izip!(
149            rand_distr::Normal::new(0.0, 3.2f32)
150                .unwrap()
151                .sample_iter(&mut self.rng),
152            container.iter_mut()
153        )
154        .for_each(|(oi, v)| {
155            let oi = oi.round();
156            let is_neg = oi.is_sign_negative() && oi != 0.0;
157            if is_neg {
158                *v = parameters - (oi.abs() as u32);
159            } else {
160                *v = oi as u32;
161            }
162        });
163    }
164}
165
166impl WithLocal for DefaultSecureRng {
167    fn with_local<F, R>(func: F) -> R
168    where
169        F: Fn(&Self) -> R,
170    {
171        DEFAULT_RNG.with_borrow(|r| func(r))
172    }
173
174    fn with_local_mut<F, R>(func: F) -> R
175    where
176        F: Fn(&mut Self) -> R,
177    {
178        DEFAULT_RNG.with_borrow_mut(|r| func(r))
179    }
180}