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