1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxAutomorphismInplace, VecZnxEncodeCoeffsi64,
  4        VecZnxSwithcDegree, VmpPMatAlloc, VmpPMatPrepare, ZnxView,
  5    },
  6    layouts::{Backend, Module, ScratchOwned},
  7    oep::{
  8        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
  9        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 10    },
 11};
 12use sampling::source::Source;
 13
 14use crate::layouts::{Infos, LWECiphertext, LWEPlaintext, LWESecret, LWESwitchingKey, prepared::LWESwitchingKeyExec};
 15
 16use crate::trait_families::{GGLWEEncryptSkFamily, GLWEDecryptFamily, GLWEKeyswitchFamily, GLWESecretExecModuleFamily};
 17
 18pub fn test_lwe_keyswitch<B: Backend>(module: &Module<B>)
 19where
 20    Module<B>: GGLWEEncryptSkFamily<B>
 21        + GLWEDecryptFamily<B>
 22        + VecZnxSwithcDegree
 23        + VecZnxAddScalarInplace
 24        + GLWEKeyswitchFamily<B>
 25        + VecZnxEncodeCoeffsi64
 26        + VecZnxAutomorphismInplace
 27        + VmpPMatAlloc<B>
 28        + VmpPMatPrepare<B>
 29        + GLWESecretExecModuleFamily<B>,
 30    B: TakeScalarZnxImpl<B>
 31        + TakeVecZnxDftImpl<B>
 32        + ScratchAvailableImpl<B>
 33        + TakeVecZnxImpl<B>
 34        + TakeVecZnxBigImpl<B>
 35        + TakeSvpPPolImpl<B>
 36        + ScratchOwnedAllocImpl<B>
 37        + ScratchOwnedBorrowImpl<B>,
 38{
 39    let n: usize = module.n();
 40    let basek: usize = 17;
 41    let sigma: f64 = 3.2;
 42
 43    let n_lwe_in: usize = 22;
 44    let n_lwe_out: usize = 30;
 45    let k_lwe_ct: usize = 2 * basek;
 46    let k_lwe_pt: usize = 8;
 47
 48    let k_ksk: usize = k_lwe_ct + basek;
 49
 50    let mut source_xs: Source = Source::new([0u8; 32]);
 51    let mut source_xa: Source = Source::new([0u8; 32]);
 52    let mut source_xe: Source = Source::new([0u8; 32]);
 53
 54    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 55        LWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k_ksk)
 56            | LWECiphertext::keyswitch_scratch_space(module, n, basek, k_lwe_ct, k_lwe_ct, k_ksk),
 57    );
 58
 59    let mut sk_lwe_in: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_in);
 60    sk_lwe_in.fill_ternary_prob(0.5, &mut source_xs);
 61
 62    let mut sk_lwe_out: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_out);
 63    sk_lwe_out.fill_ternary_prob(0.5, &mut source_xs);
 64
 65    let data: i64 = 17;
 66
 67    let mut lwe_pt_in: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
 68    module.encode_coeff_i64(basek, &mut lwe_pt_in.data, 0, k_lwe_pt, 0, data, k_lwe_pt);
 69
 70    let mut lwe_ct_in: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe_in, basek, k_lwe_ct);
 71    lwe_ct_in.encrypt_sk(
 72        module,
 73        &lwe_pt_in,
 74        &sk_lwe_in,
 75        &mut source_xa,
 76        &mut source_xe,
 77        sigma,
 78    );
 79
 80    let mut ksk: LWESwitchingKey<Vec<u8>> = LWESwitchingKey::alloc(n, basek, k_ksk, lwe_ct_in.size());
 81
 82    ksk.encrypt_sk(
 83        module,
 84        &sk_lwe_in,
 85        &sk_lwe_out,
 86        &mut source_xa,
 87        &mut source_xe,
 88        sigma,
 89        scratch.borrow(),
 90    );
 91
 92    let mut lwe_ct_out: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe_out, basek, k_lwe_ct);
 93
 94    let ksk_exec: LWESwitchingKeyExec<Vec<u8>, B> = LWESwitchingKeyExec::from(module, &ksk, scratch.borrow());
 95
 96    lwe_ct_out.keyswitch(module, &lwe_ct_in, &ksk_exec, scratch.borrow());
 97
 98    let mut lwe_pt_out: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_ct);
 99    lwe_ct_out.decrypt(module, &mut lwe_pt_out, &sk_lwe_out);
100
101    assert_eq!(lwe_pt_in.data.at(0, 0)[0], lwe_pt_out.data.at(0, 0)[0]);
102}