1use poulpy_hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
  4        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphismInplace, VecZnxBigAddInplace,
  5        VecZnxBigAddSmallInplace, VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxCopy,
  6        VecZnxDftAllocBytes, VecZnxDftApply, VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeInplace,
  7        VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing, VmpApplyDftToDft, VmpApplyDftToDftAdd,
  8        VmpApplyDftToDftTmpBytes, VmpPMatAlloc, VmpPrepare, 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    LWECiphertext, LWECiphertextLayout, LWEPlaintext, LWESecret, LWESwitchingKey, LWESwitchingKeyLayout,
 20    prepared::{LWESwitchingKeyPrepared, PrepareAlloc},
 21};
 22
 23pub fn test_lwe_keyswitch<B>(module: &Module<B>)
 24where
 25    Module<B>: VecZnxDftAllocBytes
 26        + VecZnxBigNormalize<B>
 27        + VecZnxDftApply<B>
 28        + SvpApplyDftToDftInplace<B>
 29        + VecZnxIdftApplyConsume<B>
 30        + VecZnxFillUniform
 31        + VecZnxSubInplace
 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        + VecZnxSwitchRing
 52        + VecZnxAutomorphismInplace<B>
 53        + ZnNormalizeInplace<B>
 54        + ZnFillUniform
 55        + ZnAddNormal
 56        + VecZnxCopy,
 57    B: Backend
 58        + TakeVecZnxDftImpl<B>
 59        + TakeVecZnxBigImpl<B>
 60        + TakeSvpPPolImpl<B>
 61        + ScratchOwnedAllocImpl<B>
 62        + ScratchOwnedBorrowImpl<B>
 63        + ScratchAvailableImpl<B>
 64        + TakeScalarZnxImpl<B>
 65        + TakeVecZnxImpl<B>,
 66{
 67    let n: usize = module.n();
 68    let base2k: usize = 17;
 69
 70    let n_lwe_in: usize = 22;
 71    let n_lwe_out: usize = 30;
 72    let k_lwe_ct: usize = 2 * base2k;
 73    let k_lwe_pt: usize = 8;
 74
 75    let k_ksk: usize = k_lwe_ct + base2k;
 76    let rows: usize = k_lwe_ct.div_ceil(base2k);
 77
 78    let mut source_xs: Source = Source::new([0u8; 32]);
 79    let mut source_xa: Source = Source::new([0u8; 32]);
 80    let mut source_xe: Source = Source::new([0u8; 32]);
 81
 82    let key_apply_infos: LWESwitchingKeyLayout = LWESwitchingKeyLayout {
 83        n: n.into(),
 84        base2k: base2k.into(),
 85        k: k_ksk.into(),
 86        rows: rows.into(),
 87    };
 88
 89    let lwe_in_infos: LWECiphertextLayout = LWECiphertextLayout {
 90        n: n_lwe_in.into(),
 91        base2k: base2k.into(),
 92        k: k_lwe_ct.into(),
 93    };
 94
 95    let lwe_out_infos: LWECiphertextLayout = LWECiphertextLayout {
 96        n: n_lwe_out.into(),
 97        k: k_lwe_ct.into(),
 98        base2k: base2k.into(),
 99    };
100
101    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
102        LWESwitchingKey::encrypt_sk_scratch_space(module, &key_apply_infos)
103            | LWECiphertext::keyswitch_scratch_space(module, &lwe_out_infos, &lwe_in_infos, &key_apply_infos),
104    );
105
106    let mut sk_lwe_in: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_in.into());
107    sk_lwe_in.fill_ternary_prob(0.5, &mut source_xs);
108
109    let mut sk_lwe_out: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_out.into());
110    sk_lwe_out.fill_ternary_prob(0.5, &mut source_xs);
111
112    let data: i64 = 17;
113
114    let mut lwe_pt_in: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc_with(base2k.into(), k_lwe_pt.into());
115    lwe_pt_in.encode_i64(data, k_lwe_pt.into());
116
117    let mut lwe_ct_in: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(&lwe_in_infos);
118    lwe_ct_in.encrypt_sk(
119        module,
120        &lwe_pt_in,
121        &sk_lwe_in,
122        &mut source_xa,
123        &mut source_xe,
124    );
125
126    let mut ksk: LWESwitchingKey<Vec<u8>> = LWESwitchingKey::alloc(&key_apply_infos);
127
128    ksk.encrypt_sk(
129        module,
130        &sk_lwe_in,
131        &sk_lwe_out,
132        &mut source_xa,
133        &mut source_xe,
134        scratch.borrow(),
135    );
136
137    let mut lwe_ct_out: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(&lwe_out_infos);
138
139    let ksk_prepared: LWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
140
141    lwe_ct_out.keyswitch(module, &lwe_ct_in, &ksk_prepared, scratch.borrow());
142
143    let mut lwe_pt_out: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(&lwe_out_infos);
144    lwe_ct_out.decrypt(module, &mut lwe_pt_out, &sk_lwe_out);
145
146    assert_eq!(lwe_pt_in.data.at(0, 0)[0], lwe_pt_out.data.at(0, 0)[0]);
147}