1use std::time::Instant;
  2
  3use poulpy_backend::FFT64Ref;
  4use poulpy_core::{
  5    GGSWNoise, GLWEDecrypt, GLWEEncryptSk, GLWENoise, ScratchTakeCore,
  6    layouts::{GGSWLayout, GLWELayout, GLWESecret, GLWESecretPreparedFactory, LWEInfos, LWESecret, prepared::GLWESecretPrepared},
  7};
  8use poulpy_hal::{
  9    api::{ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow},
 10    layouts::{Backend, Module, Scratch, ScratchOwned},
 11    source::Source,
 12};
 13use rand::RngCore;
 14
 15use crate::tfhe::{
 16    bdd_arithmetic::{
 17        Add, BDDKey, BDDKeyEncryptSk, BDDKeyLayout, BDDKeyPrepared, BDDKeyPreparedFactory, ExecuteBDDCircuit2WTo1W,
 18        FheUintBlockDebugPrepare, FheUintBlocks, FheUintBlocksPrepare, FheUintBlocksPrepared, FheUintBlocksPreparedDebug,
 19        FheUintBlocksPreparedEncryptSk, FheUintBlocksPreparedFactory, Sub, TEST_BDD_KEY_LAYOUT, TEST_BLOCK_SIZE, TEST_GGSW_INFOS,
 20        TEST_GLWE_INFOS, TEST_N_LWE,
 21    },
 22    blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory, CGGI},
 23};
 24
 25#[test]
 26fn test_bdd_2w_to_1w_fft64_ref() {
 27    test_bdd_2w_to_1w::<FFT64Ref, CGGI>()
 28}
 29
 30fn test_bdd_2w_to_1w<BE: Backend, BRA: BlindRotationAlgo>()
 31where
 32    Module<BE>: ModuleNew<BE>
 33        + GLWESecretPreparedFactory<BE>
 34        + GLWEDecrypt<BE>
 35        + GLWENoise<BE>
 36        + FheUintBlocksPreparedFactory<u32, BE>
 37        + FheUintBlocksPreparedEncryptSk<u32, BE>
 38        + FheUintBlockDebugPrepare<BRA, u32, BE>
 39        + BDDKeyEncryptSk<BRA, BE>
 40        + BDDKeyPreparedFactory<BRA, BE>
 41        + GGSWNoise<BE>
 42        + FheUintBlocksPrepare<BRA, u32, BE>
 43        + ExecuteBDDCircuit2WTo1W<u32, BE>
 44        + GLWEEncryptSk<BE>,
 45    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
 46    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 47    Scratch<BE>: ScratchTakeCore<BE>,
 48{
 49    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
 50    let ggsw_infos: GGSWLayout = TEST_GGSW_INFOS;
 51
 52    let n_glwe: usize = glwe_infos.n().into();
 53
 54    let module: Module<BE> = Module::<BE>::new(n_glwe as u64);
 55    let mut source: Source = Source::new([6u8; 32]);
 56    let mut source_xs: Source = Source::new([1u8; 32]);
 57    let mut source_xa: Source = Source::new([2u8; 32]);
 58    let mut source_xe: Source = Source::new([3u8; 32]);
 59
 60    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
 61
 62    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_infos);
 63    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 64    let mut sk_glwe_prep: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(&module, &glwe_infos);
 65    sk_glwe_prep.prepare(&module, &sk_glwe);
 66
 67    let a: u32 = 645139204;
 68    let b: u32 = 0;
 69
 70    println!("a: {a}");
 71    println!("b: {b}");
 72
 73    let n_lwe: u32 = TEST_N_LWE;
 74    let block_size: u32 = TEST_BLOCK_SIZE;
 75
 76    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
 77    sk_lwe.fill_binary_block(block_size as usize, &mut source_xs);
 78
 79    let bdd_key_infos: BDDKeyLayout = TEST_BDD_KEY_LAYOUT;
 80
 81    let mut bdd_key: BDDKey<Vec<u8>, BRA> = BDDKey::alloc_from_infos(&bdd_key_infos);
 82
 83    let now: Instant = Instant::now();
 84    source.fill_bytes(&mut scratch.borrow().data);
 85    scratch.borrow().data.fill(0);
 86    bdd_key.encrypt_sk(
 87        &module,
 88        &sk_lwe,
 89        &sk_glwe,
 90        &mut source_xa,
 91        &mut source_xe,
 92        scratch.borrow(),
 93    );
 94    let mut bdd_key_prepared: BDDKeyPrepared<Vec<u8>, BRA, BE> = BDDKeyPrepared::alloc_from_infos(&module, &bdd_key_infos);
 95    source.fill_bytes(&mut scratch.borrow().data);
 96    bdd_key_prepared.prepare(&module, &bdd_key, scratch.borrow());
 97    println!("BDD-KGEN: {} ms", now.elapsed().as_millis());
 98
 99    let mut sum_enc: FheUintBlocks<Vec<u8>, u32> = FheUintBlocks::<Vec<u8>, u32>::alloc(&module, &glwe_infos);
100    let mut a_enc_prep: FheUintBlocksPrepared<Vec<u8>, u32, BE> =
101        FheUintBlocksPrepared::<Vec<u8>, u32, BE>::alloc(&module, &ggsw_infos);
102    let mut b_enc_prep: FheUintBlocksPrepared<Vec<u8>, u32, BE> =
103        FheUintBlocksPrepared::<Vec<u8>, u32, BE>::alloc(&module, &ggsw_infos);
104
105    source.fill_bytes(&mut scratch.borrow().data);
106    a_enc_prep.encrypt_sk(
107        &module,
108        a,
109        &sk_glwe_prep,
110        &mut source_xa,
111        &mut source_xe,
112        scratch.borrow(),
113    );
114    source.fill_bytes(&mut scratch.borrow().data);
115    b_enc_prep.encrypt_sk(
116        &module,
117        b,
118        &sk_glwe_prep,
119        &mut source_xa,
120        &mut source_xe,
121        scratch.borrow(),
122    );
123
124    // d + a
125    sum_enc.add(&module, &a_enc_prep, &b_enc_prep, scratch.borrow());
126
127    println!(
128        "other have: {:032b}",
129        sum_enc.decrypt(&module, &sk_glwe_prep, scratch.borrow())
130    );
131
132    println!("other want: {:032b}", b.wrapping_add(a));
133
134    // print a, b, and d
135    println!("a: {a}");
136    println!("b: {b}");
137}