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