1use crate::ffi::svp;
  2use crate::{Free, Module, VecZnx, VecZnxDft};
  3
  4use crate::Infos;
  5use rand::seq::SliceRandom;
  6use rand_core::RngCore;
  7use rand_distr::{Distribution, WeightedIndex};
  8use sampling::source::Source;
  9
 10pub struct Scalar(pub Vec<i64>);
 11
 12impl Module {
 13    pub fn new_scalar(&self) -> Scalar {
 14        Scalar::new(self.n())
 15    }
 16}
 17
 18impl Scalar {
 19    pub fn new(n: usize) -> Self {
 20        Self(vec![i64::default(); Self::buffer_size(n)])
 21    }
 22
 23    pub fn buffer_size(n: usize) -> usize {
 24        n
 25    }
 26
 27    pub fn from_buffer(&mut self, n: usize, buf: &[i64]) {
 28        let size: usize = Self::buffer_size(n);
 29        assert!(
 30            buf.len() >= size,
 31            "invalid buffer: buf.len()={} < self.buffer_size(n={})={}",
 32            buf.len(),
 33            n,
 34            size
 35        );
 36        self.0 = Vec::from(&buf[..size])
 37    }
 38
 39    pub fn as_ptr(&self) -> *const i64 {
 40        self.0.as_ptr()
 41    }
 42
 43    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
 44        let choices: [i64; 3] = [-1, 0, 1];
 45        let weights: [f64; 3] = [prob / 2.0, 1.0 - prob, prob / 2.0];
 46        let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
 47        self.0
 48            .iter_mut()
 49            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
 50    }
 51
 52    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
 53        self.0[..hw]
 54            .iter_mut()
 55            .for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
 56        self.0.shuffle(source);
 57    }
 58}
 59
 60pub struct SvpPPol(pub *mut svp::svp_ppol_t, pub usize);
 61
 62/// A prepared [crate::Scalar] for [SvpPPolOps::svp_apply_dft].
 63/// An [SvpPPol] an be seen as a [VecZnxDft] of one limb.
 64/// The backend array of an [SvpPPol] is allocated in C and must be freed manually.
 65impl SvpPPol {
 66    /// Returns the ring degree of the [SvpPPol].
 67    pub fn n(&self) -> usize {
 68        self.1
 69    }
 70
 71    /// Returns the number of limbs of the [SvpPPol], which is always 1.
 72    pub fn limbs(&self) -> usize {
 73        1
 74    }
 75}
 76
 77pub trait SvpPPolOps {
 78    /// Prepares a [crate::Scalar] for a [SvpPPolOps::svp_apply_dft].
 79    fn svp_prepare(&self, svp_ppol: &mut SvpPPol, a: &Scalar);
 80
 81    /// Allocates a new [SvpPPol].
 82    fn svp_new_ppol(&self) -> SvpPPol;
 83
 84    /// Applies the [SvpPPol] x [VecZnxDft] product, where each limb of
 85    /// the [VecZnxDft] is multiplied with [SvpPPol].
 86    fn svp_apply_dft(&self, c: &mut VecZnxDft, a: &SvpPPol, b: &VecZnx);
 87}
 88
 89impl SvpPPolOps for Module {
 90    fn svp_prepare(&self, svp_ppol: &mut SvpPPol, a: &Scalar) {
 91        unsafe { svp::svp_prepare(self.0, svp_ppol.0, a.as_ptr()) }
 92    }
 93
 94    fn svp_new_ppol(&self) -> SvpPPol {
 95        unsafe { SvpPPol(svp::new_svp_ppol(self.0), self.n()) }
 96    }
 97
 98    fn svp_apply_dft(&self, c: &mut VecZnxDft, a: &SvpPPol, b: &VecZnx) {
 99        let limbs: u64 = b.limbs() as u64;
100        assert!(
101            c.limbs() as u64 >= limbs,
102            "invalid c_vector: c_vector.limbs()={} < b.limbs()={}",
103            c.limbs(),
104            limbs
105        );
106        unsafe { svp::svp_apply_dft(self.0, c.0, limbs, a.0, b.as_ptr(), limbs, b.n() as u64) }
107    }
108}