1use poulpy_hal::{
 2    layouts::{
 3        Backend, Module, ScalarZnxToRef, SvpPPol, SvpPPolOwned, SvpPPolToMut, SvpPPolToRef, VecZnxDft, VecZnxDftToMut,
 4        VecZnxDftToRef, ZnxInfos, ZnxView, ZnxViewMut,
 5    },
 6    oep::{SvpApplyImpl, SvpApplyInplaceImpl, SvpPPolAllocBytesImpl, SvpPPolAllocImpl, SvpPPolFromBytesImpl, SvpPrepareImpl},
 7};
 8
 9use crate::cpu_spqlios::{
10    FFT64,
11    ffi::{svp, vec_znx_dft::vec_znx_dft_t},
12};
13
14unsafe impl SvpPPolFromBytesImpl<Self> for FFT64 {
15    fn svp_ppol_from_bytes_impl(n: usize, cols: usize, bytes: Vec<u8>) -> SvpPPolOwned<Self> {
16        SvpPPolOwned::from_bytes(n, cols, bytes)
17    }
18}
19
20unsafe impl SvpPPolAllocImpl<Self> for FFT64 {
21    fn svp_ppol_alloc_impl(n: usize, cols: usize) -> SvpPPolOwned<Self> {
22        SvpPPolOwned::alloc(n, cols)
23    }
24}
25
26unsafe impl SvpPPolAllocBytesImpl<Self> for FFT64 {
27    fn svp_ppol_alloc_bytes_impl(n: usize, cols: usize) -> usize {
28        FFT64::layout_prep_word_count() * n * cols * size_of::<f64>()
29    }
30}
31
32unsafe impl SvpPrepareImpl<Self> for FFT64 {
33    fn svp_prepare_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
34    where
35        R: SvpPPolToMut<Self>,
36        A: ScalarZnxToRef,
37    {
38        unsafe {
39            svp::svp_prepare(
40                module.ptr(),
41                res.to_mut().at_mut_ptr(res_col, 0) as *mut svp::svp_ppol_t,
42                a.to_ref().at_ptr(a_col, 0),
43            )
44        }
45    }
46}
47
48unsafe impl SvpApplyImpl<Self> for FFT64 {
49    fn svp_apply_impl<R, A, B>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
50    where
51        R: VecZnxDftToMut<Self>,
52        A: SvpPPolToRef<Self>,
53        B: VecZnxDftToRef<Self>,
54    {
55        let mut res: VecZnxDft<&mut [u8], Self> = res.to_mut();
56        let a: SvpPPol<&[u8], Self> = a.to_ref();
57        let b: VecZnxDft<&[u8], Self> = b.to_ref();
58        unsafe {
59            svp::svp_apply_dft_to_dft(
60                module.ptr(),
61                res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
62                res.size() as u64,
63                res.cols() as u64,
64                a.at_ptr(a_col, 0) as *const svp::svp_ppol_t,
65                b.at_ptr(b_col, 0) as *const vec_znx_dft_t,
66                b.size() as u64,
67                b.cols() as u64,
68            )
69        }
70    }
71}
72
73unsafe impl SvpApplyInplaceImpl for FFT64 {
74    fn svp_apply_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
75    where
76        R: VecZnxDftToMut<Self>,
77        A: SvpPPolToRef<Self>,
78    {
79        let mut res: VecZnxDft<&mut [u8], Self> = res.to_mut();
80        let a: SvpPPol<&[u8], Self> = a.to_ref();
81        unsafe {
82            svp::svp_apply_dft_to_dft(
83                module.ptr(),
84                res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
85                res.size() as u64,
86                res.cols() as u64,
87                a.at_ptr(a_col, 0) as *const svp::svp_ppol_t,
88                res.at_ptr(res_col, 0) as *const vec_znx_dft_t,
89                res.size() as u64,
90                res.cols() as u64,
91            )
92        }
93    }
94}