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                    n,
 55                    res_size,
 56                    a_size,
 57                    mat_rows,
 58                    mat_cols_in,
 59                    mat_cols_out,
 60                    mat_size,
 61                ) | module.vec_znx_big_normalize_tmp_bytes(n),
 62            );
 63
 64            let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, a_cols, a_size);
 65
 66            (0..a_cols).for_each(|i| {
 67                a.at_mut(i, a_size - 1)[i + 1] = 1;
 68            });
 69
 70            let mut vmp: VmpPMat<Vec<u8>, B> = module.vmp_pmat_alloc(n, mat_rows, mat_cols_in, mat_cols_out, mat_size);
 71
 72            let mut c_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(n, mat_cols_out, mat_size);
 73            let mut c_big: VecZnxBig<Vec<u8>, B> = module.vec_znx_big_alloc(n, mat_cols_out, mat_size);
 74
 75            let mut mat: MatZnx<Vec<u8>> = MatZnx::alloc(n, mat_rows, mat_cols_in, mat_cols_out, mat_size);
 76
 77            // Construts a [VecZnxMatDft] that performs cyclic rotations on each submatrix.
 78            (0..a.size()).for_each(|row_i| {
 79                (0..mat_cols_in).for_each(|col_in_i| {
 80                    (0..mat_cols_out).for_each(|col_out_i| {
 81                        let idx = 1 + col_in_i * mat_cols_out + col_out_i;
 82                        mat.at_mut(row_i, col_in_i).at_mut(col_out_i, row_i)[idx] = 1_i64; // X^{idx}
 83                    });
 84                });
 85            });
 86
 87            module.vmp_prepare(&mut vmp, &mat, scratch.borrow());
 88
 89            let mut a_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(n, a_cols, a_size);
 90            (0..a_cols).for_each(|i| {
 91                module.vec_znx_dft_from_vec_znx(1, 0, &mut a_dft, i, &a, i);
 92            });
 93
 94            module.vmp_apply(&mut c_dft, &a_dft, &vmp, scratch.borrow());
 95
 96            let mut res_have_vi64: Vec<i64> = vec![i64::default(); n];
 97
 98            let mut res_have: VecZnx<Vec<u8>> = VecZnx::alloc(n, res_cols, res_size);
 99            (0..mat_cols_out).for_each(|i| {
100                module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut c_big, i, &mut c_dft, i);
101                module.vec_znx_big_normalize(basek, &mut res_have, i, &c_big, i, scratch.borrow());
102            });
103
104            (0..mat_cols_out).for_each(|col_i| {
105                let mut res_want_vi64: Vec<i64> = vec![i64::default(); n];
106                (0..a_cols).for_each(|i| {
107                    res_want_vi64[(i + 1) + (1 + i * mat_cols_out + col_i)] = 1;
108                });
109                res_have.decode_vec_i64(basek, col_i, basek * a_size, &mut res_have_vi64);
110                assert_eq!(res_have_vi64, res_want_vi64);
111            });
112        });
113    });
114}