1use std::marker::PhantomData;
  2
  3use crate::ffi::svp::{self, svp_ppol_t};
  4use crate::ffi::vec_znx_dft::vec_znx_dft_t;
  5use crate::{Backend, FFT64, Module, VecZnx, VecZnxDft, ZnxLayout, assert_alignement};
  6
  7use crate::{ZnxInfos, alloc_aligned, cast_mut};
  8use rand::seq::SliceRandom;
  9use rand_core::RngCore;
 10use rand_distr::{Distribution, weighted::WeightedIndex};
 11use sampling::source::Source;
 12
 13pub struct Scalar {
 14    pub n: usize,
 15    pub data: Vec<i64>,
 16    pub ptr: *mut i64,
 17}
 18
 19impl<B: Backend> Module<B> {
 20    pub fn new_scalar(&self) -> Scalar {
 21        Scalar::new(self.n())
 22    }
 23}
 24
 25impl Scalar {
 26    pub fn new(n: usize) -> Self {
 27        let mut data: Vec<i64> = alloc_aligned::<i64>(n);
 28        let ptr: *mut i64 = data.as_mut_ptr();
 29        Self {
 30            n: n,
 31            data: data,
 32            ptr: ptr,
 33        }
 34    }
 35
 36    pub fn n(&self) -> usize {
 37        self.n
 38    }
 39
 40    pub fn bytes_of(n: usize) -> usize {
 41        n * std::mem::size_of::<i64>()
 42    }
 43
 44    pub fn from_bytes(n: usize, bytes: &mut [u8]) -> Self {
 45        let size: usize = Self::bytes_of(n);
 46        debug_assert!(
 47            bytes.len() == size,
 48            "invalid buffer: bytes.len()={} < self.bytes_of(n={})={}",
 49            bytes.len(),
 50            n,
 51            size
 52        );
 53        #[cfg(debug_assertions)]
 54        {
 55            assert_alignement(bytes.as_ptr())
 56        }
 57        unsafe {
 58            let bytes_i64: &mut [i64] = cast_mut::<u8, i64>(bytes);
 59            let ptr: *mut i64 = bytes_i64.as_mut_ptr();
 60            Self {
 61                n: n,
 62                data: Vec::from_raw_parts(bytes_i64.as_mut_ptr(), bytes.len(), bytes.len()),
 63                ptr: ptr,
 64            }
 65        }
 66    }
 67
 68    pub fn from_bytes_borrow(n: usize, bytes: &mut [u8]) -> Self {
 69        let size: usize = Self::bytes_of(n);
 70        debug_assert!(
 71            bytes.len() == size,
 72            "invalid buffer: bytes.len()={} < self.bytes_of(n={})={}",
 73            bytes.len(),
 74            n,
 75            size
 76        );
 77        #[cfg(debug_assertions)]
 78        {
 79            assert_alignement(bytes.as_ptr())
 80        }
 81        let bytes_i64: &mut [i64] = cast_mut::<u8, i64>(bytes);
 82        let ptr: *mut i64 = bytes_i64.as_mut_ptr();
 83        Self {
 84            n: n,
 85            data: Vec::new(),
 86            ptr: ptr,
 87        }
 88    }
 89
 90    pub fn as_ptr(&self) -> *const i64 {
 91        self.ptr
 92    }
 93
 94    pub fn raw(&self) -> &[i64] {
 95        unsafe { std::slice::from_raw_parts(self.ptr, self.n) }
 96    }
 97
 98    pub fn raw_mut(&self) -> &mut [i64] {
 99        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.n) }
