1use crate::{
  2    api::{
  3        ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxBigAlloc, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
  4        VecZnxDftAlloc, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigTmpA, VmpApply, VmpApplyTmpBytes, VmpPMatAlloc, VmpPrepare,
  5    },
  6    layouts::{MatZnx, Module, ScratchOwned, VecZnx, VecZnxBig, VecZnxDft, VmpPMat, ZnxInfos, ZnxViewMut},
  7    oep::{
  8        ModuleNewImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, VecZnxBigAllocImpl, VecZnxBigNormalizeImpl,
  9        VecZnxBigNormalizeTmpBytesImpl, VecZnxDftAllocImpl, VecZnxDftFromVecZnxImpl, VecZnxDftToVecZnxBigTmpAImpl, VmpApplyImpl,
 10        VmpApplyTmpBytesImpl, VmpPMatAllocImpl, VmpPMatPrepareImpl,
 11    },
 12};
 13
 14use crate::layouts::Backend;
 15
 16pub fn test_vmp_apply<B>()
 17where
 18    B: Backend
 19        + ModuleNewImpl<B>
 20        + VmpApplyTmpBytesImpl<B>
 21        + VecZnxBigNormalizeTmpBytesImpl<B>
 22        + VmpPMatAllocImpl<B>
 23        + VecZnxDftAllocImpl<B>
 24        + VecZnxBigAllocImpl<B>
 25        + VmpPMatPrepareImpl<B>
 26        + VecZnxDftFromVecZnxImpl<B>
 27        + VmpApplyImpl<B>
 28        + VecZnxDftToVecZnxBigTmpAImpl<B>
 29        + ScratchOwnedAllocImpl<B>
 30        + ScratchOwnedBorrowImpl<B>
 31        + VecZnxBigNormalizeImpl<B>,
 32{
 33    let log_n: i32 = 5;
 34    let n: usize = 1 << log_n;
 35
 36    let module: Module<B> = Module::<B>::new(n as u64);
 37    let basek: usize = 15;
 38    let a_size: usize = 5;
 39    let mat_size: usize = 6;
 40    let res_size: usize = a_size;
 41
 42    [1, 2].iter().for_each(|cols_in| {
 43        [1, 2].iter().for_each(|cols_out| {
 44            let a_cols: usize = *cols_in;
 45            let res_cols: usize = *cols_out;
 46
 47            let mat_rows: usize = a_size;
 48            let mat_cols_in: usize = a_cols;
 49            let mat_cols_out: usize = res_cols;
 50
 51            let mut scratch = ScratchOwned::alloc(
 52                module.vmp_apply_tmp_bytes(
 53                    res_size,
 54                    a_size,
 55                    mat_rows,
 56                    mat_cols_in,
 57                    mat_cols_out,
 58                    mat_size,
 59                ) | module.vec_znx_big_normalize_tmp_bytes(),
 60            );
 61
 62            let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, a_cols, a_size);
 63
 64            (0..a_cols).for_each(|i| {
 65                a.at_mut(i, a_size - 1)[i + 1] = 1;
 66            });
 67
 68            let mut vmp: VmpPMat<Vec<u8>, B> = module.vmp_pmat_alloc(mat_rows, mat_cols_in, mat_cols_out, mat_size);
 69
 70            let mut c_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(mat_cols_out, mat_size);
 71            let mut c_big: VecZnxBig<Vec<u8>, B> = module.vec_znx_big_alloc(mat_cols_out, mat_size);
 72
 73            let mut mat: MatZnx<Vec<u8>> = MatZnx::alloc(n, mat_rows, mat_cols_in, mat_cols_out, mat_size);
 74
 75            // Construts a [VecZnxMatDft] that performs cyclic rotations on each submatrix.
 76            (0..a.size()).for_each(|row_i| {
 77                (0..mat_cols_in).for_each(|col_in_i| {
 78                    (0..mat_cols_out).for_each(|col_out_i| {
 79                        let idx = 1 + col_in_i * mat_cols_out + col_out_i;
 80                        mat.at_mut(row_i, col_in_i).at_mut(col_out_i, row_i)[idx] = 1_i64; // X^{idx}
 81                    });
 82                });
 83            });
 84
 85            module.vmp_prepare(&mut vmp, &mat, scratch.borrow());
 86
 87            let mut a_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(a_cols, a_size);
 88            (0..a_cols).for_each(|i| {
 89                module.vec_znx_dft_from_vec_znx(1, 0, &mut a_dft, i, &a, i);
 90            });
 91
 92            module.vmp_apply(&mut c_dft, &a_dft, &vmp, scratch.borrow());
 93
 94            let mut res_have_vi64: Vec<i64> = vec![i64::default(); n];
 95
 96            let mut res_have: VecZnx<Vec<u8>> = VecZnx::alloc(n, res_cols, res_size);
 97            (0..mat_cols_out).for_each(|i| {
 98                module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut c_big, i, &mut c_dft, i);
 99                module.vec_znx_big_normalize(basek, &mut res_have, i, &c_big, i, scratch.borrow());
100            });
101
102            (0..mat_cols_out).for_each(|col_i| {
103                let mut res_want_vi64: Vec<i64> = vec![i64::default(); n];
104                (0..a_cols).for_each(|i| {
105                    res_want_vi64[(i + 1) + (1 + i * mat_cols_out + col_i)] = 1;
106                });
107                res_have.decode_vec_i64(basek, col_i, basek * a_size, &mut res_have_vi64);
108                assert_eq!(res_have_vi64, res_want_vi64);
109            });
110        });
111    });
112}