1use poulpy_hal::{
  2    api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxNormalize},
  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, LWESwitchingKeyPreparedFactory,
 11        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        + LWESwitchingKeyPreparedFactory<BE>
 21        + LWEDecrypt<BE>
 22        + VecZnxNormalize<BE>,
 23    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 24    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 25{
 26    let n: usize = module.n();
 27    let in_base2k: usize = 17;
 28    let out_base2k: usize = 15;
 29    let key_base2k: usize = 13;
 30
 31    let n_lwe_in: usize = module.n() >> 1;
 32    let n_lwe_out: usize = module.n() >> 1;
 33    let k_lwe_ct: usize = 102;
 34    let k_lwe_pt: usize = 8;
 35
 36    let k_ksk: usize = k_lwe_ct + key_base2k;
 37    let dnum: usize = k_lwe_ct.div_ceil(key_base2k);
 38
 39    let mut source_xs: Source = Source::new([0u8; 32]);
 40    let mut source_xa: Source = Source::new([0u8; 32]);
 41    let mut source_xe: Source = Source::new([0u8; 32]);
 42
 43    let key_apply_infos: LWESwitchingKeyLayout = LWESwitchingKeyLayout {
 44        n: n.into(),
 45        base2k: key_base2k.into(),
 46        k: k_ksk.into(),
 47        dnum: dnum.into(),
 48    };
 49
 50    let lwe_in_infos: LWELayout = LWELayout {
 51        n: n_lwe_in.into(),
 52        base2k: in_base2k.into(),
 53        k: k_lwe_ct.into(),
 54    };
 55
 56    let lwe_out_infos: LWELayout = LWELayout {
 57        n: n_lwe_out.into(),
 58        k: k_lwe_ct.into(),
 59        base2k: out_base2k.into(),
 60    };
 61
 62    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
 63        LWESwitchingKey::encrypt_sk_tmp_bytes(module, &key_apply_infos)
 64            | LWE::keyswitch_tmp_bytes(module, &lwe_out_infos, &lwe_in_infos, &key_apply_infos),
 65    );
 66
 67    let mut sk_lwe_in: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_in.into());
 68    sk_lwe_in.fill_ternary_prob(0.5, &mut source_xs);
 69
 70    let mut sk_lwe_out: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_out.into());
 71    sk_lwe_out.fill_ternary_prob(0.5, &mut source_xs);
 72
 73    let data: i64 = 17;
 74
 75    let mut lwe_pt_in: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(in_base2k.into(), k_lwe_pt.into());
 76    lwe_pt_in.encode_i64(data, k_lwe_pt.into());
 77
 78    let mut lwe_ct_in: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_in_infos);
 79    lwe_ct_in.encrypt_sk(
 80        module,
 81        &lwe_pt_in,
 82        &sk_lwe_in,
 83        &mut source_xa,
 84        &mut source_xe,
 85        scratch.borrow(),
 86    );
 87
 88    let mut ksk: LWESwitchingKey<Vec<u8>> = LWESwitchingKey::alloc_from_infos(&key_apply_infos);
 89
 90    ksk.encrypt_sk(
 91        module,
 92        &sk_lwe_in,
 93        &sk_lwe_out,
 94        &mut source_xa,
 95        &mut source_xe,
 96        scratch.borrow(),
 97    );
 98
 99    let mut lwe_ct_out: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_out_infos);
100
101    let mut ksk_prepared: LWESwitchingKeyPrepared<Vec<u8>, BE> = LWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
102    ksk_prepared.prepare(module, &ksk, scratch.borrow());
103
104    lwe_ct_out.keyswitch(module, &lwe_ct_in, &ksk_prepared, scratch.borrow());
105
106    let mut lwe_pt_out: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc_from_infos(&lwe_out_infos);
107    lwe_ct_out.decrypt(module, &mut lwe_pt_out, &sk_lwe_out, scratch.borrow());
108
109    let mut lwe_pt_want: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc_from_infos(&lwe_out_infos);
110    module.vec_znx_normalize(
111        lwe_pt_want.data_mut(),
112        out_base2k,
113        0,
114        0,
115        lwe_pt_in.data(),
116        in_base2k,
117        0,
118        scratch.borrow(),
119    );
120
121    assert_eq!(lwe_pt_want.data.at(0, 0)[0], lwe_pt_out.data.at(0, 0)[0]);
122}