1use poulpy_hal::{
  2    api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow},
  3    layouts::{Backend, Module, Scratch, ScratchOwned, ZnxView},
  4    source::Source,
  5};
  6
  7use crate::{
  8    LWEDecrypt, LWEEncryptSk, LWEKeySwitch, LWESwitchingKeyEncrypt, ScratchTakeCore,
  9    layouts::{
 10        LWE, LWELayout, LWEPlaintext, LWESecret, LWESwitchingKey, LWESwitchingKeyLayout, LWESwitchingKeyPreparedAlloc,
 11        prepared::LWESwitchingKeyPrepared,
 12    },
 13};
 14
 15pub fn test_lwe_keyswitch<BE: Backend>(module: &Module<BE>)
 16where
 17    Module<BE>:
 18        LWEKeySwitch<BE> + LWESwitchingKeyEncrypt<BE> + LWEEncryptSk<BE> + LWESwitchingKeyPreparedAlloc<BE> + LWEDecrypt<BE>,
 19    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 20    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 21{
 22    let n: usize = module.n();
 23    let base2k: usize = 17;
 24
 25    let n_lwe_in: usize = 22;
 26    let n_lwe_out: usize = 30;
 27    let k_lwe_ct: usize = 2 * base2k;
 28    let k_lwe_pt: usize = 8;
 29
 30    let k_ksk: usize = k_lwe_ct + base2k;
 31    let dnum: usize = k_lwe_ct.div_ceil(base2k);
 32
 33    let mut source_xs: Source = Source::new([0u8; 32]);
 34    let mut source_xa: Source = Source::new([0u8; 32]);
 35    let mut source_xe: Source = Source::new([0u8; 32]);
 36
 37    let key_apply_infos: LWESwitchingKeyLayout = LWESwitchingKeyLayout {
 38        n: n.into(),
 39        base2k: base2k.into(),
 40        k: k_ksk.into(),
 41        dnum: dnum.into(),
 42    };
 43
 44    let lwe_in_infos: LWELayout = LWELayout {
 45        n: n_lwe_in.into(),
 46        base2k: base2k.into(),
 47        k: k_lwe_ct.into(),
 48    };
 49
 50    let lwe_out_infos: LWELayout = LWELayout {
 51        n: n_lwe_out.into(),
 52        k: k_lwe_ct.into(),
 53        base2k: base2k.into(),
 54    };
 55
 56    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
 57        LWESwitchingKey::encrypt_sk_tmp_bytes(module, &key_apply_infos)
 58            | LWE::keyswitch_tmp_bytes(module, &lwe_out_infos, &lwe_in_infos, &key_apply_infos),
 59    );
 60
 61    let mut sk_lwe_in: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_in.into());
 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.into());
 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(base2k.into(), k_lwe_pt.into());
 70    lwe_pt_in.encode_i64(data, k_lwe_pt.into());
 71
 72    let mut lwe_ct_in: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_in_infos);
 73    lwe_ct_in.encrypt_sk(
 74        module,
 75        &lwe_pt_in,
 76        &sk_lwe_in,
 77        &mut source_xa,
 78        &mut source_xe,
 79    );
 80
 81    let mut ksk: LWESwitchingKey<Vec<u8>> = LWESwitchingKey::alloc_from_infos(&key_apply_infos);
 82
 83    ksk.encrypt_sk(
 84        module,
 85        &sk_lwe_in,
 86        &sk_lwe_out,
 87        &mut source_xa,
 88        &mut source_xe,
 89        scratch.borrow(),
 90    );
 91
 92    let mut lwe_ct_out: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_out_infos);
 93
 94    let mut ksk_prepared: LWESwitchingKeyPrepared<Vec<u8>, BE> = LWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
 95    ksk_prepared.prepare(module, &ksk, scratch.borrow());
 96
 97    lwe_ct_out.keyswitch(module, &lwe_ct_in, &ksk_prepared, scratch.borrow());
 98
 99    let mut lwe_pt_out: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc_from_infos(&lwe_out_infos);
100    lwe_ct_out.decrypt(module, &mut lwe_pt_out, &sk_lwe_out);
101
102    assert_eq!(lwe_pt_in.data.at(0, 0)[0], lwe_pt_out.data.at(0, 0)[0]);
103}