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