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