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, ScalarZnx, ScalarZnxDft, ScalarZnxDftOwned, ScalarZnxDftToMut, ScalarZnxDftToRef, ScalarZnxToMut,
  6    ScalarZnxToRef, Scratch, VecZnxDft, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef, VecZnxOps,
  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    fn svp_idft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
 33    where
 34        R: ScalarZnxToMut,
 35        A: ScalarZnxDftToRef<BACKEND>;
 36}
 37
 38impl<B: Backend> ScalarZnxDftAlloc<B> for Module<B> {
 39    fn new_scalar_znx_dft(&self, cols: usize) -> ScalarZnxDftOwned<B> {
 40        ScalarZnxDftOwned::new(self, cols)
 41    }
 42
 43    fn bytes_of_scalar_znx_dft(&self, cols: usize) -> usize {
 44        ScalarZnxDftOwned::bytes_of(self, cols)
 45    }
 46
 47    fn new_scalar_znx_dft_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxDftOwned<B> {
 48        ScalarZnxDftOwned::new_from_bytes(self, cols, bytes)
 49    }
 50}
 51
 52impl ScalarZnxDftOps<FFT64> for Module<FFT64> {
 53    fn svp_idft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
 54    where
 55        R: ScalarZnxToMut,
 56        A: ScalarZnxDftToRef<FFT64>,
 57    {
 58        let res_mut: &mut ScalarZnx<&mut [u8]> = &mut res.to_mut();
 59        let a_ref: &ScalarZnxDft<&[u8], FFT64> = &a.to_ref();
 60        let (mut vec_znx_big, scratch1) = scratch.tmp_vec_znx_big(self, 1, 1);
 61        self.vec_znx_idft(&mut vec_znx_big, 0, a_ref, a_col, scratch1);
 62        self.vec_znx_copy(res_mut, res_col, &vec_znx_big.to_vec_znx_small(), 0);
 63    }
 64
 65    fn svp_prepare<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
 66    where
 67        R: ScalarZnxDftToMut<FFT64>,
 68        A: ScalarZnxToRef,
 69    {
 70        unsafe {
 71            svp::svp_prepare(
 72                self.ptr,
 73                res.to_mut().at_mut_ptr(res_col, 0) as *mut svp::svp_ppol_t,
 74                a.to_ref().at_ptr(a_col, 0),
 75            )
 76        }
 77    }
 78
 79    fn svp_apply<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
 80    where
 81        R: VecZnxDftToMut<FFT64>,
 82        A: ScalarZnxDftToRef<FFT64>,
 83        B: VecZnxDftToRef<FFT64>,
 84    {
 85        let mut res: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
 86        let a: ScalarZnxDft<&[u8], FFT64> = a.to_ref();
 87        let b: VecZnxDft<&[u8], FFT64> = b.to_ref();
 88        unsafe {
 89            svp::svp_apply_dft_to_dft(
 90                self.ptr,
 91                res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
 92                res.size() as u64,
 93                res.cols() as u64,
 94                a.at_ptr(a_col, 0) as *const svp::svp_ppol_t,
 95                b.at_ptr(b_col, 0) as *const vec_znx_dft_t,
 96                b.size() as u64,
 97                b.cols() as u64,
 98            )
 99        }
100    }
101
102    fn svp_apply_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
103    where
104        R: VecZnxDftToMut<FFT64>,
105        A: ScalarZnxDftToRef<FFT64>,
106    {
107        let mut res: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
108        let a: ScalarZnxDft<&[u8], FFT64> = a.to_ref();
109        unsafe {
110            svp::svp_apply_dft_to_dft(
111                self.ptr,
112                res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
113                res.size() as u64,
114                res.cols() as u64,
115                a.at_ptr(a_col, 0) as *const svp::svp_ppol_t,
116                res.at_ptr(res_col, 0) as *const vec_znx_dft_t,
117                res.size() as u64,
118                res.cols() as u64,
119            )
120        }
121    }
122}