1use poulpy_core::{
 2    GGSWNoise, GLWEDecrypt, GLWEEncryptSk, GLWENoise, SIGMA, ScratchTakeCore,
 3    layouts::{GGSWInfos, GGSWLayout, GLWEInfos, GLWELayout, GLWESecretPreparedFactory, LWEInfos, prepared::GLWESecretPrepared},
 4};
 5use poulpy_hal::{
 6    api::{ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow},
 7    layouts::{Backend, Module, Scratch, ScratchOwned, Stats},
 8    source::Source,
 9};
10use rand::RngCore;
11
12use crate::tfhe::{
13    bdd_arithmetic::{
14        BDDKeyEncryptSk, BDDKeyPrepared, BDDKeyPreparedFactory, ExecuteBDDCircuit2WTo1W, FheUint, FheUintPrepare,
15        FheUintPrepareDebug, FheUintPreparedDebug, FheUintPreparedEncryptSk, FheUintPreparedFactory,
16        tests::test_suite::{TEST_BASE2K, TEST_GGSW_INFOS, TEST_GLWE_INFOS, TestContext},
17    },
18    blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory},
19};
20
21pub fn test_bdd_prepare<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
22where
23    Module<BE>: ModuleNew<BE>
24        + GLWESecretPreparedFactory<BE>
25        + GLWEDecrypt<BE>
26        + GLWENoise<BE>
27        + FheUintPreparedFactory<u32, BE>
28        + FheUintPreparedEncryptSk<u32, BE>
29        + FheUintPrepareDebug<BRA, u32, BE>
30        + BDDKeyEncryptSk<BRA, BE>
31        + BDDKeyPreparedFactory<BRA, BE>
32        + GGSWNoise<BE>
33        + FheUintPrepare<BRA, BE>
34        + ExecuteBDDCircuit2WTo1W<BE>
35        + GLWEEncryptSk<BE>,
36    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
37    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
38    Scratch<BE>: ScratchTakeCore<BE>,
39{
40    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
41    let ggsw_infos: GGSWLayout = TEST_GGSW_INFOS;
42
43    let module: &Module<BE> = &test_context.module;
44    let sk_glwe_prep: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
45    let bdd_key_prepared: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
46
47    let mut source: Source = Source::new([6u8; 32]);
48
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(value)
55    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_infos);
56    let value: u32 = source.next_u32();
57    c_enc.encrypt_sk(
58        module,
59        value,
60        sk_glwe_prep,
61        &mut source_xa,
62        &mut source_xe,
63        scratch.borrow(),
64    );
65
66    // GGSW(0)
67    let mut c_enc_prep_debug: FheUintPreparedDebug<Vec<u8>, u32> =
68        FheUintPreparedDebug::<Vec<u8>, u32>::alloc_from_infos(module, &ggsw_infos);
69
70    let mut scratch_2 = ScratchOwned::alloc(module.fhe_uint_prepare_tmp_bytes(7, 1, &c_enc_prep_debug, &c_enc, bdd_key_prepared));
71
72    // GGSW(value)
73    c_enc_prep_debug.prepare(module, &c_enc, bdd_key_prepared, scratch_2.borrow());
74
75    let max_noise = |col_i: usize| {
76        let mut noise: f64 = -(ggsw_infos.size() as f64 * TEST_BASE2K as f64) + SIGMA.log2() + 2.0;
77        noise += 0.5 * ggsw_infos.log_n() as f64;
78        if col_i != 0 {
79            noise += 0.5 * ggsw_infos.log_n() as f64
80        }
81        noise
82    };
83
84    for row in 0..c_enc_prep_debug.dnum().as_usize() {
85        for col in 0..c_enc_prep_debug.rank().as_usize() + 1 {
86            let stats: Vec<Stats> = c_enc_prep_debug.noise(module, row, col, value, sk_glwe_prep, scratch.borrow());
87            for (i, stat) in stats.iter().enumerate() {
88                let noise_have: f64 = stat.std().log2();
89                let noise_max: f64 = max_noise(col);
90                assert!(noise_have <= noise_max, "bit: {i} noise_have: {noise_have} > noise_max: {noise_max}")
91            }
92        }
93    }
94}