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