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