1use std::time::Instant;
  2
  3use poulpy_hal::{
  4    api::{ModuleN, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxRotateInplace},
  5    layouts::{Backend, ScalarZnx, Scratch, ScratchOwned, ZnxView, ZnxViewMut},
  6    source::Source,
  7};
  8
  9use crate::tfhe::{
 10    blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory, BlindRotationKeyLayout},
 11    circuit_bootstrapping::{
 12        CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
 13        CircuitBootstrappingKeyPrepared, CircuitBootstrappingKeyPreparedFactory, CirtuitBootstrappingExecute,
 14    },
 15};
 16
 17use poulpy_core::{
 18    GGSWNoise, GLWEDecrypt, GLWEEncryptSk, GLWEExternalProduct, LWEEncryptSk, ScratchTakeCore,
 19    layouts::{
 20        Dsize, GGLWEToGGSWKeyLayout, GGSWLayout, GGSWPreparedFactory, GLWEAutomorphismKeyLayout, GLWESecretPreparedFactory,
 21        LWELayout,
 22    },
 23};
 24
 25use poulpy_core::layouts::{
 26    GGSW, GLWE, GLWEPlaintext, GLWESecret, LWE, LWEPlaintext, LWESecret,
 27    prepared::{GGSWPrepared, GLWESecretPrepared},
 28};
 29
 30pub fn test_circuit_bootstrapping_to_exponent<BE: Backend, M, BRA: BlindRotationAlgo>(module: &M)
 31where
 32    M: ModuleN
 33        + GLWESecretPreparedFactory<BE>
 34        + GLWEExternalProduct<BE>
 35        + GLWEDecrypt<BE>
 36        + LWEEncryptSk<BE>
 37        + CircuitBootstrappingKeyEncryptSk<BRA, BE>
 38        + CircuitBootstrappingKeyPreparedFactory<BRA, BE>
 39        + CirtuitBootstrappingExecute<BRA, BE>
 40        + GGSWPreparedFactory<BE>
 41        + GGSWNoise<BE>
 42        + GLWEEncryptSk<BE>
 43        + VecZnxRotateInplace<BE>,
 44    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, // TODO find a way to remove this bound or move it to CBT KEY
 45    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 46    Scratch<BE>: ScratchTakeCore<BE>,
 47{
 48    let n_glwe: usize = module.n();
 49    let base2k: usize = 17;
 50    let extension_factor: usize = 1;
 51    let rank: usize = 1;
 52
 53    let n_lwe: usize = 77;
 54    let k_lwe_pt: usize = 4;
 55    let k_lwe_ct: usize = 22;
 56    let block_size: usize = 7;
 57
 58    let k_brk: usize = 5 * base2k;
 59    let rows_brk: usize = 4;
 60
 61    let k_atk: usize = 5 * base2k;
 62    let rows_atk: usize = 4;
 63
 64    let k_tsk: usize = 5 * base2k;
 65    let rows_tsk: usize = 4;
 66
 67    let k_ggsw_res: usize = 4 * base2k;
 68    let rows_ggsw_res: usize = 2;
 69
 70    let lwe_infos: LWELayout = LWELayout {
 71        n: n_lwe.into(),
 72        k: k_lwe_ct.into(),
 73        base2k: base2k.into(),
 74    };
 75
 76    let cbt_infos: CircuitBootstrappingKeyLayout = CircuitBootstrappingKeyLayout {
 77        layout_brk: BlindRotationKeyLayout {
 78            n_glwe: n_glwe.into(),
 79            n_lwe: n_lwe.into(),
 80            base2k: base2k.into(),
 81            k: k_brk.into(),
 82            dnum: rows_brk.into(),
 83            rank: rank.into(),
 84        },
 85        layout_atk: GLWEAutomorphismKeyLayout {
 86            n: n_glwe.into(),
 87            base2k: base2k.into(),
 88            k: k_atk.into(),
 89            dnum: rows_atk.into(),
 90            rank: rank.into(),
 91            dsize: Dsize(1),
 92        },
 93        layout_tsk: GGLWEToGGSWKeyLayout {
 94            n: n_glwe.into(),
 95            base2k: base2k.into(),
 96            k: k_tsk.into(),
 97            dnum: rows_tsk.into(),
 98            dsize: Dsize(1),
 99            rank: rank.into(),
100        },
101    };
102
103    let ggsw_infos: GGSWLayout = GGSWLayout {
104        n: n_glwe.into(),
105        base2k: base2k.into(),
106        k: k_ggsw_res.into(),
107        dnum: rows_ggsw_res.into(),
108        dsize: Dsize(1),
109        rank: rank.into(),
110    };
111
112    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 23);
113
114    let mut source_xs: Source = Source::new([1u8; 32]);
115    let mut source_xa: Source = Source::new([1u8; 32]);
116    let mut source_xe: Source = Source::new([1u8; 32]);
117
118    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
119    sk_lwe.fill_binary_block(block_size, &mut source_xs);
120
121    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe.into(), rank.into());
122    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
123
124    let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
125    sk_glwe_prepared.prepare(module, &sk_glwe);
126
127    let data: i64 = 1;
128
129    let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(base2k.into(), k_lwe_pt.into());
130    pt_lwe.encode_i64(data, (k_lwe_pt + 1).into());
131
132    println!("pt_lwe: {pt_lwe}");
133
134    let mut ct_lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
135    ct_lwe.encrypt_sk(module, &pt_lwe, &sk_lwe, &mut source_xa, &mut source_xe);
136
137    let now: Instant = Instant::now();
138    let mut cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::alloc_from_infos(&cbt_infos);
139    println!("CBT-ALLOC: {} ms", now.elapsed().as_millis());
140
141    let now: Instant = Instant::now();
142    cbt_key.encrypt_sk(
143        module,
144        &sk_lwe,
145        &sk_glwe,
146        &mut source_xa,
147        &mut source_xe,
148        scratch.borrow(),
149    );
150    println!("CBT-ENCRYPT: {} ms", now.elapsed().as_millis());
151
152    let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
153
154    let log_gap_out = 1;
155
156    let mut cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> =
157        CircuitBootstrappingKeyPrepared::alloc_from_infos(module, &cbt_infos);
158    cbt_prepared.prepare(module, &cbt_key, scratch.borrow());
159
160    let now: Instant = Instant::now();
161    cbt_prepared.execute_to_exponent(
162        module,
163        log_gap_out,
164        &mut res,
165        &ct_lwe,
166        k_lwe_pt,
167        extension_factor,
168        scratch.borrow(),
169    );
170    println!("CBT: {} ms", now.elapsed().as_millis());
171
172    // X^{data * 2^log_gap_out}
173    let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n_glwe, 1);
174    pt_ggsw.at_mut(0, 0)[0] = 1;
175    module.vec_znx_rotate_inplace(
176        data * (1 << log_gap_out),
177        &mut pt_ggsw.as_vec_znx_mut(),
178        0,
179        scratch.borrow(),
180    );
181
182    res.print_noise(module, &sk_glwe_prepared, &pt_ggsw);
183
184    let mut ct_glwe: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&ggsw_infos);
185    let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
186    pt_glwe.data.at_mut(0, 0)[0] = 1 << (base2k - 2);
187
188    ct_glwe.encrypt_sk(
189        module,
190        &pt_glwe,
191        &sk_glwe_prepared,
192        &mut source_xa,
193        &mut source_xe,
194        scratch.borrow(),
195    );
196
197    let mut res_prepared: GGSWPrepared<Vec<u8>, BE> = GGSWPrepared::alloc_from_infos(module, &res);
198    res_prepared.prepare(module, &res, scratch.borrow());
199
200    ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
201
202    let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
203    ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
204
205    // Parameters are set such that the first limb should be noiseless.
206    let mut pt_want: Vec<i64> = vec![0i64; module.n()];
207    pt_want[data as usize * (1 << log_gap_out)] = pt_glwe.data.at(0, 0)[0];
208    assert_eq!(pt_res.data.at(0, 0), pt_want);
209}
210
211pub fn test_circuit_bootstrapping_to_constant<BE: Backend, M, BRA: BlindRotationAlgo>(module: &M)
212where
213    M: ModuleN
214        + GLWESecretPreparedFactory<BE>
215        + GLWEExternalProduct<BE>
216        + GLWEDecrypt<BE>
217        + LWEEncryptSk<BE>
218        + CircuitBootstrappingKeyEncryptSk<BRA, BE>
219        + CircuitBootstrappingKeyPreparedFactory<BRA, BE>
220        + CirtuitBootstrappingExecute<BRA, BE>
221        + GGSWPreparedFactory<BE>
222        + GGSWNoise<BE>
223        + GLWEEncryptSk<BE>
224        + VecZnxRotateInplace<BE>,
225    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, // TODO find a way to remove this bound or move it to CBT KEY
226    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
227    Scratch<BE>: ScratchTakeCore<BE>,
228{
229    let n_glwe: usize = module.n();
230    let base2k: usize = 14;
231    let extension_factor: usize = 1;
232    let rank: usize = 1;
233
234    let n_lwe: usize = 77;
235    let k_lwe_pt: usize = 1;
236    let k_lwe_ct: usize = 13;
237    let block_size: usize = 7;
238
239    let k_brk: usize = 5 * base2k;
240    let rows_brk: usize = 3;
241
242    let k_atk: usize = 5 * base2k;
243    let rows_atk: usize = 4;
244
245    let k_tsk: usize = 5 * base2k;
246    let rows_tsk: usize = 4;
247
248    let k_ggsw_res: usize = 4 * base2k;
249    let rows_ggsw_res: usize = 3;
250
251    let lwe_infos: LWELayout = LWELayout {
252        n: n_lwe.into(),
253        k: k_lwe_ct.into(),
254        base2k: base2k.into(),
255    };
256
257    let cbt_infos: CircuitBootstrappingKeyLayout = CircuitBootstrappingKeyLayout {
258        layout_brk: BlindRotationKeyLayout {
259            n_glwe: n_glwe.into(),
260            n_lwe: n_lwe.into(),
261            base2k: base2k.into(),
262            k: k_brk.into(),
263            dnum: rows_brk.into(),
264            rank: rank.into(),
265        },
266        layout_atk: GLWEAutomorphismKeyLayout {
267            n: n_glwe.into(),
268            base2k: base2k.into(),
269            k: k_atk.into(),
270            dnum: rows_atk.into(),
271            rank: rank.into(),
272            dsize: Dsize(1),
273        },
274        layout_tsk: GGLWEToGGSWKeyLayout {
275            n: n_glwe.into(),
276            base2k: base2k.into(),
277            k: k_tsk.into(),
278            dnum: rows_tsk.into(),
279            dsize: Dsize(1),
280            rank: rank.into(),
281        },
282    };
283
284    let ggsw_infos: GGSWLayout = GGSWLayout {
285        n: n_glwe.into(),
286        base2k: base2k.into(),
287        k: k_ggsw_res.into(),
288        dnum: rows_ggsw_res.into(),
289        dsize: Dsize(1),
290        rank: rank.into(),
291    };
292
293    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 23);
294
295    let mut source_xs: Source = Source::new([1u8; 32]);
296    let mut source_xa: Source = Source::new([1u8; 32]);
297    let mut source_xe: Source = Source::new([1u8; 32]);
298
299    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
300    sk_lwe.fill_binary_block(block_size, &mut source_xs);
301
302    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe.into(), rank.into());
303    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
304
305    let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
306    sk_glwe_prepared.prepare(module, &sk_glwe);
307
308    let data: i64 = 1;
309
310    let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(base2k.into(), k_lwe_pt.into());
311    pt_lwe.encode_i64(data, (k_lwe_pt + 1).into());
312
313    println!("pt_lwe: {pt_lwe}");
314
315    let mut ct_lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
316    ct_lwe.encrypt_sk(module, &pt_lwe, &sk_lwe, &mut source_xa, &mut source_xe);
317
318    let now: Instant = Instant::now();
319    let mut cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::alloc_from_infos(&cbt_infos);
320    println!("CBT-ALLOC: {} ms", now.elapsed().as_millis());
321
322    let now: Instant = Instant::now();
323    cbt_key.encrypt_sk(
324        module,
325        &sk_lwe,
326        &sk_glwe,
327        &mut source_xa,
328        &mut source_xe,
329        scratch.borrow(),
330    );
331    println!("CBT-ENCRYPT: {} ms", now.elapsed().as_millis());
332
333    let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
334
335    let mut cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> =
336        CircuitBootstrappingKeyPrepared::alloc_from_infos(module, &cbt_infos);
337    cbt_prepared.prepare(module, &cbt_key, scratch.borrow());
338
339    let now: Instant = Instant::now();
340    cbt_prepared.execute_to_constant(
341        module,
342        &mut res,
343        &ct_lwe,
344        k_lwe_pt,
345        extension_factor,
346        scratch.borrow(),
347    );
348    println!("CBT: {} ms", now.elapsed().as_millis());
349
350    // X^{data * 2^log_gap_out}
351    let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n_glwe, 1);
352    pt_ggsw.at_mut(0, 0)[0] = data;
353
354    res.print_noise(module, &sk_glwe_prepared, &pt_ggsw);
355
356    let mut ct_glwe: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&ggsw_infos);
357    let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
358    pt_glwe.data.at_mut(0, 0)[0] = 1 << (base2k - k_lwe_pt - 1);
359
360    ct_glwe.encrypt_sk(
361        module,
362        &pt_glwe,
363        &sk_glwe_prepared,
364        &mut source_xa,
365        &mut source_xe,
366        scratch.borrow(),
367    );
368
369    let mut res_prepared: GGSWPrepared<Vec<u8>, BE> = GGSWPrepared::alloc_from_infos(module, &res);
370    res_prepared.prepare(module, &res, scratch.borrow());
371
372    ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
373
374    let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
375    ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
376
377    // Parameters are set such that the first limb should be noiseless.
378    let mut pt_want: Vec<i64> = vec![0i64; module.n()];
379    pt_want[0] = pt_glwe.data.at(0, 0)[0] * data;
380    assert_eq!(pt_res.data.at(0, 0), pt_want);
381}