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