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, ScalarZnxDft, ScalarZnxDftOwned, ScalarZnxDftToMut, ScalarZnxDftToRef, ScalarZnxToRef, VecZnxDft,
  6    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: ScalarZnxToRef;
 20
 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: VecZnxDftToRef<BACKEND>;
 26
 27    fn svp_apply_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
 28    where
 29        R: VecZnxDftToMut<BACKEND>,
 30        A: ScalarZnxDftToRef<BACKEND>;
 31}
 32
 33impl<B: Backend> ScalarZnxDftAlloc<B> for Module<B> {
 34    fn new_scalar_znx_dft(&self, cols: usize) -> ScalarZnxDftOwned<B> {
 35        ScalarZnxDftOwned::new(self, cols)
 36    }
 37
 38    fn bytes_of_scalar_znx_dft(&self, cols: usize) -> usize {
 39        ScalarZnxDftOwned::bytes_of(self, cols)
 40    }
 41
 42    fn new_scalar_znx_dft_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxDftOwned<B> {
 43        ScalarZnxDftOwned::new_from_bytes(self, cols, bytes)
 44    }
 45}
 46
 47impl ScalarZnxDftOps<FFT64> for Module<FFT64> {
 48    fn svp_prepare<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
 49    where
 50        R: ScalarZnxDftToMut<FFT64>,
 51        A: ScalarZnxToRef,
 52    {
 53        unsafe {
 54            svp::svp_prepare(
 55                self.ptr,
 56                res.to_mut().at_mut_ptr(res_col, 0) as *mut svp::svp_ppol_t,
 57                a.to_ref().at_ptr(a_col, 0),
 58            )
 59        }
 60    }
 61
 62    fn svp_apply<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
 63    where
 64        R: VecZnxDftToMut<FFT64>,
 65        A: ScalarZnxDftToRef<FFT64>,
 66        B: VecZnxDftToRef<FFT64>,
 67    {
 68        let mut res: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
 69        let a: ScalarZnxDft<&[u8], FFT64> = a.to_ref();
 70        let b: VecZnxDft<&[u8], FFT64> = b.to_ref();
 71        unsafe {
 72            svp::svp_apply_dft_to_dft(
 73                self.ptr,
 74                res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
 75                res.size() as u64,
 76                res.cols() as u64,
 77                a.at_ptr(a_col, 0) as *const svp::svp_ppol_t,
 78                b.at_ptr(b_col, 0) as *const vec_znx_dft_t,
 79                b.size() as u64,
 80                b.cols() as u64,
 81            )
 82        }
 83    }
 84
 85    fn svp_apply_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
 86    where
 87        R: VecZnxDftToMut<FFT64>,
 88        A: ScalarZnxDftToRef<FFT64>,
 89    {
 90        let mut res: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
 91        let a: ScalarZnxDft<&[u8], FFT64> = a.to_ref();
 92        unsafe {
 93            svp::svp_apply_dft_to_dft(
 94                self.ptr,
 95                res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
 96                res.size() as u64,
 97                res.cols() as u64,
 98                a.at_ptr(a_col, 0) as *const svp::svp_ppol_t,
 99                res.at_ptr(res_col, 0) as *const vec_znx_dft_t,
100                res.size() as u64,
101                res.cols() as u64,
102            )
103        }
104    }
105}