1use crate::ffi::vec_znx_big;
 2use crate::ffi::vec_znx_dft;
 3use crate::{Module, VecZnxBig};
 4
 5pub struct VecZnxDft(pub *mut vec_znx_dft::vec_znx_dft_t, pub usize);
 6
 7impl VecZnxDft {
 8    pub fn as_vec_znx_big(&mut self) -> VecZnxBig {
 9        VecZnxBig(self.0 as *mut vec_znx_big::vec_znx_bigcoeff_t, self.1)
10    }
11    pub fn limbs(&self) -> usize {
12        self.1
13    }
14}
15
16impl Module {
17    // Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
18    pub fn new_vec_znx_dft(&self, limbs: usize) -> VecZnxDft {
19        unsafe { VecZnxDft(vec_znx_dft::new_vec_znx_dft(self.0, limbs as u64), limbs) }
20    }
21
22    // b <- IDFT(a), uses a as scratch space.
23    pub fn vec_znx_idft_tmp_a(&self, b: &mut VecZnxBig, a: &mut VecZnxDft, a_limbs: usize) {
24        assert!(
25            b.limbs() >= a_limbs,
26            "invalid c_vector: b_vector.limbs()={} < a_limbs={}",
27            b.limbs(),
28            a_limbs
29        );
30        unsafe { vec_znx_dft::vec_znx_idft_tmp_a(self.0, b.0, a_limbs as u64, a.0, a_limbs as u64) }
31    }
32
33    // Returns the size of the scratch space for [vec_znx_idft].
34    pub fn vec_znx_idft_tmp_bytes(&self) -> usize {
35        unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.0) as usize }
36    }
37
38    // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
39    pub fn vec_znx_idft(
40        &self,
41        b_vector: &mut VecZnxBig,
42        a_vector: &mut VecZnxDft,
43        a_limbs: usize,
44        tmp_bytes: &mut [u8],
45    ) {
46        assert!(
47            b_vector.limbs() >= a_limbs,
48            "invalid c_vector: b_vector.limbs()={} < a_limbs={}",
49            b_vector.limbs(),
50            a_limbs
51        );
52        assert!(
53            a_vector.limbs() >= a_limbs,
54            "invalid c_vector: c_vector.limbs()={} < a_limbs={}",
55            a_vector.limbs(),
56            a_limbs
57        );
58        assert!(
59            tmp_bytes.len() <= self.vec_znx_idft_tmp_bytes(),
60            "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
61            tmp_bytes.len(),
62            self.vec_znx_idft_tmp_bytes()
63        );
64        unsafe {
65            vec_znx_dft::vec_znx_idft(
66                self.0,
67                b_vector.0,
68                a_limbs as u64,
69                a_vector.0,
70                a_limbs as u64,
71                tmp_bytes.as_mut_ptr(),
72            )
73        }
74    }
75}