1use crate::{Infos, Module, VecZnxApi};
 2use rand_distr::{Distribution, Normal};
 3use sampling::source::Source;
 4
 5pub trait Sampling<T: VecZnxApi + Infos> {
 6    /// Fills the first `limbs` limbs with uniform values in \[-2^{log_base2k-1}, 2^{log_base2k-1}\]
 7    fn fill_uniform(&self, log_base2k: usize, a: &mut T, limbs: usize, source: &mut Source);
 8
 9    /// Adds vector sampled according to the provided distribution, scaled by 2^{-log_k} and bounded to \[-bound, bound\].
10    fn add_dist_f64<D: Distribution<f64>>(
11        &self,
12        log_base2k: usize,
13        a: &mut T,
14        log_k: usize,
15        source: &mut Source,
16        dist: D,
17        bound: f64,
18    );
19
20    /// Adds a discrete normal vector scaled by 2^{-log_k} with the provided standard deviation and bounded to \[-bound, bound\].
21    fn add_normal(
22        &self,
23        log_base2k: usize,
24        a: &mut T,
25        log_k: usize,
26        source: &mut Source,
27        sigma: f64,
28        bound: f64,
29    );
30}
31
32impl<T: VecZnxApi + Infos> Sampling<T> for Module {
33    fn fill_uniform(&self, log_base2k: usize, a: &mut T, limbs: usize, source: &mut Source) {
34        let base2k: u64 = 1 << log_base2k;
35        let mask: u64 = base2k - 1;
36        let base2k_half: i64 = (base2k >> 1) as i64;
37        let size: usize = a.n() * limbs;
38        a.raw_mut()[..size]
39            .iter_mut()
40            .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
41    }
42
43    fn add_dist_f64<D: Distribution<f64>>(
44        &self,
45        log_base2k: usize,
46        a: &mut T,
47        log_k: usize,
48        source: &mut Source,
49        dist: D,
50        bound: f64,
51    ) {
52        assert!(
53            (bound.log2().ceil() as i64) < 64,
54            "invalid bound: ceil(log2(bound))={} > 63",
55            (bound.log2().ceil() as i64)
56        );
57
58        let log_base2k_rem: usize = log_k % log_base2k;
59
60        if log_base2k_rem != 0 {
61            a.at_mut(a.limbs() - 1).iter_mut().for_each(|a| {
62                let mut dist_f64: f64 = dist.sample(source);
63                while dist_f64.abs() > bound {
64                    dist_f64 = dist.sample(source)
65                }
66                *a += (dist_f64.round() as i64) << log_base2k_rem
67            });
68        } else {
69            a.at_mut(a.limbs() - 1).iter_mut().for_each(|a| {
70                let mut dist_f64: f64 = dist.sample(source);
71                while dist_f64.abs() > bound {
72                    dist_f64 = dist.sample(source)
73                }
74                *a += dist_f64.round() as i64
75            });
76        }
77    }
78
79    fn add_normal(
80        &self,
81        log_base2k: usize,
82        a: &mut T,
83        log_k: usize,
84        source: &mut Source,
85        sigma: f64,
86        bound: f64,
87    ) {
88        self.add_dist_f64(
89            log_base2k,
90            a,
91            log_k,
92            source,
93            Normal::new(0.0, sigma).unwrap(),
94            bound,
95        );
96    }
97}