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