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::bin_fhe::{
 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, GGSWInfos, GGSWLayout, GGSWPreparedFactory, GLWEAutomorphismKeyLayout, GLWEInfos,
 21        GLWESecretPreparedFactory, 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_res: usize = 15;
 50    let base2k_lwe: usize = 14;
 51    let base2k_brk: usize = 13;
 52    let base2k_tsk: usize = 12;
 53    let base2k_atk: usize = 11;
 54    let extension_factor: usize = 1;
 55    let rank: usize = 1;
 56
 57    let n_lwe: usize = 77;
 58    let k_lwe_pt: usize = 4;
 59    let k_lwe_ct: usize = 22;
 60    let block_size: usize = 7;
 61
 62    let k_ggsw_res: usize = 4 * base2k_res;
 63    let rows_ggsw_res: usize = 3;
 64
 65    let k_brk: usize = k_ggsw_res + base2k_brk;
 66    let rows_brk: usize = 4;
 67
 68    let k_atk: usize = k_ggsw_res + base2k_tsk;
 69    let rows_atk: usize = 4;
 70
 71    let k_tsk: usize = k_ggsw_res + base2k_atk;
 72    let rows_tsk: usize = 4;
 73
 74    let lwe_infos: LWELayout = LWELayout {
 75        n: n_lwe.into(),
 76        k: k_lwe_ct.into(),
 77        base2k: base2k_lwe.into(),
 78    };
 79
 80    let cbt_infos: CircuitBootstrappingKeyLayout = CircuitBootstrappingKeyLayout {
 81        brk_layout: BlindRotationKeyLayout {
 82            n_glwe: n_glwe.into(),
 83            n_lwe: n_lwe.into(),
 84            base2k: base2k_brk.into(),
 85            k: k_brk.into(),
 86            dnum: rows_brk.into(),
 87            rank: rank.into(),
 88        },
 89        atk_layout: GLWEAutomorphismKeyLayout {
 90            n: n_glwe.into(),
 91            base2k: base2k_atk.into(),
 92            k: k_atk.into(),
 93            dnum: rows_atk.into(),
 94            rank: rank.into(),
 95            dsize: Dsize(1),
 96        },
 97        tsk_layout: GGLWEToGGSWKeyLayout {
 98            n: n_glwe.into(),
 99            base2k: base2k_tsk.into(),
100            k: k_tsk.into(),
101            dnum: rows_tsk.into(),
102            dsize: Dsize(1),
103            rank: rank.into(),
104        },
105    };
106
107    let ggsw_infos: GGSWLayout = GGSWLayout {
108        n: n_glwe.into(),
109        base2k: base2k_res.into(),
110        k: k_ggsw_res.into(),
111        dnum: rows_ggsw_res.into(),
112        dsize: Dsize(1),
113        rank: rank.into(),
114    };
115
116    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 23);
117
118    let mut source_xs: Source = Source::new([1u8; 32]);
119    let mut source_xa: Source = Source::new([1u8; 32]);
120    let mut source_xe: Source = Source::new([1u8; 32]);
121
122    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
123    sk_lwe.fill_binary_block(block_size, &mut source_xs);
124
125    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe.into(), rank.into());
126    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
127
128    let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
129    sk_glwe_prepared.prepare(module, &sk_glwe);
130
131    let data: i64 = 1;
132
133    let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(base2k_lwe.into(), k_lwe_pt.into());
134    pt_lwe.encode_i64(data, (k_lwe_pt + 1).into());
135
136    println!("pt_lwe: {pt_lwe}");
137
138    let mut ct_lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
139    ct_lwe.encrypt_sk(
140        module,
141        &pt_lwe,
142        &sk_lwe,
143        &mut source_xa,
144        &mut source_xe,
145        scratch.borrow(),
146    );
147
148    let now: Instant = Instant::now();
149    let mut cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::alloc_from_infos(&cbt_infos);
150    println!("CBT-ALLOC: {} ms", now.elapsed().as_millis());
151
152    let now: Instant = Instant::now();
153    cbt_key.encrypt_sk(
154        module,
155        &sk_lwe,
156        &sk_glwe,
157        &mut source_xa,
158        &mut source_xe,
159        scratch.borrow(),
160    );
161    println!("CBT-ENCRYPT: {} ms", now.elapsed().as_millis());
162
163    let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
164
165    let log_gap_out = 1;
166
167    let mut cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> =
168        CircuitBootstrappingKeyPrepared::alloc_from_infos(module, &cbt_infos);
169    cbt_prepared.prepare(module, &cbt_key, scratch.borrow());
170
171    let now: Instant = Instant::now();
172    cbt_prepared.execute_to_exponent(
173        module,
174        log_gap_out,
175        &mut res,
176        &ct_lwe,
177        k_lwe_pt,
178        extension_factor,
179        scratch.borrow(),
180    );
181    println!("CBT: {} ms", now.elapsed().as_millis());
182
183    // X^{data * 2^log_gap_out}
184    let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n_glwe, 1);
185    pt_ggsw.at_mut(0, 0)[0] = 1;
186    module.vec_znx_rotate_inplace(
187        data * (1 << log_gap_out),
188        &mut pt_ggsw.as_vec_znx_mut(),
189        0,
190        scratch.borrow(),
191    );
192
193    for row in 0..res.dnum().as_usize() {
194        for col in 0..res.rank().as_usize() + 1 {
195            println!(
196                "row:{row} col:{col} -> {}",
197                res.noise(
198                    module,
199                    row,
200                    col,
201                    &pt_ggsw,
202                    &sk_glwe_prepared,
203                    scratch.borrow()
204                )
205                .std()
206                .log2()
207            )
208        }
209    }
210    let mut ct_glwe: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&ggsw_infos);
211    let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
212    pt_glwe.data.at_mut(0, 0)[0] = 1 << (base2k_res - 2);
213
214    ct_glwe.encrypt_sk(
215        module,
216        &pt_glwe,
217        &sk_glwe_prepared,
218        &mut source_xa,
219        &mut source_xe,
220        scratch.borrow(),
221    );
222
223    let mut res_prepared: GGSWPrepared<Vec<u8>, BE> = GGSWPrepared::alloc_from_infos(module, &res);
224    res_prepared.prepare(module, &res, scratch.borrow());
225
226    ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
227
228    let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
229    ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
230
231    // Parameters are set such that the first limb should be noiseless.
232    let mut pt_want: Vec<i64> = vec![0i64; module.n()];
233    pt_want[data as usize * (1 << log_gap_out)] = pt_glwe.data.at(0, 0)[0];
234    assert_eq!(pt_res.data.at(0, 0), pt_want);
235}
236
237pub fn test_circuit_bootstrapping_to_constant<BE: Backend, M, BRA: BlindRotationAlgo>(module: &M)
238where
239    M: ModuleN
240        + GLWESecretPreparedFactory<BE>
241        + GLWEExternalProduct<BE>
242        + GLWEDecrypt<BE>
243        + LWEEncryptSk<BE>
244        + CircuitBootstrappingKeyEncryptSk<BRA, BE>
245        + CircuitBootstrappingKeyPreparedFactory<BRA, BE>
246        + CirtuitBootstrappingExecute<BRA, BE>
247        + GGSWPreparedFactory<BE>
248        + GGSWNoise<BE>
249        + GLWEEncryptSk<BE>
250        + VecZnxRotateInplace<BE>,
251    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, // TODO find a way to remove this bound or move it to CBT KEY
252    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
253    Scratch<BE>: ScratchTakeCore<BE>,
254{
255    let n_glwe: usize = module.n();
256    let base2k_res: usize = 15;
257    let base2k_lwe: usize = 14;
258    let base2k_brk: usize = 13;
259    let base2k_tsk: usize = 12;
260    let base2k_atk: usize = 11;
261    let extension_factor: usize = 1;
262    let rank: usize = 1;
263
264    let n_lwe: usize = 77;
265    let k_lwe_pt: usize = 1;
266    let k_lwe_ct: usize = 13;
267    let block_size: usize = 7;
268
269    let k_ggsw_res: usize = 4 * base2k_res;
270    let rows_ggsw_res: usize = 3;
271
272    let k_brk: usize = k_ggsw_res + base2k_brk;
273    let rows_brk: usize = 4;
274
275    let k_atk: usize = k_ggsw_res + base2k_tsk;
276    let rows_atk: usize = 4;
277
278    let k_tsk: usize = k_ggsw_res + base2k_atk;
279    let rows_tsk: usize = 4;
280
281    let lwe_infos: LWELayout = LWELayout {
282        n: n_lwe.into(),
283        k: k_lwe_ct.into(),
284        base2k: base2k_lwe.into(),
285    };
286
287    let cbt_infos: CircuitBootstrappingKeyLayout = CircuitBootstrappingKeyLayout {
288        brk_layout: BlindRotationKeyLayout {
289            n_glwe: n_glwe.into(),
290            n_lwe: n_lwe.into(),
291            base2k: base2k_brk.into(),
292            k: k_brk.into(),
293            dnum: rows_brk.into(),
294            rank: rank.into(),
295        },
296        atk_layout: GLWEAutomorphismKeyLayout {
297            n: n_glwe.into(),
298            base2k: base2k_atk.into(),
299            k: k_atk.into(),
300            dnum: rows_atk.into(),
301            rank: rank.into(),
302            dsize: Dsize(1),
303        },
304        tsk_layout: GGLWEToGGSWKeyLayout {
305            n: n_glwe.into(),
306            base2k: base2k_tsk.into(),
307            k: k_tsk.into(),
308            dnum: rows_tsk.into(),
309            dsize: Dsize(1),
310            rank: rank.into(),
311        },
312    };
313
314    let ggsw_infos: GGSWLayout = GGSWLayout {
315        n: n_glwe.into(),
316        base2k: base2k_res.into(),
317        k: k_ggsw_res.into(),
318        dnum: rows_ggsw_res.into(),
319        dsize: Dsize(1),
320        rank: rank.into(),
321    };
322
323    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 23);
324
325    let mut source_xs: Source = Source::new([1u8; 32]);
326    let mut source_xa: Source = Source::new([1u8; 32]);
327    let mut source_xe: Source = Source::new([1u8; 32]);
328
329    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
330    sk_lwe.fill_binary_block(block_size, &mut source_xs);
331
332    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe.into(), rank.into());
333    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
334
335    let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
336    sk_glwe_prepared.prepare(module, &sk_glwe);
337
338    let data: i64 = 1;
339
340    let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(base2k_lwe.into(), k_lwe_pt.into());
341    pt_lwe.encode_i64(data, (k_lwe_pt + 1).into());
342
343    println!("pt_lwe: {pt_lwe}");
344
345    let mut ct_lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
346    ct_lwe.encrypt_sk(
347        module,
348        &pt_lwe,
349        &sk_lwe,
350        &mut source_xa,
351        &mut source_xe,
352        scratch.borrow(),
353    );
354
355    let now: Instant = Instant::now();
356    let mut cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::alloc_from_infos(&cbt_infos);
357    println!("CBT-ALLOC: {} ms", now.elapsed().as_millis());
358
359    let now: Instant = Instant::now();
360    cbt_key.encrypt_sk(
361        module,
362        &sk_lwe,
363        &sk_glwe,
364        &mut source_xa,
365        &mut source_xe,
366        scratch.borrow(),
367    );
368    println!("CBT-ENCRYPT: {} ms", now.elapsed().as_millis());
369
370    let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
371
372    let mut cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> =
373        CircuitBootstrappingKeyPrepared::alloc_from_infos(module, &cbt_infos);
374    cbt_prepared.prepare(module, &cbt_key, scratch.borrow());
375
376    let now: Instant = Instant::now();
377    cbt_prepared.execute_to_constant(
378        module,
379        &mut res,
380        &ct_lwe,
381        k_lwe_pt,
382        extension_factor,
383        scratch.borrow(),
384    );
385    println!("CBT: {} ms", now.elapsed().as_millis());
386
387    // X^{data * 2^log_gap_out}
388    let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n_glwe, 1);
389    pt_ggsw.at_mut(0, 0)[0] = data;
390
391    for row in 0..res.dnum().as_usize() {
392        for col in 0..res.rank().as_usize() + 1 {
393            println!(
394                "row:{row} col:{col} -> {}",
395                res.noise(
396                    module,
397                    row,
398                    col,
399                    &pt_ggsw,
400                    &sk_glwe_prepared,
401                    scratch.borrow()
402                )
403                .std()
404                .log2()
405            )
406        }
407    }
408
409    let mut ct_glwe: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&ggsw_infos);
410    let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
411    pt_glwe.data.at_mut(0, 0)[0] = 1 << (base2k_res - k_lwe_pt - 1);
412
413    ct_glwe.encrypt_sk(
414        module,
415        &pt_glwe,
416        &sk_glwe_prepared,
417        &mut source_xa,
418        &mut source_xe,
419        scratch.borrow(),
420    );
421
422    let mut res_prepared: GGSWPrepared<Vec<u8>, BE> = GGSWPrepared::alloc_from_infos(module, &res);
423    res_prepared.prepare(module, &res, scratch.borrow());
424
425    ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
426
427    let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
428    ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
429
430    // Parameters are set such that the first limb should be noiseless.
431    let mut pt_want: Vec<i64> = vec![0i64; module.n()];
432    pt_want[0] = pt_glwe.data.at(0, 0)[0] * data;
433    println!("pt_res: {pt_res}");
434    assert_eq!(pt_res.data.at(0, 0), pt_want);
435}