1use crate::ffi::svp;
  2use crate::{alias_mut_slice_to_vec, assert_alignement, Module, VecZnxApi, VecZnxDft};
  3
  4use crate::{alloc_aligned, cast, 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(alloc_aligned::<i64>(n))
 21    }
 22
 23    pub fn n(&self) -> usize {
 24        self.0.len()
 25    }
 26
 27    pub fn buffer_size(n: usize) -> usize {
 28        n
 29    }
 30
 31    pub fn from_buffer(&mut self, n: usize, buf: &mut [u8]) {
 32        let size: usize = Self::buffer_size(n);
 33        debug_assert!(
 34            buf.len() >= size,
 35            "invalid buffer: buf.len()={} < self.buffer_size(n={})={}",
 36            buf.len(),
 37            n,
 38            size
 39        );
 40        #[cfg(debug_assertions)]
 41        {
 42            assert_alignement(buf.as_ptr())
 43        }
 44        self.0 = alias_mut_slice_to_vec(cast::<u8, i64>(&buf[..size]))
 45    }
 46
 47    pub fn as_ptr(&self) -> *const i64 {
 48        self.0.as_ptr()
 49    }
 50
 51    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
 52        let choices: [i64; 3] = [-1, 0, 1];
 53        let weights: [f64; 3] = [prob / 2.0, 1.0 - prob, prob / 2.0];
 54        let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
 55        self.0
 56            .iter_mut()
 57            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
 58    }
 59
 60    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
 61        assert!(hw <= self.n());
 62        self.0[..hw]
 63            .iter_mut()
 64            .for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
 65        self.0.shuffle(source);
 66    }
 67}
 68
 69pub struct SvpPPol(pub *mut svp::svp_ppol_t, pub usize);
 70
 71/// A prepared [crate::Scalar] for [SvpPPolOps::svp_apply_dft].
 72/// An [SvpPPol] an be seen as a [VecZnxDft] of one limb.
 73/// The backend array of an [SvpPPol] is allocated in C and must be freed manually.
 74impl SvpPPol {
 75    /// Returns the ring degree of the [SvpPPol].
 76    pub fn n(&self) -> usize {
 77        self.1
 78    }
 79
 80    pub fn from_bytes(size: usize, bytes: &mut [u8]) -> SvpPPol {
 81        #[cfg(debug_assertions)]
 82        {
 83            assert_alignement(bytes.as_ptr())
 84        }
 85        debug_assert!(bytes.len() << 3 >= size);
 86        SvpPPol(bytes.as_mut_ptr() as *mut svp::svp_ppol_t, size)
 87    }
 88
 89    /// Returns the number of cols of the [SvpPPol], which is always 1.
 90    pub fn cols(&self) -> usize {
 91        1
 92    }
 93}
 94
 95pub trait SvpPPolOps {
 96    /// Allocates a new [SvpPPol].
 97    fn new_svp_ppol(&self) -> SvpPPol;
 98
 99    /// Returns the minimum number of bytes necessary to allocate
100    /// a new [SvpPPol] through [SvpPPol::from_bytes].
101    fn bytes_of_svp_ppol(&self) -> usize;
102
103    /// Prepares a [crate::Scalar] for a [SvpPPolOps::svp_apply_dft].
104    fn svp_prepare(&self, svp_ppol: &mut SvpPPol, a: &Scalar);
105
106    /// Applies the [SvpPPol] x [VecZnxDft] product, where each limb of
107    /// the [VecZnxDft] is multiplied with [SvpPPol].
108    fn svp_apply_dft<T: VecZnxApi + Infos>(
109        &self,
110        c: &mut VecZnxDft,
111        a: &SvpPPol,
112        b: &T,
113        b_cols: usize,
114    );
115}
116
117impl SvpPPolOps for Module {
118    fn new_svp_ppol(&self) -> SvpPPol {
119        unsafe { SvpPPol(svp::new_svp_ppol(self.0), self.n()) }
120    }
121
122    fn bytes_of_svp_ppol(&self) -> usize {
123        unsafe { svp::bytes_of_svp_ppol(self.0) as usize }
124    }
125
126    fn svp_prepare(&self, svp_ppol: &mut SvpPPol, a: &Scalar) {
127        unsafe { svp::svp_prepare(self.0, svp_ppol.0, a.as_ptr()) }
128    }
129
130    fn svp_apply_dft<T: VecZnxApi + Infos>(
131        &self,
132        c: &mut VecZnxDft,
133        a: &SvpPPol,
134        b: &T,
135        b_cols: usize,
136    ) {
137        debug_assert!(
138            c.cols() >= b_cols,
139            "invalid c_vector: c_vector.cols()={} < b.cols()={}",
140            c.cols(),
141            b_cols
142        );
143        unsafe {
144            svp::svp_apply_dft(
145                self.0,
146                c.0,
147                b_cols as u64,
148                a.0,
149                b.as_ptr(),
150                b_cols as u64,
151                b.n() as u64,
152            )
153        }
154    }
155}