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