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