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