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, 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 the minimum number of bytes necessary to allocate
34    /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
35    pub fn bytes_of_vec_znx_dft(&self, limbs: usize) -> usize {
36        unsafe { bytes_of_vec_znx_dft(self.0, limbs as u64) as usize }
37    }
38
39    // b <- IDFT(a), uses a as scratch space.
40    pub fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig, a: &mut VecZnxDft, a_limbs: usize) {
41        assert!(
42            b.limbs() >= a_limbs,
43            "invalid c_vector: b_vector.limbs()={} < a_limbs={}",
44            b.limbs(),
45            a_limbs
46        );
47        unsafe { vec_znx_dft::vec_znx_idft_tmp_a(self.0, b.0, a_limbs as u64, a.0, a_limbs as u64) }
48    }
49
50    // Returns the size of the scratch space for [vec_znx_idft].
51    pub fn vec_znx_idft_tmp_bytes(&self) -> usize {
52        unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.0) as usize }
53    }
54
55    // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
56    pub fn vec_znx_idft(
57        &self,
58        b_vector: &mut VecZnxBig,
59        a_vector: &mut VecZnxDft,
60        a_limbs: usize,
61        tmp_bytes: &mut [u8],
62    ) {
63        assert!(
64            b_vector.limbs() >= a_limbs,
65            "invalid c_vector: b_vector.limbs()={} < a_limbs={}",
66            b_vector.limbs(),
67            a_limbs
68        );
69        assert!(
70            a_vector.limbs() >= a_limbs,
71            "invalid c_vector: c_vector.limbs()={} < a_limbs={}",
72            a_vector.limbs(),
73            a_limbs
74        );
75        assert!(
76            tmp_bytes.len() <= self.vec_znx_idft_tmp_bytes(),
77            "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
78            tmp_bytes.len(),
79            self.vec_znx_idft_tmp_bytes()
80        );
81        unsafe {
82            vec_znx_dft::vec_znx_idft(
83                self.0,
84                b_vector.0,
85                a_limbs as u64,
86                a_vector.0,
87                a_limbs as u64,
88                tmp_bytes.as_mut_ptr(),
89            )
90        }
91    }
92}