1use poulpy_backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApply, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
  4        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace,
  5        VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAdd, VecZnxDftAddInplace,
  6        VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftSubABInplace, VecZnxDftToVecZnxBig, VecZnxDftToVecZnxBigConsume,
  7        VecZnxDftToVecZnxBigTmpBytes, VecZnxDftZero, VecZnxFillUniform, VecZnxMulXpMinusOneInplace, VecZnxNormalize,
  8        VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxRotateInplace, VecZnxSub, VecZnxSubABInplace,
  9        VecZnxSwithcDegree, VmpApply, VmpApplyAdd, VmpApplyTmpBytes, VmpPMatAlloc, VmpPrepare, ZnxView,
 10    },
 11    layouts::{Backend, Module, ScratchOwned},
 12    oep::{
 13        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl,
 14        TakeVecZnxDftSliceImpl, TakeVecZnxImpl, TakeVecZnxSliceImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
 15    },
 16    source::Source,
 17};
 18
 19use crate::tfhe::blind_rotation::{
 20    BlincRotationExecute, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk, BlindRotationKeyPrepared, CGGI,
 21    LookUpTable, cggi_blind_rotate_scratch_space, mod_switch_2n,
 22};
 23
 24use poulpy_core::layouts::{
 25    GLWECiphertext, GLWEPlaintext, GLWESecret, Infos, LWECiphertext, LWECiphertextToRef, LWEPlaintext, LWESecret,
 26    prepared::{GLWESecretPrepared, PrepareAlloc},
 27};
 28
 29pub fn test_blind_rotation<B>(module: &Module<B>, n_lwe: usize, block_size: usize, extension_factor: usize)
 30where
 31    Module<B>: VecZnxBigAllocBytes
 32        + VecZnxDftAllocBytes
 33        + SvpPPolAllocBytes
 34        + VmpApplyTmpBytes
 35        + VecZnxBigNormalizeTmpBytes
 36        + VecZnxDftToVecZnxBigTmpBytes
 37        + VecZnxDftToVecZnxBig<B>
 38        + VecZnxDftAdd<B>
 39        + VecZnxDftAddInplace<B>
 40        + VecZnxDftFromVecZnx<B>
 41        + VecZnxDftZero<B>
 42        + SvpApply<B>
 43        + VecZnxDftSubABInplace<B>
 44        + VecZnxBigAddSmallInplace<B>
 45        + VecZnxRotate
 46        + VecZnxAddInplace
 47        + VecZnxSubABInplace
 48        + VecZnxNormalize<B>
 49        + VecZnxNormalizeInplace<B>
 50        + VecZnxCopy
 51        + VecZnxMulXpMinusOneInplace
 52        + SvpPrepare<B>
 53        + SvpPPolAlloc<B>
 54        + SvpApplyInplace<B>
 55        + VecZnxDftToVecZnxBigConsume<B>
 56        + VecZnxBigAddInplace<B>
 57        + VecZnxBigNormalize<B>
 58        + VecZnxNormalizeTmpBytes
 59        + VecZnxFillUniform
 60        + VecZnxAddNormal
 61        + VecZnxAddScalarInplace
 62        + VecZnxRotateInplace
 63        + VecZnxSwithcDegree
 64        + VecZnxSub
 65        + VmpPMatAlloc<B>
 66        + VmpPrepare<B>
 67        + VmpApply<B>
 68        + VmpApplyAdd<B>,
 69    B: Backend
 70        + VecZnxDftAllocBytesImpl<B>
 71        + VecZnxBigAllocBytesImpl<B>
 72        + ScratchOwnedAllocImpl<B>
 73        + ScratchOwnedBorrowImpl<B>
 74        + TakeVecZnxDftImpl<B>
 75        + TakeVecZnxBigImpl<B>
 76        + TakeVecZnxDftSliceImpl<B>
 77        + ScratchAvailableImpl<B>
 78        + TakeVecZnxImpl<B>
 79        + TakeVecZnxSliceImpl<B>,
 80{
 81    let n: usize = module.n();
 82    let basek: usize = 19;
 83    let k_lwe: usize = 24;
 84    let k_brk: usize = 3 * basek;
 85    let rows_brk: usize = 2; // Ensures first limb is noise-free.
 86    let k_lut: usize = basek;
 87    let k_res: usize = 2 * basek;
 88    let rank: usize = 1;
 89
 90    let log_message_modulus = 4;
 91
 92    let message_modulus: usize = 1 << log_message_modulus;
 93
 94    let mut source_xs: Source = Source::new([2u8; 32]);
 95    let mut source_xe: Source = Source::new([2u8; 32]);
 96    let mut source_xa: Source = Source::new([1u8; 32]);
 97
 98    let mut scratch: ScratchOwned<B> = ScratchOwned::<B>::alloc(BlindRotationKey::generate_from_sk_scratch_space(
 99        module, n, basek, k_brk, rank,
100    ));
101
102    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
103    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
104    let sk_glwe_dft: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch.borrow());
105
106    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
107    sk_lwe.fill_binary_block(block_size, &mut source_xs);
108
109    let mut scratch_br: ScratchOwned<B> = ScratchOwned::<B>::alloc(cggi_blind_rotate_scratch_space(
110        module,
111        n,
112        block_size,
113        extension_factor,
114        basek,
115        k_res,
116        k_brk,
117        rows_brk,
118        rank,
119    ));
120
121    let mut brk: BlindRotationKey<Vec<u8>, CGGI> =
122        BlindRotationKey::<Vec<u8>, CGGI>::alloc(n, n_lwe, basek, k_brk, rows_brk, rank);
123
124    brk.encrypt_sk(
125        module,
126        &sk_glwe_dft,
127        &sk_lwe,
128        &mut source_xa,
129        &mut source_xe,
130        3.2,
131        scratch.borrow(),
132    );
133
134    let mut lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe);
135
136    let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe);
137
138    let x: i64 = 15 % (message_modulus as i64);
139
140    pt_lwe.encode_i64(x, log_message_modulus + 1);
141
142    lwe.encrypt_sk(
143        module,
144        &pt_lwe,
145        &sk_lwe,
146        &mut source_xa,
147        &mut source_xe,
148        3.2,
149    );
150
151    let f = |x: i64| -> i64 { 2 * x + 1 };
152
153    let mut f_vec: Vec<i64> = vec![0i64; message_modulus];
154    f_vec
155        .iter_mut()
156        .enumerate()
157        .for_each(|(i, x)| *x = f(i as i64));
158
159    let mut lut: LookUpTable = LookUpTable::alloc(n, basek, k_lut, extension_factor);
160    lut.set(module, &f_vec, log_message_modulus + 1);
161
162    let mut res: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_res, rank);
163
164    let brk_prepared: BlindRotationKeyPrepared<Vec<u8>, CGGI, B> = brk.prepare_alloc(module, scratch.borrow());
165
166    brk_prepared.execute(module, &mut res, &lwe, &lut, scratch_br.borrow());
167
168    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_res);
169
170    res.decrypt(module, &mut pt_have, &sk_glwe_dft, scratch.borrow());
171
172    let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
173
174    mod_switch_2n(
175        2 * lut.domain_size(),
176        &mut lwe_2n,
177        &lwe.to_ref(),
178        lut.rotation_direction(),
179    );
180
181    let pt_want: i64 = (lwe_2n[0]
182        + lwe_2n[1..]
183            .iter()
184            .zip(sk_lwe.raw())
185            .map(|(x, y)| x * y)
186            .sum::<i64>())
187        & (2 * lut.domain_size() - 1) as i64;
188
189    lut.rotate(module, pt_want);
190
191    // First limb should be exactly equal (test are parameterized such that the noise does not reach
192    // the first limb)
193    assert_eq!(pt_have.data.at(0, 0), lut.data[0].at(0, 0));
194
195    // Verify that it effectively compute f(x)
196    let mut have: i64 = pt_have.decode_coeff_i64(log_message_modulus + 1, 0);
197
198    // Get positive representative and assert equality
199    have = (have + message_modulus as i64) % (message_modulus as i64);
200
201    assert_eq!(have, f(x) % (message_modulus as i64));
202}