1use crate::ffi::svp;
 2use crate::ffi::vec_znx_dft::vec_znx_dft_t;
 3use crate::znx_base::{ZnxInfos, ZnxView, ZnxViewMut};
 4use crate::{
 5    Backend, FFT64, Module, ScalarToRef, ScalarZnxDft, ScalarZnxDftOwned, ScalarZnxDftToMut, ScalarZnxDftToRef,
 6    VecZnxDft, VecZnxDftToMut, VecZnxDftToRef,
 7};
 8
 9pub trait ScalarZnxDftAlloc<B: Backend> {
10    fn new_scalar_znx_dft(&self, cols: usize) -> ScalarZnxDftOwned<B>;
11    fn bytes_of_scalar_znx_dft(&self, cols: usize) -> usize;
12    fn new_scalar_znx_dft_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxDftOwned<B>;
13}
14
15pub trait ScalarZnxDftOps<BACKEND: Backend> {
16    fn svp_prepare<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
17    where
18        R: ScalarZnxDftToMut<BACKEND>,
19        A: ScalarToRef;
20    fn svp_apply<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
21    where
22        R: VecZnxDftToMut<BACKEND>,
23        A: ScalarZnxDftToRef<BACKEND>,
24        B: VecZnxDftToRef<FFT64>;
25    fn svp_apply_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
26    where
27        R: VecZnxDftToMut<BACKEND>,
28        A: ScalarZnxDftToRef<BACKEND>;
29}
30
31impl<B: Backend> ScalarZnxDftAlloc<B> for Module<B> {
32    fn new_scalar_znx_dft(&self, cols: usize) -> ScalarZnxDftOwned<B> {
33        ScalarZnxDftOwned::new(self, cols)
34    }
35
36    fn bytes_of_scalar_znx_dft(&self, cols: usize) -> usize {
37        ScalarZnxDftOwned::bytes_of(self, cols)
38    }
39
40    fn new_scalar_znx_dft_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxDftOwned<B> {
41        ScalarZnxDftOwned::new_from_bytes(self, cols, bytes)
42    }
43}
44
45impl ScalarZnxDftOps<FFT64> for Module<FFT64> {
46    fn svp_prepare<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
47    where
48        R: ScalarZnxDftToMut<FFT64>,
49        A: ScalarToRef,
50    {
51        unsafe {
52            svp::svp_prepare(
53                self.ptr,
54                res.to_mut().at_mut_ptr(res_col, 0) as *mut svp::svp_ppol_t,
55                a.to_ref().at_ptr(a_col, 0),
56            )
57        }
58    }
59
60    fn svp_apply<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
61    where
62        R: VecZnxDftToMut<FFT64>,
63        A: ScalarZnxDftToRef<FFT64>,
64        B: VecZnxDftToRef<FFT64>,
65    {
66        let mut res: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
67        let a: ScalarZnxDft<&[u8], FFT64> = a.to_ref();
68        let b: VecZnxDft<&[u8], FFT64> = b.to_ref();
69        unsafe {
70            svp::svp_apply_dft_to_dft(
71                self.ptr,
72                res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
73                res.size() as u64,
74                a.at_ptr(a_col, 0) as *const svp::svp_ppol_t,
75                b.at_ptr(b_col, 0) as *const vec_znx_dft_t,
76                b.size() as u64,
77            )
78        }
79    }
80
81    fn svp_apply_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
82    where
83        R: VecZnxDftToMut<FFT64>,
84        A: ScalarZnxDftToRef<FFT64>,
85    {
86        let mut res: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
87        let a: ScalarZnxDft<&[u8], FFT64> = a.to_ref();
88        unsafe {
89            svp::svp_apply_dft_to_dft(
90                self.ptr,
91                res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
92                res.size() as u64,
93                a.at_ptr(a_col, 0) as *const svp::svp_ppol_t,
94                res.at_ptr(res_col, 0) as *const vec_znx_dft_t,
95                res.size() as u64,
96            )
97        }
98    }
99}