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