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