1use crate::ffi::vec_znx_big::vec_znx_big_t;
  2use crate::ffi::vec_znx_dft;
  3use crate::ffi::vec_znx_dft::{bytes_of_vec_znx_dft, vec_znx_dft_t};
  4use crate::{alloc_aligned, VecZnx, DEFAULTALIGN};
  5use crate::{assert_alignement, Infos, Module, VecZnxBig, BACKEND};
  6
  7pub struct VecZnxDft {
  8    pub data: Vec<u8>,
  9    pub ptr: *mut u8,
 10    pub n: usize,
 11    pub cols: usize,
 12    pub backend: BACKEND,
 13}
 14
 15impl VecZnxDft {
 16    /// Returns a new [VecZnxDft] with the provided data as backing array.
 17    /// User must ensure that data is properly alligned and that
 18    /// the size of data is at least equal to [Module::bytes_of_vec_znx_dft].
 19    pub fn from_bytes(module: &Module, cols: usize, bytes: &mut [u8]) -> VecZnxDft {
 20        #[cfg(debug_assertions)]
 21        {
 22            assert_eq!(bytes.len(), module.bytes_of_vec_znx_dft(cols));
 23            assert_alignement(bytes.as_ptr())
 24        }
 25        unsafe {
 26            VecZnxDft {
 27                data: Vec::from_raw_parts(bytes.as_mut_ptr(), bytes.len(), bytes.len()),
 28                ptr: bytes.as_mut_ptr(),
 29                n: module.n(),
 30                cols: cols,
 31                backend: module.backend,
 32            }
 33        }
 34    }
 35
 36    pub fn from_bytes_borrow(module: &Module, cols: usize, bytes: &mut [u8]) -> VecZnxDft {
 37        #[cfg(debug_assertions)]
 38        {
 39            assert_eq!(bytes.len(), module.bytes_of_vec_znx_dft(cols));
 40            assert_alignement(bytes.as_ptr());
 41        }
 42        VecZnxDft {
 43            data: Vec::new(),
 44            ptr: bytes.as_mut_ptr(),
 45            n: module.n(),
 46            cols: cols,
 47            backend: module.backend,
 48        }
 49    }
 50
 51    /// Cast a [VecZnxDft] into a [VecZnxBig].
 52    /// The returned [VecZnxBig] shares the backing array
 53    /// with the original [VecZnxDft].
 54    pub fn as_vec_znx_big(&mut self) -> VecZnxBig {
 55        VecZnxBig {
 56            data: Vec::new(),
 57            ptr: self.ptr,
 58            n: self.n,
 59            cols: self.cols,
 60            backend: self.backend,
 61        }
 62    }
 63
 64    pub fn n(&self) -> usize {
 65        self.n
 66    }
 67
 68    pub fn cols(&self) -> usize {
 69        self.cols
 70    }
 71
 72    pub fn backend(&self) -> BACKEND {
 73        self.backend
 74    }
 75
 76    /// Returns a non-mutable reference of `T` of the entire contiguous array of the [VecZnxDft].
 77    /// When using [`crate::FFT64`] as backend, `T` should be [f64].
 78    /// When using [`crate::NTT120`] as backend, `T` should be [i64].
 79    /// The length of the returned array is cols * n.
 80    pub fn raw<T>(&self, module: &Module) -> &[T] {
 81        let ptr: *const T = self.ptr as *const T;
 82        let len: usize = (self.cols() * module.n() * 8) / std::mem::size_of::<T>();
 83        unsafe { &std::slice::from_raw_parts(ptr, len) }
 84    }
 85
 86    pub fn at<T>(&self, module: &Module, col_i: usize) -> &[T] {
 87        &self.raw::<T>(module)[col_i * module.n()..(col_i + 1) * module.n()]
 88    }
 89
 90    /// Returns a mutable reference of `T` of the entire contiguous array of the [VecZnxDft].
 91    /// When using [`crate::FFT64`] as backend, `T` should be [f64].
 92    /// When using [`crate::NTT120`] as backend, `T` should be [i64].
 93    /// The length of the returned array is cols * n.
 94    pub fn raw_mut<T>(&self, module: &Module) -> &mut [T] {
 95        let ptr: *mut T = self.ptr as *mut T;
 96        let len: usize = (self.cols() * module.n() * 8) / std::mem::size_of::<T>();
 97        unsafe { std::slice::from_raw_parts_mut(ptr, len) }
 98    }
 99
100    pub fn at_mut<T>(&self, module: &Module, col_i: usize) -> &mut [T] {
101        &mut self.raw_mut::<T>(module)[col_i * module.n()..(col_i + 1) * module.n()]
102    }
103}
104
105pub trait VecZnxDftOps {
106    /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
107    fn new_vec_znx_dft(&self, cols: usize) -> VecZnxDft;
108
109    /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
110    ///
111    /// # Arguments
112    ///
113    /// * `cols`: the number of cols of the [VecZnxDft].
114    /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
115    ///
116    /// # Panics
117    /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
118    fn new_vec_znx_dft_from_bytes(&self, cols: usize, bytes: &mut [u8]) -> VecZnxDft;
119
120    /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
121    ///
122    /// # Arguments
123    ///
124    /// * `cols`: the number of cols of the [VecZnxDft].
125    /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
126    ///
127    /// # Panics
128    /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
129    fn bytes_of_vec_znx_dft(&self, cols: usize) -> usize;
130
131    /// Returns the minimum number of bytes necessary to allocate
132    /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
133    fn vec_znx_idft_tmp_bytes(&self) -> usize;
134
135    /// b <- IDFT(a), uses a as scratch space.
136    fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig, a: &mut VecZnxDft, a_cols: usize);
137
138    fn vec_znx_idft(&self, b: &mut VecZnxBig, a: &VecZnxDft, a_cols: usize, tmp_bytes: &mut [u8]);
139
140    fn vec_znx_dft(&self, b: &mut VecZnxDft, a: &VecZnx, a_cols: usize);
141
142    fn vec_znx_dft_automorphism(
143        &self,
144        k: i64,
145        b: &mut VecZnxDft,
146        b_cols: usize,
147        a: &VecZnxDft,
148        a_cols: usize,
149    );
150
151    fn vec_znx_dft_automorphism_inplace(
152        &self,
153        k: i64,
154        a: &mut VecZnxDft,
155        a_cols: usize,
156        tmp_bytes: &mut [u8],
157    );
158
159    fn vec_znx_dft_automorphism_tmp_bytes(&self) -> usize;
160}
161
162impl VecZnxDftOps for Module {
163    fn new_vec_znx_dft(&self, cols: usize) -> VecZnxDft {
164        let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_vec_znx_dft(cols));
165        let ptr: *mut u8 = data.as_mut_ptr();
166        VecZnxDft {
167            data: data,
168            ptr: ptr,
169            n: self.n(),
170            cols: cols,
171            backend: self.backend(),
172        }
173    }
174
175    fn new_vec_znx_dft_from_bytes(&self, cols: usize, tmp_bytes: &mut [u8]) -> VecZnxDft {
176        debug_assert!(
177            tmp_bytes.len() >= Self::bytes_of_vec_znx_dft(self, cols),
178            "invalid bytes: bytes.len()={} < bytes_of_vec_znx_dft={}",
179            tmp_bytes.len(),
180            Self::bytes_of_vec_znx_dft(self, cols)
181        );
182        #[cfg(debug_assertions)]
183        {
184            assert_alignement(tmp_bytes.as_ptr())
185        }
186        VecZnxDft::from_bytes(self, cols, tmp_bytes)
187    }
188
189    fn bytes_of_vec_znx_dft(&self, cols: usize) -> usize {
190        unsafe { bytes_of_vec_znx_dft(self.ptr, cols as u64) as usize }
191    }
192
193    fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig, a: &mut VecZnxDft, a_cols: usize) {
194        debug_assert!(
195            b.cols() >= a_cols,
196            "invalid c_vector: b_vector.cols()={} < a_cols={}",
197            b.cols(),
198            a_cols
199        );
200        unsafe {
201            vec_znx_dft::vec_znx_idft_tmp_a(
202                self.ptr,
203                b.ptr as *mut vec_znx_big_t,
204                b.cols() as u64,
205                a.ptr as *mut vec_znx_dft_t,
206                a_cols as u64,
207            )
208        }
209    }
210
211    fn vec_znx_idft_tmp_bytes(&self) -> usize {
212        unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
213    }
214
215    /// b <- DFT(a)
216    ///
217    /// # Panics
218    /// If b.cols < a_cols
219    fn vec_znx_dft(&self, b: &mut VecZnxDft, a: &VecZnx, a_cols: usize) {
220        debug_assert!(
221            b.cols() >= a_cols,
222            "invalid a_cols: b.cols()={} < a_cols={}",
223            b.cols(),
224            a_cols
225        );
226        unsafe {
227            vec_znx_dft::vec_znx_dft(
228                self.ptr,
229                b.ptr as *mut vec_znx_dft_t,
230                b.cols() as u64,
231                a.as_ptr(),
232                a_cols as u64,
233                a.n() as u64,
234            )
235        }
236    }
237
238    // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
239    fn vec_znx_idft(&self, b: &mut VecZnxBig, a: &VecZnxDft, a_cols: usize, tmp_bytes: &mut [u8]) {
240        #[cfg(debug_assertions)]
241        {
242            assert!(
243                b.cols() >= a_cols,
244                "invalid c_vector: b.cols()={} < a_cols={}",
245                b.cols(),
246                a_cols
247            );
248            assert!(
249                a.cols() >= a_cols,
250                "invalid c_vector: a.cols()={} < a_cols={}",
251                a.cols(),
252                a_cols
253            );
254            assert!(
255                tmp_bytes.len() >= Self::vec_znx_idft_tmp_bytes(self),
256                "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
257                tmp_bytes.len(),
258                Self::vec_znx_idft_tmp_bytes(self)
259            );
260            assert_alignement(tmp_bytes.as_ptr())
261        }
262        unsafe {
263            vec_znx_dft::vec_znx_idft(
264                self.ptr,
265                b.ptr as *mut vec_znx_big_t,
266                a.cols() as u64,
267                a.ptr as *const vec_znx_dft_t,
268                a_cols as u64,
269                tmp_bytes.as_mut_ptr(),
270            )
271        }
272    }
273
274    fn vec_znx_dft_automorphism(
275        &self,
276        k: i64,
277        b: &mut VecZnxDft,
278        b_cols: usize,
279        a: &VecZnxDft,
280        a_cols: usize,
281    ) {
282        #[cfg(debug_assertions)]
283        {
284            assert!(
285                b.cols() >= a_cols,
286                "invalid c_vector: b.cols()={} < a_cols={}",
287                b.cols(),
288                a_cols
289            );
290            assert!(
291                a.cols() >= a_cols,
292                "invalid c_vector: a.cols()={} < a_cols={}",
293                a.cols(),
294                a_cols
295            );
296        }
297        unsafe {
298            vec_znx_dft::vec_znx_dft_automorphism(
299                self.ptr,
300                k,
301                b.ptr as *mut vec_znx_dft_t,
302                b_cols as u64,
303                a.ptr as *const vec_znx_dft_t,
304                a_cols as u64,
305                [0u8; 0].as_mut_ptr(),
306            );
307        }
308    }
309
310    fn vec_znx_dft_automorphism_inplace(
311        &self,
312        k: i64,
313        a: &mut VecZnxDft,
314        a_cols: usize,
315        tmp_bytes: &mut [u8],
316    ) {
317        #[cfg(debug_assertions)]
318        {
319            assert!(
320                a.cols() >= a_cols,
321                "invalid c_vector: a.cols()={} < a_cols={}",
322                a.cols(),
323                a_cols
324            );
325            assert!(
326                tmp_bytes.len() >= Self::vec_znx_dft_automorphism_tmp_bytes(self),
327                "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_dft_automorphism_tmp_bytes()={}",
328                tmp_bytes.len(),
329                Self::vec_znx_dft_automorphism_tmp_bytes(self)
330            );
331            assert_alignement(tmp_bytes.as_ptr())
332        }
333        unsafe {
334            vec_znx_dft::vec_znx_dft_automorphism(
335                self.ptr,
336                k,
337                a.ptr as *mut vec_znx_dft_t,
338                a_cols as u64,
339                a.ptr as *const vec_znx_dft_t,
340                a_cols as u64,
341                tmp_bytes.as_mut_ptr(),
342            );
343        }
344    }
345
346    fn vec_znx_dft_automorphism_tmp_bytes(&self) -> usize {
347        unsafe {
348            std::cmp::max(
349                vec_znx_dft::vec_znx_dft_automorphism_tmp_bytes(self.ptr) as usize,
350                DEFAULTALIGN,
351            )
352        }
353    }
354}
355
356#[cfg(test)]
357mod tests {
358    use crate::{
359        alloc_aligned, Module, Sampling, VecZnx, VecZnxDft, VecZnxDftOps, VecZnxOps, BACKEND,
360    };
361    use itertools::izip;
362    use sampling::source::{new_seed, Source};
363
364    #[test]
365    fn test_automorphism_dft() {
366        let module: Module = Module::new(128, BACKEND::FFT64);
367
368        let cols: usize = 2;
369        let log_base2k: usize = 17;
370        let mut a: VecZnx = module.new_vec_znx(cols);
371        let mut a_dft: VecZnxDft = module.new_vec_znx_dft(cols);
372        let mut b_dft: VecZnxDft = module.new_vec_znx_dft(cols);
373
374        let mut source: Source = Source::new(new_seed());
375        module.fill_uniform(log_base2k, &mut a, cols, &mut source);
376
377        let mut tmp_bytes: Vec<u8> = alloc_aligned(module.vec_znx_dft_automorphism_tmp_bytes());
378
379        let p: i64 = -5;
380
381        // a_dft <- DFT(a)
382        module.vec_znx_dft(&mut a_dft, &a, cols);
383
384        // a_dft <- AUTO(a_dft)
385        module.vec_znx_dft_automorphism_inplace(p, &mut a_dft, cols, &mut tmp_bytes);
386
387        // a <- AUTO(a)
388        module.vec_znx_automorphism_inplace(p, &mut a, cols);
389
390        // b_dft <- DFT(AUTO(a))
391        module.vec_znx_dft(&mut b_dft, &a, cols);
392
393        let a_f64: &[f64] = a_dft.raw(&module);
394        let b_f64: &[f64] = b_dft.raw(&module);
395        izip!(a_f64.iter(), b_f64.iter()).for_each(|(ai, bi)| {
396            assert!((ai - bi).abs() <= 1e-9, "{:+e} > 1e-9", (ai - bi).abs());
397        });
398
399        module.free()
400    }
401}