1use poulpy_hal::{
  2    api::{
  3        DFT, IDFTConsume, ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
  4        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphismInplace, VecZnxBigAddInplace,
  5        VecZnxBigAddSmallInplace, VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAllocBytes,
  6        VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace,
  7        VecZnxSwithcDegree, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPMatAlloc, VmpPrepare,
  8        ZnAddNormal, ZnFillUniform, ZnNormalizeInplace,
  9    },
 10    layouts::{Backend, Module, ScratchOwned, ZnxView},
 11    oep::{
 12        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
 13        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 14    },
 15    source::Source,
 16};
 17
 18use crate::layouts::{
 19    Infos, LWECiphertext, LWEPlaintext, LWESecret, LWESwitchingKey,
 20    prepared::{LWESwitchingKeyPrepared, PrepareAlloc},
 21};
 22
 23pub fn test_lwe_keyswitch<B>(module: &Module<B>)
 24where
 25    Module<B>: VecZnxDftAllocBytes
 26        + VecZnxBigNormalize<B>
 27        + DFT<B>
 28        + SvpApplyInplace<B>
 29        + IDFTConsume<B>
 30        + VecZnxFillUniform
 31        + VecZnxSubABInplace
 32        + VecZnxAddInplace
 33        + VecZnxNormalizeInplace<B>
 34        + VecZnxAddNormal
 35        + VecZnxNormalize<B>
 36        + VecZnxSub
 37        + SvpPrepare<B>
 38        + SvpPPolAllocBytes
 39        + SvpPPolAlloc<B>
 40        + VecZnxBigAllocBytes
 41        + VecZnxBigAddInplace<B>
 42        + VecZnxBigAddSmallInplace<B>
 43        + VecZnxNormalizeTmpBytes
 44        + VecZnxAddScalarInplace
 45        + VmpPMatAlloc<B>
 46        + VmpPrepare<B>
 47        + VmpApplyDftToDftTmpBytes
 48        + VmpApplyDftToDft<B>
 49        + VmpApplyDftToDftAdd<B>
 50        + VecZnxBigNormalizeTmpBytes
 51        + VecZnxSwithcDegree
 52        + VecZnxAutomorphismInplace
 53        + ZnNormalizeInplace<B>
 54        + ZnFillUniform
 55        + ZnAddNormal,
 56    B: Backend
 57        + TakeVecZnxDftImpl<B>
 58        + TakeVecZnxBigImpl<B>
 59        + TakeSvpPPolImpl<B>
 60        + ScratchOwnedAllocImpl<B>
 61        + ScratchOwnedBorrowImpl<B>
 62        + ScratchAvailableImpl<B>
 63        + TakeScalarZnxImpl<B>
 64        + TakeVecZnxImpl<B>,
 65{
 66    let n: usize = module.n();
 67    let basek: usize = 17;
 68
 69    let n_lwe_in: usize = 22;
 70    let n_lwe_out: usize = 30;
 71    let k_lwe_ct: usize = 2 * basek;
 72    let k_lwe_pt: usize = 8;
 73
 74    let k_ksk: usize = k_lwe_ct + basek;
 75
 76    let mut source_xs: Source = Source::new([0u8; 32]);
 77    let mut source_xa: Source = Source::new([0u8; 32]);
 78    let mut source_xe: Source = Source::new([0u8; 32]);
 79
 80    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 81        LWESwitchingKey::encrypt_sk_scratch_space(module, basek, k_ksk)
 82            | LWECiphertext::keyswitch_scratch_space(module, basek, k_lwe_ct, k_lwe_ct, k_ksk),
 83    );
 84
 85    let mut sk_lwe_in: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_in);
 86    sk_lwe_in.fill_ternary_prob(0.5, &mut source_xs);
 87
 88    let mut sk_lwe_out: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_out);
 89    sk_lwe_out.fill_ternary_prob(0.5, &mut source_xs);
 90
 91    let data: i64 = 17;
 92
 93    let mut lwe_pt_in: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
 94
 95    lwe_pt_in.encode_i64(data, k_lwe_pt);
 96
 97    let mut lwe_ct_in: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe_in, basek, k_lwe_ct);
 98    lwe_ct_in.encrypt_sk(
 99        module,
100        &lwe_pt_in,
101        &sk_lwe_in,
102        &mut source_xa,
103        &mut source_xe,
104    );
105
106    let mut ksk: LWESwitchingKey<Vec<u8>> = LWESwitchingKey::alloc(n, basek, k_ksk, lwe_ct_in.size());
107
108    ksk.encrypt_sk(
109        module,
110        &sk_lwe_in,
111        &sk_lwe_out,
112        &mut source_xa,
113        &mut source_xe,
114        scratch.borrow(),
115    );
116
117    let mut lwe_ct_out: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe_out, basek, k_lwe_ct);
118
119    let ksk_prepared: LWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
120
121    lwe_ct_out.keyswitch(module, &lwe_ct_in, &ksk_prepared, scratch.borrow());
122
123    let mut lwe_pt_out: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_ct);
124    lwe_ct_out.decrypt(module, &mut lwe_pt_out, &sk_lwe_out);
125
126    assert_eq!(lwe_pt_in.data.at(0, 0)[0], lwe_pt_out.data.at(0, 0)[0]);
127}