1use crate::ffi::svp::{self, svp_ppol_t};
 2use crate::ffi::vec_znx_dft::vec_znx_dft_t;
 3use crate::znx_base::{ZnxInfos, ZnxView, ZnxViewMut};
 4use crate::{Backend, FFT64, Module, Scalar, ScalarZnxDft, ScalarZnxDftOwned, VecZnx, VecZnxDft};
 5
 6pub trait ScalarZnxDftAlloc<B> {
 7    fn new_scalar_znx_dft(&self, cols: usize) -> ScalarZnxDftOwned<B>;
 8    fn bytes_of_scalar_znx_dft(&self, cols: usize) -> usize;
 9    fn new_scalar_znx_dft_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxDftOwned<B>;
10    // fn new_scalar_znx_dft_from_bytes_borrow(&self, cols: usize, bytes: &mut [u8]) -> ScalarZnxDft<B>;
11}
12
13pub trait ScalarZnxDftOps<DataMut, Data, B: Backend> {
14    fn svp_prepare(&self, res: &mut ScalarZnxDft<DataMut, B>, res_col: usize, a: &Scalar<Data>, a_col: usize);
15    fn svp_apply_dft(
16        &self,
17        res: &mut VecZnxDft<DataMut, B>,
18        res_col: usize,
19        a: &ScalarZnxDft<Data, B>,
20        a_col: usize,
21        b: &VecZnx<Data>,
22        b_col: usize,
23    );
24}
25
26impl<B: Backend> ScalarZnxDftAlloc<B> for Module<B> {
27    fn new_scalar_znx_dft(&self, cols: usize) -> ScalarZnxDftOwned<B> {
28        ScalarZnxDftOwned::new(self, cols)
29    }
30
31    fn bytes_of_scalar_znx_dft(&self, cols: usize) -> usize {
32        ScalarZnxDftOwned::bytes_of(self, cols)
33    }
34
35    fn new_scalar_znx_dft_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxDftOwned<B> {
36        ScalarZnxDftOwned::new_from_bytes(self, cols, bytes)
37    }
38
39    // fn new_scalar_znx_dft_from_bytes_borrow(&self, cols: usize, bytes: &mut [u8]) -> ScalarZnxDft<FFT64> {
40    //     ScalarZnxDft::from_bytes_borrow(self, SCALAR_ZNX_DFT_ROWS, cols, SCALAR_ZNX_DFT_SIZE, bytes)
41    // }
42}
43
44impl<DataMut, Data> ScalarZnxDftOps<DataMut, Data, FFT64> for Module<FFT64>
45where
46    DataMut: AsMut<[u8]> + AsRef<[u8]>,
47    Data: AsRef<[u8]>,
48{
49    fn svp_prepare(&self, res: &mut ScalarZnxDft<DataMut, FFT64>, res_col: usize, a: &Scalar<Data>, a_col: usize) {
50        unsafe {
51            svp::svp_prepare(
52                self.ptr,
53                res.at_mut_ptr(res_col, 0) as *mut svp_ppol_t,
54                a.at_ptr(a_col, 0),
55            )
56        }
57    }
58
59    fn svp_apply_dft(
60        &self,
61        res: &mut VecZnxDft<DataMut, FFT64>,
62        res_col: usize,
63        a: &ScalarZnxDft<Data, FFT64>,
64        a_col: usize,
65        b: &VecZnx<Data>,
66        b_col: usize,
67    ) {
68        unsafe {
69            svp::svp_apply_dft(
70                self.ptr,
71                res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
72                res.size() as u64,
73                a.at_ptr(a_col, 0) as *const svp_ppol_t,
74                b.at_ptr(b_col, 0),
75                b.size() as u64,
76                b.sl() as u64,
77            )
78        }
79    }
80}