1use poulpy_core::{
  2    GGSWNoise, GLWEDecrypt, GLWEEncryptSk, GLWENoise, SIGMA, ScratchTakeCore,
  3    layouts::{GGSWLayout, GLWELayout, GLWESecret, GLWESecretPreparedFactory, LWEInfos, LWESecret, prepared::GLWESecretPrepared},
  4};
  5use poulpy_hal::{
  6    api::{ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow},
  7    layouts::{Backend, Module, Scratch, ScratchOwned},
  8    source::Source,
  9};
 10use rand::RngCore;
 11
 12use crate::tfhe::{
 13    bdd_arithmetic::{
 14        BDDKey, BDDKeyEncryptSk, BDDKeyLayout, BDDKeyPreparedFactory, BddKeyPrepared, ExecuteBDDCircuit2WTo1W, FheUint,
 15        FheUintBlockDebugPrepare, FheUintBlocksPrepare, FheUintBlocksPreparedEncryptSk, FheUintBlocksPreparedFactory,
 16        FheUintPreparedDebug,
 17        tests::test_suite::{TEST_BASE2K, TEST_BDD_KEY_LAYOUT, TEST_BLOCK_SIZE, TEST_GGSW_INFOS, TEST_GLWE_INFOS, TEST_N_LWE},
 18    },
 19    blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory},
 20};
 21
 22pub fn test_bdd_prepare<BRA: BlindRotationAlgo, BE: Backend>()
 23where
 24    Module<BE>: ModuleNew<BE>
 25        + GLWESecretPreparedFactory<BE>
 26        + GLWEDecrypt<BE>
 27        + GLWENoise<BE>
 28        + FheUintBlocksPreparedFactory<u32, BE>
 29        + FheUintBlocksPreparedEncryptSk<u32, BE>
 30        + FheUintBlockDebugPrepare<BRA, u32, BE>
 31        + BDDKeyEncryptSk<BRA, BE>
 32        + BDDKeyPreparedFactory<BRA, BE>
 33        + GGSWNoise<BE>
 34        + FheUintBlocksPrepare<BRA, u32, BE>
 35        + ExecuteBDDCircuit2WTo1W<u32, BE>
 36        + GLWEEncryptSk<BE>,
 37    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
 38    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 39    Scratch<BE>: ScratchTakeCore<BE>,
 40{
 41    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
 42    let ggsw_infos: GGSWLayout = TEST_GGSW_INFOS;
 43
 44    let n_glwe: usize = glwe_infos.n().into();
 45
 46    let module: Module<BE> = Module::<BE>::new(n_glwe as u64);
 47    let mut source: Source = Source::new([6u8; 32]);
 48    let mut source_xs: Source = Source::new([1u8; 32]);
 49    let mut source_xa: Source = Source::new([2u8; 32]);
 50    let mut source_xe: Source = Source::new([3u8; 32]);
 51
 52    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
 53
 54    // GLWE Secret
 55    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_infos);
 56    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 57
 58    let mut sk_glwe_prep: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(&module, &glwe_infos);
 59    sk_glwe_prep.prepare(&module, &sk_glwe);
 60
 61    // LWE Secret
 62    let n_lwe: u32 = TEST_N_LWE;
 63    let block_size: u32 = TEST_BLOCK_SIZE;
 64    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
 65    sk_lwe.fill_binary_block(block_size as usize, &mut source_xs);
 66
 67    // CBT KEY
 68    let bdd_key_infos: BDDKeyLayout = TEST_BDD_KEY_LAYOUT;
 69
 70    let mut bdd_key: BDDKey<Vec<u8>, BRA> = BDDKey::alloc_from_infos(&bdd_key_infos);
 71
 72    source.fill_bytes(&mut scratch.borrow().data);
 73    scratch.borrow().data.fill(0);
 74    bdd_key.encrypt_sk(
 75        &module,
 76        &sk_lwe,
 77        &sk_glwe,
 78        &mut source_xa,
 79        &mut source_xe,
 80        scratch.borrow(),
 81    );
 82    let mut bdd_key_prepared: BddKeyPrepared<Vec<u8>, BRA, BE> = BddKeyPrepared::alloc_from_infos(&module, &bdd_key_infos);
 83    source.fill_bytes(&mut scratch.borrow().data);
 84    bdd_key_prepared.prepare(&module, &bdd_key, scratch.borrow());
 85
 86    // GLWE(value)
 87    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&module, &glwe_infos);
 88    let value: u32 = source.next_u32();
 89    c_enc.encrypt_sk(
 90        &module,
 91        value,
 92        &sk_glwe_prep,
 93        &mut source_xa,
 94        &mut source_xe,
 95        scratch.borrow(),
 96    );
 97
 98    // GGSW(0)
 99    let mut c_enc_prep_debug: FheUintPreparedDebug<Vec<u8>, u32> =
100        FheUintPreparedDebug::<Vec<u8>, u32>::alloc_from_infos(&module, &ggsw_infos);
101
102    // GGSW(value)
103    c_enc_prep_debug.prepare(&module, &c_enc, &bdd_key_prepared, scratch.borrow());
104
105    let max_noise = |col_i: usize| {
106        let mut noise: f64 = -(ggsw_infos.size() as f64 * TEST_BASE2K as f64) + SIGMA.log2() + 1.0;
107        noise += 0.5 * ggsw_infos.log_n() as f64;
108        if col_i != 0 {
109            noise += 0.5 * ggsw_infos.log_n() as f64
110        }
111        noise
112    };
113
114    // c_enc_prep_debug.print_noise(&module, &sk_glwe_prep, value);
115
116    c_enc_prep_debug.assert_noise(&module, &sk_glwe_prep, value, &max_noise);
117}