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