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, VecZnx,
 6    VecZnxDft, VecZnxDftToMut, VecZnxToRef, ZnxSliceSize,
 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    // fn new_scalar_znx_dft_from_bytes_borrow(&self, cols: usize, bytes: &mut [u8]) -> ScalarZnxDft<B>;
14}
15
16pub trait ScalarZnxDftOps<BACKEND: Backend> {
17    fn svp_prepare<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
18    where
19        R: ScalarZnxDftToMut<BACKEND>,
20        A: ScalarToRef;
21    fn svp_apply<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
22    where
23        R: VecZnxDftToMut<BACKEND>,
24        A: ScalarZnxDftToRef<BACKEND>,
25        B: VecZnxToRef;
26}
27
28impl<B: Backend> ScalarZnxDftAlloc<B> for Module<B> {
29    fn new_scalar_znx_dft(&self, cols: usize) -> ScalarZnxDftOwned<B> {
30        ScalarZnxDftOwned::new(self, cols)
31    }
32
33    fn bytes_of_scalar_znx_dft(&self, cols: usize) -> usize {
34        ScalarZnxDftOwned::bytes_of(self, cols)
35    }
36
37    fn new_scalar_znx_dft_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxDftOwned<B> {
38        ScalarZnxDftOwned::new_from_bytes(self, cols, bytes)
39    }
40}
41
42impl ScalarZnxDftOps<FFT64> for Module<FFT64> {
43    fn svp_prepare<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
44    where
45        R: ScalarZnxDftToMut<FFT64>,
46        A: ScalarToRef,
47    {
48        unsafe {
49            svp::svp_prepare(
50                self.ptr,
51                res.to_mut().at_mut_ptr(res_col, 0) as *mut svp::svp_ppol_t,
52                a.to_ref().at_ptr(a_col, 0),
53            )
54        }
55    }
56
57    fn svp_apply<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
58    where
59        R: VecZnxDftToMut<FFT64>,
60        A: ScalarZnxDftToRef<FFT64>,
61        B: VecZnxToRef,
62    {
63        let mut res: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
64        let a: ScalarZnxDft<&[u8], FFT64> = a.to_ref();
65        let b: VecZnx<&[u8]> = b.to_ref();
66        unsafe {
67            svp::svp_apply_dft(
68                self.ptr,
69                res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
70                res.size() as u64,
71                a.at_ptr(a_col, 0) as *const svp::svp_ppol_t,
72                b.at_ptr(b_col, 0),
73                b.size() as u64,
74                b.sl() as u64,
75            )
76        }
77    }
78}