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