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