1use crate::{Backend, Module, VecZnx, ZnxInfos, ZnxLayout};
  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(&self, log_base2k: usize, a: &mut VecZnx, col_i: 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<D: Distribution<f64>>(
 11        &self,
 12        log_base2k: usize,
 13        a: &mut VecZnx,
 14        col_i: usize,
 15        log_k: usize,
 16        source: &mut Source,
 17        dist: D,
 18        bound: f64,
 19    );
 20
 21    /// Adds a discrete normal vector scaled by 2^{-log_k} with the provided standard deviation and bounded to \[-bound, bound\].
 22    fn add_normal(
 23        &self,
 24        log_base2k: usize,
 25        a: &mut VecZnx,
 26        col_i: usize,
 27        log_k: usize,
 28        source: &mut Source,
 29        sigma: f64,
 30        bound: f64,
 31    );
 32}
 33
 34impl<B: Backend> Sampling for Module<B> {
 35    fn fill_uniform(&self, log_base2k: usize, a: &mut VecZnx, col_i: usize, limbs: usize, source: &mut Source) {
 36        let base2k: u64 = 1 << log_base2k;
 37        let mask: u64 = base2k - 1;
 38        let base2k_half: i64 = (base2k >> 1) as i64;
 39        (0..limbs).for_each(|j| {
 40            a.at_poly_mut(col_i, j)
 41                .iter_mut()
 42                .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
 43        })
 44    }
 45
 46    fn add_dist_f64<D: Distribution<f64>>(
 47        &self,
 48        log_base2k: usize,
 49        a: &mut VecZnx,
 50        col_i: usize,
 51        log_k: usize,
 52        source: &mut Source,
 53        dist: D,
 54        bound: f64,
 55    ) {
 56        assert!(
 57            (bound.log2().ceil() as i64) < 64,
 58            "invalid bound: ceil(log2(bound))={} > 63",
 59            (bound.log2().ceil() as i64)
 60        );
 61
 62        let log_base2k_rem: usize = log_k % log_base2k;
 63
 64        if log_base2k_rem != 0 {
 65            a.at_poly_mut(col_i, a.limbs() - 1)
 66                .iter_mut()
 67                .for_each(|a| {
 68                    let mut dist_f64: f64 = dist.sample(source);
 69                    while dist_f64.abs() > bound {
 70                        dist_f64 = dist.sample(source)
 71                    }
 72                    *a += (dist_f64.round() as i64) << log_base2k_rem;
 73                });
 74        } else {
 75            a.at_poly_mut(col_i, a.limbs() - 1)
 76                .iter_mut()
 77                .for_each(|a| {
 78                    let mut dist_f64: f64 = dist.sample(source);
 79                    while dist_f64.abs() > bound {
 80                        dist_f64 = dist.sample(source)
 81                    }
 82                    *a += dist_f64.round() as i64
 83                });
 84        }
 85    }
 86
 87    fn add_normal(
 88        &self,
 89        log_base2k: usize,
 90        a: &mut VecZnx,
 91        col_i: usize,
 92        log_k: usize,
 93        source: &mut Source,
 94        sigma: f64,
 95        bound: f64,
 96    ) {
 97        self.add_dist_f64(
 98            log_base2k,
 99            a,
100            col_i,
101            log_k,
102            source,
103            Normal::new(0.0, sigma).unwrap(),
104            bound,
105        );
106    }
107}