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