100    }
101
102    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
103        let choices: [i64; 3] = [-1, 0, 1];
104        let weights: [f64; 3] = [prob / 2.0, 1.0 - prob, prob / 2.0];
105        let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
106        self.data
107            .iter_mut()
108            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
109    }
110
111    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
112        assert!(hw <= self.n());
113        self.data[..hw]
114            .iter_mut()
115            .for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
116        self.data.shuffle(source);
117    }
118
119    pub fn as_vec_znx(&self) -> VecZnx {
120        VecZnx {
121            n: self.n,
122            cols: 1,
123            limbs: 1,
124            data: Vec::new(),
125            ptr: self.ptr,
126        }
127    }
128}
129
130pub trait ScalarOps {
131    fn bytes_of_scalar(&self) -> usize;
132    fn new_scalar(&self) -> Scalar;
133    fn new_scalar_from_bytes(&self, bytes: &mut [u8]) -> Scalar;
134    fn new_scalar_from_bytes_borrow(&self, tmp_bytes: &mut [u8]) -> Scalar;
135}
136impl<B: Backend> ScalarOps for Module<B> {
137    fn bytes_of_scalar(&self) -> usize {
138        Scalar::bytes_of(self.n())
139    }
140    fn new_scalar(&self) -> Scalar {
141        Scalar::new(self.n())
142    }
143    fn new_scalar_from_bytes(&self, bytes: &mut [u8]) -> Scalar {
144        Scalar::from_bytes(self.n(), bytes)
145    }
146    fn new_scalar_from_bytes_borrow(&self, tmp_bytes: &mut [u8]) -> Scalar {
147        Scalar::from_bytes_borrow(self.n(), tmp_bytes)
148    }
149}
150
151pub struct ScalarZnxDft<B: Backend> {
152    pub n: usize,
153    pub data: Vec<u8>,
154    pub ptr: *mut u8,
155    _marker: PhantomData<B>,
156}
157
158/// A prepared [crate::Scalar] for [SvpPPolOps::svp_apply_dft].
159/// An [SvpPPol] an be seen as a [VecZnxDft] of one limb.
160impl ScalarZnxDft<FFT64> {
161    pub fn new(module: &Module<FFT64>) -> Self {
162        module.new_svp_ppol()
163    }
164
165    /// Returns the ring degree of the [SvpPPol].
166    pub fn n(&self) -> usize {
167        self.n
168    }
169
170    pub fn bytes_of(module: &Module<FFT64>) -> usize {
171        module.bytes_of_svp_ppol()
172    }
173
174    pub fn from_bytes(module: &Module<FFT64>, bytes: &mut [u8]) -> Self {
175        #[cfg(debug_assertions)]
176        {
177            assert_alignement(bytes.as_ptr());
178            assert_eq!(bytes.len(), module.bytes_of_svp_ppol());
179        }
180        unsafe {
181            Self {
182                n: module.n(),
183                data: Vec::from_raw_parts(bytes.as_mut_ptr(), bytes.len(), bytes.len()),
184                ptr: bytes.as_mut_ptr(),
185                _marker: PhantomData,
186            }
187        }
188    }
189
190    pub fn from_bytes_borrow(module: &Module<FFT64>, tmp_bytes: &mut [u8]) -> Self {
191        #[cfg(debug_assertions)]
192        {
193            assert_alignement(tmp_bytes.as_ptr());
194            assert_eq!(tmp_bytes.len(), module.bytes_of_svp_ppol());
195        }
196        Self {
197            n: module.n(),
198            data: Vec::new(),
199            ptr: tmp_bytes.as_mut_ptr(),
200            _marker: PhantomData,
201        }
202    }
203
204    /// Returns the number of cols of the [SvpPPol], which is always 1.
205    pub fn cols(&self) -> usize {
206        1
207    }
208}
209
210pub trait ScalarZnxDftOps<B: Backend> {
211    /// Allocates a new [SvpPPol].
212    fn new_svp_ppol(&self) -> ScalarZnxDft<B>;
213
214    /// Returns the minimum number of bytes necessary to allocate
215    /// a new [SvpPPol] through [SvpPPol::from_bytes] ro.
216    fn bytes_of_svp_ppol(&self) -> usize;
217
218    /// Allocates a new [SvpPPol] from an array of bytes.
219    /// The array of bytes is owned by the [SvpPPol].
220    /// The method will panic if bytes.len() < [SvpPPolOps::bytes_of_svp_ppol]
221    fn new_svp_ppol_from_bytes(&self, bytes: &mut [u8]) -> ScalarZnxDft<B>;
222
223    /// Allocates a new [SvpPPol] from an array of bytes.
224    /// The array of bytes is borrowed by the [SvpPPol].
225    /// The method will panic if bytes.len() < [SvpPPolOps::bytes_of_svp_ppol]
226    fn new_svp_ppol_from_bytes_borrow(&self, tmp_bytes: &mut [u8]) -> ScalarZnxDft<B>;
227
228    /// Prepares a [crate::Scalar] for a [SvpPPolOps::svp_apply_dft].
229    fn svp_prepare(&self, svp_ppol: &mut ScalarZnxDft<B>, a: &Scalar);
230
231    /// Applies the [SvpPPol] x [VecZnxDft] product, where each limb of
232    /// the [VecZnxDft] is multiplied with [SvpPPol].
233    fn svp_apply_dft(&self, c: &mut VecZnxDft<B>, a: &ScalarZnxDft<B>, b: &VecZnx);
234}
235
236impl ScalarZnxDftOps<FFT64> for Module<FFT64> {
237    fn new_svp_ppol(&self) -> ScalarZnxDft<FFT64> {
238        let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_svp_ppol());
239        let ptr: *mut u8 = data.as_mut_ptr();
240        ScalarZnxDft::<FFT64> {
241            data: data,
242            ptr: ptr,
243            n: self.n(),
244            _marker: PhantomData,
245        }
246    }
247
248    fn bytes_of_svp_ppol(&self) -> usize {
249        unsafe { svp::bytes_of_svp_ppol(self.ptr) as usize }
250    }
251
252    fn new_svp_ppol_from_bytes(&self, bytes: &mut [u8]) -> ScalarZnxDft<FFT64> {
253        ScalarZnxDft::from_bytes(self, bytes)
254    }
255
256    fn new_svp_ppol_from_bytes_borrow(&self, tmp_bytes: &mut [u8]) -> ScalarZnxDft<FFT64> {
257        ScalarZnxDft::from_bytes_borrow(self, tmp_bytes)
258    }
259
260    fn svp_prepare(&self, svp_ppol: &mut ScalarZnxDft<FFT64>, a: &Scalar) {
261        unsafe { svp::svp_prepare(self.ptr, svp_ppol.ptr as *mut svp_ppol_t, a.as_ptr()) }
262    }
263
264    fn svp_apply_dft(&self, c: &mut VecZnxDft<FFT64>, a: &ScalarZnxDft<FFT64>, b: &VecZnx) {
265        unsafe {
266            svp::svp_apply_dft(
267                self.ptr,
268                c.ptr as *mut vec_znx_dft_t,
269                c.cols() as u64,
270                a.ptr as *const svp_ppol_t,
271                b.as_ptr(),
272                b.cols() as u64,
273                b.n() as u64,
274            )
275        }
276    }
277}