1use crate::ffi::vec_znx_big;
  2use crate::ffi::vec_znx_dft;
  3use crate::ffi::vec_znx_dft::bytes_of_vec_znx_dft;
  4use crate::{assert_alignement, Infos, Module, VecZnxApi, VecZnxBig};
  5
  6pub struct VecZnxDft(pub *mut vec_znx_dft::vec_znx_dft_t, pub usize);
  7
  8impl VecZnxDft {
  9    /// Returns a new [VecZnxDft] with the provided data as backing array.
 10    /// User must ensure that data is properly alligned and that
 11    /// the size of data is at least equal to [Module::bytes_of_vec_znx_dft].
 12    pub fn from_bytes(cols: usize, tmp_bytes: &mut [u8]) -> VecZnxDft {
 13        #[cfg(debug_assertions)]
 14        {
 15            assert_alignement(tmp_bytes.as_ptr())
 16        }
 17        VecZnxDft(
 18            tmp_bytes.as_mut_ptr() as *mut vec_znx_dft::vec_znx_dft_t,
 19            cols,
 20        )
 21    }
 22
 23    /// Cast a [VecZnxDft] into a [VecZnxBig].
 24    /// The returned [VecZnxBig] shares the backing array
 25    /// with the original [VecZnxDft].
 26    pub fn as_vec_znx_big(&mut self) -> VecZnxBig {
 27        VecZnxBig(self.0 as *mut vec_znx_big::vec_znx_bigcoeff_t, self.1)
 28    }
 29    pub fn cols(&self) -> usize {
 30        self.1
 31    }
 32}
 33
 34pub trait VecZnxDftOps {
 35    /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
 36    fn new_vec_znx_dft(&self, cols: usize) -> VecZnxDft;
 37
 38    /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
 39    ///
 40    /// # Arguments
 41    ///
 42    /// * `cols`: the number of cols of the [VecZnxDft].
 43    /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
 44    ///
 45    /// # Panics
 46    /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
 47    fn new_vec_znx_dft_from_bytes(&self, cols: usize, bytes: &mut [u8]) -> VecZnxDft;
 48
 49    /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
 50    ///
 51    /// # Arguments
 52    ///
 53    /// * `cols`: the number of cols of the [VecZnxDft].
 54    /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
 55    ///
 56    /// # Panics
 57    /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
 58    fn bytes_of_vec_znx_dft(&self, cols: usize) -> usize;
 59
 60    /// Returns the minimum number of bytes necessary to allocate
 61    /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
 62    fn vec_znx_idft_tmp_bytes(&self) -> usize;
 63
 64    /// b <- IDFT(a), uses a as scratch space.
 65    fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig, a: &mut VecZnxDft, a_limbs: usize);
 66
 67    fn vec_znx_idft(
 68        &self,
 69        b: &mut VecZnxBig,
 70        a: &mut VecZnxDft,
 71        a_limbs: usize,
 72        tmp_bytes: &mut [u8],
 73    );
 74
 75    fn vec_znx_dft<T: VecZnxApi + Infos>(&self, b: &mut VecZnxDft, a: &T, a_limbs: usize);
 76}
 77
 78impl VecZnxDftOps for Module {
 79    fn new_vec_znx_dft(&self, cols: usize) -> VecZnxDft {
 80        unsafe { VecZnxDft(vec_znx_dft::new_vec_znx_dft(self.0, cols as u64), cols) }
 81    }
 82
 83    fn new_vec_znx_dft_from_bytes(&self, cols: usize, tmp_bytes: &mut [u8]) -> VecZnxDft {
 84        debug_assert!(
 85            tmp_bytes.len() >= <Module as VecZnxDftOps>::bytes_of_vec_znx_dft(self, cols),
 86            "invalid bytes: bytes.len()={} < bytes_of_vec_znx_dft={}",
 87            tmp_bytes.len(),
 88            <Module as VecZnxDftOps>::bytes_of_vec_znx_dft(self, cols)
 89        );
 90        #[cfg(debug_assertions)]
 91        {
 92            assert_alignement(tmp_bytes.as_ptr())
 93        }
 94        VecZnxDft::from_bytes(cols, tmp_bytes)
 95    }
 96
 97    fn bytes_of_vec_znx_dft(&self, cols: usize) -> usize {
 98        unsafe { bytes_of_vec_znx_dft(self.0, cols as u64) as usize }
 99    }
100
101    fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig, a: &mut VecZnxDft, a_limbs: usize) {
102        debug_assert!(
103            b.cols() >= a_limbs,
104            "invalid c_vector: b_vector.cols()={} < a_limbs={}",
105            b.cols(),
106            a_limbs
107        );
108        unsafe {
109            vec_znx_dft::vec_znx_idft_tmp_a(self.0, b.0, b.cols() as u64, a.0, a_limbs as u64)
110        }
111    }
112
113    fn vec_znx_idft_tmp_bytes(&self) -> usize {
114        unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.0) as usize }
115    }
116
117    /// b <- DFT(a)
118    ///
119    /// # Panics
120    /// If b.cols < a_cols
121    fn vec_znx_dft<T: VecZnxApi + Infos>(&self, b: &mut VecZnxDft, a: &T, a_cols: usize) {
122        debug_assert!(
123            b.cols() >= a_cols,
124            "invalid a_cols: b.cols()={} < a_cols={}",
125            b.cols(),
126            a_cols
127        );
128        unsafe {
129            vec_znx_dft::vec_znx_dft(
130                self.0,
131                b.0,
132                b.cols() as u64,
133                a.as_ptr(),
134                a_cols as u64,
135                a.n() as u64,
136            )
137        }
138    }
139
140    // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
141    fn vec_znx_idft(
142        &self,
143        b: &mut VecZnxBig,
144        a: &mut VecZnxDft,
145        a_cols: usize,
146        tmp_bytes: &mut [u8],
147    ) {
148        debug_assert!(
149            b.cols() >= a_cols,
150            "invalid c_vector: b.cols()={} < a_cols={}",
151            b.cols(),
152            a_cols
153        );
154        debug_assert!(
155            a.cols() >= a_cols,
156            "invalid c_vector: a.cols()={} < a_cols={}",
157            a.cols(),
158            a_cols
159        );
160        debug_assert!(
161            tmp_bytes.len() <= <Module as VecZnxDftOps>::vec_znx_idft_tmp_bytes(self),
162            "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
163            tmp_bytes.len(),
164            <Module as VecZnxDftOps>::vec_znx_idft_tmp_bytes(self)
165        );
166        #[cfg(debug_assertions)]
167        {
168            assert_alignement(tmp_bytes.as_ptr())
169        }
170        unsafe {
171            vec_znx_dft::vec_znx_idft(
172                self.0,
173                b.0,
174                a.cols() as u64,
175                a.0,
176                a_cols as u64,
177                tmp_bytes.as_mut_ptr(),
178            )
179        }
180    }
181}