1use crate::{
  2    layouts::{
  3        Backend, ScalarZnx, ScalarZnxToRef, SvpPPol, SvpPPolToMut, SvpPPolToRef, VecZnx, VecZnxDft, VecZnxDftToMut,
  4        VecZnxDftToRef, VecZnxToRef, ZnxInfos, ZnxView, ZnxViewMut,
  5    },
  6    reference::fft64::reim::{ReimAddMul, ReimDFTExecute, ReimFFTTable, ReimFromZnx, ReimMul, ReimMulInplace, ReimZero},
  7};
  8
  9pub fn svp_prepare<R, A, BE>(table: &ReimFFTTable<f64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
 10where
 11    BE: Backend<ScalarPrep = f64> + ReimDFTExecute<ReimFFTTable<f64>, f64> + ReimFromZnx,
 12    R: SvpPPolToMut<BE>,
 13    A: ScalarZnxToRef,
 14{
 15    let mut res: SvpPPol<&mut [u8], BE> = res.to_mut();
 16    let a: ScalarZnx<&[u8]> = a.to_ref();
 17    BE::reim_from_znx(res.at_mut(res_col, 0), a.at(a_col, 0));
 18    BE::reim_dft_execute(table, res.at_mut(res_col, 0));
 19}
 20
 21pub fn svp_apply_dft<R, A, B, BE>(
 22    table: &ReimFFTTable<f64>,
 23    res: &mut R,
 24    res_col: usize,
 25    a: &A,
 26    a_col: usize,
 27    b: &B,
 28    b_col: usize,
 29) where
 30    BE: Backend<ScalarPrep = f64> + ReimDFTExecute<ReimFFTTable<f64>, f64> + ReimZero + ReimFromZnx + ReimMulInplace,
 31    R: VecZnxDftToMut<BE>,
 32    A: SvpPPolToRef<BE>,
 33    B: VecZnxToRef,
 34{
 35    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
 36    let a: SvpPPol<&[u8], BE> = a.to_ref();
 37    let b: VecZnx<&[u8]> = b.to_ref();
 38
 39    let res_size: usize = res.size();
 40    let b_size: usize = b.size();
 41    let min_size: usize = res_size.min(b_size);
 42
 43    let ppol: &[f64] = a.at(a_col, 0);
 44    for j in 0..min_size {
 45        let out: &mut [f64] = res.at_mut(res_col, j);
 46        BE::reim_from_znx(out, b.at(b_col, j));
 47        BE::reim_dft_execute(table, out);
 48        BE::reim_mul_inplace(out, ppol);
 49    }
 50
 51    for j in min_size..res_size {
 52        BE::reim_zero(res.at_mut(res_col, j));
 53    }
 54}
 55
 56pub fn svp_apply_dft_to_dft<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
 57where
 58    BE: Backend<ScalarPrep = f64> + ReimMul + ReimZero,
 59    R: VecZnxDftToMut<BE>,
 60    A: SvpPPolToRef<BE>,
 61    B: VecZnxDftToRef<BE>,
 62{
 63    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
 64    let a: SvpPPol<&[u8], BE> = a.to_ref();
 65    let b: VecZnxDft<&[u8], BE> = b.to_ref();
 66
 67    let res_size: usize = res.size();
 68    let b_size: usize = b.size();
 69    let min_size: usize = res_size.min(b_size);
 70
 71    let ppol: &[f64] = a.at(a_col, 0);
 72    for j in 0..min_size {
 73        BE::reim_mul(res.at_mut(res_col, j), ppol, b.at(b_col, j));
 74    }
 75
 76    for j in min_size..res_size {
 77        BE::reim_zero(res.at_mut(res_col, j));
 78    }
 79}
 80
 81pub fn svp_apply_dft_to_dft_add<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
 82where
 83    BE: Backend<ScalarPrep = f64> + ReimAddMul + ReimZero,
 84    R: VecZnxDftToMut<BE>,
 85    A: SvpPPolToRef<BE>,
 86    B: VecZnxDftToRef<BE>,
 87{
 88    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
 89    let a: SvpPPol<&[u8], BE> = a.to_ref();
 90    let b: VecZnxDft<&[u8], BE> = b.to_ref();
 91
 92    let res_size: usize = res.size();
 93    let b_size: usize = b.size();
 94    let min_size: usize = res_size.min(b_size);
 95
 96    let ppol: &[f64] = a.at(a_col, 0);
 97    for j in 0..min_size {
 98        BE::reim_addmul(res.at_mut(res_col, j), ppol, b.at(b_col, j));
 99    }
100
101    for j in min_size..res_size {
102        BE::reim_zero(res.at_mut(res_col, j));
103    }
104}
105
106pub fn svp_apply_dft_to_dft_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
107where
108    BE: Backend<ScalarPrep = f64> + ReimMulInplace,
109    R: VecZnxDftToMut<BE>,
110    A: SvpPPolToRef<BE>,
111{
112    let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
113    let a: SvpPPol<&[u8], BE> = a.to_ref();
114
115    let ppol: &[f64] = a.at(a_col, 0);
116    for j in 0..res.size() {
117        BE::reim_mul_inplace(res.at_mut(res_col, j), ppol);
118    }
119}