1use backend::hal::{
  2    api::{
  3        MatZnxAlloc, ScalarZnxAlloc, ScalarZnxAllocBytes, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace,
  4        VecZnxAlloc, VecZnxAllocBytes, VecZnxBigAlloc, VecZnxDftAlloc, VecZnxStd, VecZnxSubScalarInplace, VecZnxSwithcDegree,
  5    },
  6    layouts::{Backend, Module, ScratchOwned, VecZnxDft},
  7    oep::{
  8        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
  9        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
 10    },
 11};
 12use sampling::source::Source;
 13
 14use crate::{
 15    GGLWEEncryptSkFamily, GGLWEExecLayoutFamily, GLWEDecryptFamily, GLWEPlaintext, GLWESecret, GLWESecretExec, GLWETensorKey,
 16    GLWETensorKeyEncryptSkFamily, Infos,
 17};
 18
 19pub(crate) trait TestModuleFamily<B: Backend> = GGLWEEncryptSkFamily<B>
 20    + GLWEDecryptFamily<B>
 21    + MatZnxAlloc
 22    + ScalarZnxAlloc
 23    + ScalarZnxAllocBytes
 24    + VecZnxAllocBytes
 25    + VecZnxSwithcDegree
 26    + VecZnxAddScalarInplace
 27    + VecZnxStd
 28    + VecZnxAlloc
 29    + VecZnxSubScalarInplace;
 30
 31pub(crate) trait TestScratchFamily<B: Backend> = TakeVecZnxDftImpl<B>
 32    + TakeVecZnxBigImpl<B>
 33    + TakeSvpPPolImpl<B>
 34    + ScratchOwnedAllocImpl<B>
 35    + ScratchOwnedBorrowImpl<B>
 36    + ScratchAvailableImpl<B>
 37    + TakeScalarZnxImpl<B>
 38    + TakeVecZnxImpl<B>
 39    + VecZnxDftAllocBytesImpl<B>
 40    + VecZnxBigAllocBytesImpl<B>
 41    + TakeSvpPPolImpl<B>;
 42
 43pub(crate) fn test_encrypt_sk<B: Backend>(module: &Module<B>, basek: usize, k: usize, sigma: f64, rank: usize)
 44where
 45    Module<B>: TestModuleFamily<B>
 46        + GGLWEExecLayoutFamily<B>
 47        + GLWETensorKeyEncryptSkFamily<B>
 48        + GLWEDecryptFamily<B>
 49        + VecZnxDftAlloc<B>
 50        + VecZnxBigAlloc<B>,
 51    B: TestScratchFamily<B>,
 52{
 53    let rows: usize = k / basek;
 54
 55    let mut tensor_key: GLWETensorKey<Vec<u8>> = GLWETensorKey::alloc(&module, basek, k, rows, 1, rank);
 56
 57    let mut source_xs: Source = Source::new([0u8; 32]);
 58    let mut source_xe: Source = Source::new([0u8; 32]);
 59    let mut source_xa: Source = Source::new([0u8; 32]);
 60
 61    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GLWETensorKey::encrypt_sk_scratch_space(
 62        module,
 63        basek,
 64        tensor_key.k(),
 65        rank,
 66    ));
 67
 68    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
 69    sk.fill_ternary_prob(0.5, &mut source_xs);
 70    let mut sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(&module, &sk);
 71    sk_exec.prepare(module, &sk);
 72
 73    tensor_key.encrypt_sk(
 74        module,
 75        &sk,
 76        &mut source_xa,
 77        &mut source_xe,
 78        sigma,
 79        scratch.borrow(),
 80    );
 81
 82    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
 83
 84    let mut sk_ij_dft = module.vec_znx_dft_alloc(1, 1);
 85    let mut sk_ij_big = module.vec_znx_big_alloc(1, 1);
 86    let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, 1);
 87    let mut sk_dft: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(rank, 1);
 88
 89    (0..rank).for_each(|i| {
 90        module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data, i);
 91    });
 92
 93    (0..rank).for_each(|i| {
 94        (0..rank).for_each(|j| {
 95            module.svp_apply(&mut sk_ij_dft, 0, &sk_exec.data, j, &sk_dft, i);
 96            module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
 97            module.vec_znx_big_normalize(basek, &mut sk_ij.data, 0, &sk_ij_big, 0, scratch.borrow());
 98            (0..tensor_key.rank_in()).for_each(|col_i| {
 99                (0..tensor_key.rows()).for_each(|row_i| {
100                    tensor_key
101                        .at(i, j)
102                        .at(row_i, col_i)
103                        .decrypt(&module, &mut pt, &sk_exec, scratch.borrow());
104
105                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, col_i);
106
107                    let std_pt: f64 = module.vec_znx_std(basek, &pt.data, 0) * (k as f64).exp2();
108                    assert!((sigma - std_pt).abs() <= 0.5, "{} {}", sigma, std_pt);
109                });
110            });
111        })
112    })
113}