1use crate::ffi::vec_znx_big;
  2use crate::ffi::vec_znx_dft;
  3use crate::znx_base::ZnxAlloc;
  4use crate::znx_base::ZnxInfos;
  5use crate::znx_base::ZnxLayout;
  6use crate::znx_base::ZnxSliceSize;
  7use crate::{Backend, FFT64, Module, VecZnx, VecZnxBig, VecZnxDft, assert_alignement};
  8
  9pub trait VecZnxDftOps<B: Backend> {
 10    /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
 11    fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDft<B>;
 12
 13    /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
 14    ///
 15    /// Behavior: takes ownership of the backing array.
 16    ///
 17    /// # Arguments
 18    ///
 19    /// * `cols`: the number of cols of the [VecZnxDft].
 20    /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
 21    ///
 22    /// # Panics
 23    /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
 24    fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDft<B>;
 25
 26    /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
 27    ///
 28    /// Behavior: the backing array is only borrowed.
 29    ///
 30    /// # Arguments
 31    ///
 32    /// * `cols`: the number of cols of the [VecZnxDft].
 33    /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
 34    ///
 35    /// # Panics
 36    /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
 37    fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnxDft<B>;
 38
 39    /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
 40    ///
 41    /// # Arguments
 42    ///
 43    /// * `cols`: the number of cols of the [VecZnxDft].
 44    /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
 45    ///
 46    /// # Panics
 47    /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
 48    fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize;
 49
 50    /// Returns the minimum number of bytes necessary to allocate
 51    /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
 52    fn vec_znx_idft_tmp_bytes(&self) -> usize;
 53
 54    /// b <- IDFT(a), uses a as scratch space.
 55    fn vec_znx_idft_tmp_a(&self, res: &mut VecZnxBig<B>, res_col: usize, a: &mut VecZnxDft<B>, a_cols: usize);
 56
 57    fn vec_znx_idft(&self, res: &mut VecZnxBig<B>, res_col: usize, a: &VecZnxDft<B>, a_col: usize, tmp_bytes: &mut [u8]);
 58
 59    fn vec_znx_dft(&self, res: &mut VecZnxDft<B>, res_col: usize, a: &VecZnx, a_col: usize);
 60}
 61
 62impl VecZnxDftOps<FFT64> for Module<FFT64> {
 63    fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDft<FFT64> {
 64        VecZnxDft::<FFT64>::new(&self, 1, cols, size)
 65    }
 66
 67    fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDft<FFT64> {
 68        VecZnxDft::from_bytes(self, 1, cols, size, bytes)
 69    }
 70
 71    fn new_vec_znx_dft_from_bytes_borrow(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnxDft<FFT64> {
 72        VecZnxDft::from_bytes_borrow(self, 1, cols, size, bytes)
 73    }
 74
 75    fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize {
 76        VecZnxDft::bytes_of(&self, 1, cols, size)
 77    }
 78
 79    fn vec_znx_idft_tmp_a(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &mut VecZnxDft<FFT64>, a_col: usize) {
 80        #[cfg(debug_assertions)]
 81        {
 82            assert_eq!(res.poly_count(), a.poly_count());
 83        }
 84
 85        unsafe {
 86            vec_znx_dft::vec_znx_idft_tmp_a(
 87                self.ptr,
 88                res.at_mut_ptr(res_col * res.size(), 0) as *mut vec_znx_big::vec_znx_big_t,
 89                res.size() as u64,
 90                a.at_ptr(a_col * a.size(), 0) as *mut vec_znx_dft::vec_znx_dft_t,
 91                a.size() as u64,
 92            )
 93        }
 94    }
 95
 96    fn vec_znx_idft_tmp_bytes(&self) -> usize {
 97        unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
 98    }
 99
100    /// b <- DFT(a)
101    ///
102    /// # Panics
103    /// If b.cols < a_cols
104    fn vec_znx_dft(&self, res: &mut VecZnxDft<FFT64>, res_col: usize, a: &VecZnx, a_col: usize) {
105        unsafe {
106            vec_znx_dft::vec_znx_dft(
107                self.ptr,
108                res.at_mut_ptr(res_col * res.size(), 0) as *mut vec_znx_dft::vec_znx_dft_t,
109                res.size() as u64,
110                a.at_ptr(a_col, 0),
111                a.size() as u64,
112                a.sl() as u64,
113            )
114        }
115    }
116
117    // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
118    fn vec_znx_idft(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnxDft<FFT64>, a_col: usize, tmp_bytes: &mut [u8]) {
119        #[cfg(debug_assertions)]
120        {
121            assert!(
122                tmp_bytes.len() >= Self::vec_znx_idft_tmp_bytes(self),
123                "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
124                tmp_bytes.len(),
125                Self::vec_znx_idft_tmp_bytes(self)
126            );
127            assert_alignement(tmp_bytes.as_ptr())
128        }
129        unsafe {
130            vec_znx_dft::vec_znx_idft(
131                self.ptr,
132                res.at_mut_ptr(res_col * res.size(), 0) as *mut vec_znx_big::vec_znx_big_t,
133                res.size() as u64,
134                a.at_ptr(a_col * res.size(), 0) as *const vec_znx_dft::vec_znx_dft_t,
135                a.size() as u64,
136                tmp_bytes.as_mut_ptr(),
137            )
138        }
139    }
140}