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::{Module, VecZnx, 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(limbs: usize, data: &mut [u8]) -> VecZnxDft {
 13        VecZnxDft(data.as_mut_ptr() as *mut vec_znx_dft::vec_znx_dft_t, limbs)
 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 limbs(&self) -> usize {
 23        self.1
 24    }
 25}
 26
 27impl Module {
 28    // Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
 29    pub fn new_vec_znx_dft(&self, limbs: usize) -> VecZnxDft {
 30        unsafe { VecZnxDft(vec_znx_dft::new_vec_znx_dft(self.0, limbs as u64), limbs) }
 31    }
 32
 33    /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
 34    ///
 35    /// # Arguments
 36    ///
 37    /// * `limbs`: the number of limbs of the [VecZnxDft].
 38    /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
 39    ///
 40    /// # Panics
 41    /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
 42    pub fn new_vec_znx_dft_from_bytes(&self, limbs: usize, bytes: &mut [u8]) -> VecZnxDft {
 43        assert!(
 44            bytes.len() >= self.bytes_of_vec_znx_dft(limbs),
 45            "invalid bytes: bytes.len()={} < bytes_of_vec_znx_dft={}",
 46            bytes.len(),
 47            self.bytes_of_vec_znx_dft(limbs)
 48        );
 49        VecZnxDft::from_bytes(limbs, bytes)
 50    }
 51
 52    /// Returns the minimum number of bytes necessary to allocate
 53    /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
 54    pub fn bytes_of_vec_znx_dft(&self, limbs: usize) -> usize {
 55        unsafe { bytes_of_vec_znx_dft(self.0, limbs as u64) as usize }
 56    }
 57
 58    // b <- IDFT(a), uses a as scratch space.
 59    pub fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig, a: &mut VecZnxDft, a_limbs: usize) {
 60        assert!(
 61            b.limbs() >= a_limbs,
 62            "invalid c_vector: b_vector.limbs()={} < a_limbs={}",
 63            b.limbs(),
 64            a_limbs
 65        );
 66        unsafe {
 67            vec_znx_dft::vec_znx_idft_tmp_a(self.0, b.0, b.limbs() as u64, a.0, a_limbs as u64)
 68        }
 69    }
 70
 71    // Returns the size of the scratch space for [vec_znx_idft].
 72    pub fn vec_znx_idft_tmp_bytes(&self) -> usize {
 73        unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.0) as usize }
 74    }
 75
 76    /// b <- DFT(a)
 77    ///
 78    /// # Panics
 79    /// If b.limbs < a_limbs
 80    pub fn vec_znx_dft(&self, b: &mut VecZnxDft, a: &VecZnx, a_limbs: usize) {
 81        assert!(
 82            b.limbs() >= a_limbs,
 83            "invalid a_limbs: b.limbs()={} < a_limbs={}",
 84            b.limbs(),
 85            a_limbs
 86        );
 87        unsafe {
 88            vec_znx_dft::vec_znx_dft(
 89                self.0,
 90                b.0,
 91                b.limbs() as u64,
 92                a.as_ptr(),
 93                a_limbs as u64,
 94                a.n as u64,
 95            )
 96        }
 97    }
 98
 99    // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
100    pub fn vec_znx_idft(
101        &self,
102        b: &mut VecZnxBig,
103        a: &mut VecZnxDft,
104        a_limbs: usize,
105        tmp_bytes: &mut [u8],
106    ) {
107        assert!(
108            b.limbs() >= a_limbs,
109            "invalid c_vector: b.limbs()={} < a_limbs={}",
110            b.limbs(),
111            a_limbs
112        );
113        assert!(
114            a.limbs() >= a_limbs,
115            "invalid c_vector: a.limbs()={} < a_limbs={}",
116            a.limbs(),
117            a_limbs
118        );
119        assert!(
120            tmp_bytes.len() <= self.vec_znx_idft_tmp_bytes(),
121            "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
122            tmp_bytes.len(),
123            self.vec_znx_idft_tmp_bytes()
124        );
125        unsafe {
126            vec_znx_dft::vec_znx_idft(
127                self.0,
128                b.0,
129                a.limbs() as u64,
130                a.0,
131                a_limbs as u64,
132                tmp_bytes.as_mut_ptr(),
133            )
134        }
135    }
136}