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