1use poulpy_core::{
 2    GGSWNoise, GLWEDecrypt, GLWEEncryptSk, GLWENoise, ScratchTakeCore,
 3    layouts::{GGSWLayout, GLWELayout, GLWESecret, GLWESecretPreparedFactory, LWEInfos, 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        BDDKeyEncryptSk, BDDKeyPreparedFactory, ExecuteBDDCircuit2WTo1W, FheUintBlockDebugPrepare, FheUintBlocks,
15        FheUintBlocksPrepare, FheUintBlocksPrepared, FheUintBlocksPreparedEncryptSk, FheUintBlocksPreparedFactory, Srl,
16        tests::test_suite::{TEST_GGSW_INFOS, TEST_GLWE_INFOS},
17    },
18    blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory},
19};
20
21pub fn test_bdd_srl<BRA: BlindRotationAlgo, BE: Backend>()
22where
23    Module<BE>: ModuleNew<BE>
24        + GLWESecretPreparedFactory<BE>
25        + GLWEDecrypt<BE>
26        + GLWENoise<BE>
27        + FheUintBlocksPreparedFactory<u32, BE>
28        + FheUintBlocksPreparedEncryptSk<u32, BE>
29        + FheUintBlockDebugPrepare<BRA, u32, BE>
30        + BDDKeyEncryptSk<BRA, BE>
31        + BDDKeyPreparedFactory<BRA, BE>
32        + GGSWNoise<BE>
33        + FheUintBlocksPrepare<BRA, u32, BE>
34        + ExecuteBDDCircuit2WTo1W<u32, 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 n_glwe: usize = glwe_infos.n().into();
44
45    let module: Module<BE> = Module::<BE>::new(n_glwe as u64);
46    let mut source: Source = Source::new([6u8; 32]);
47    let mut source_xs: Source = Source::new([1u8; 32]);
48    let mut source_xa: Source = Source::new([2u8; 32]);
49    let mut source_xe: Source = Source::new([3u8; 32]);
50
51    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
52
53    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_infos);
54    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
55    let mut sk_glwe_prep: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(&module, &glwe_infos);
56    sk_glwe_prep.prepare(&module, &sk_glwe);
57
58    let mut res: FheUintBlocks<Vec<u8>, u32> = FheUintBlocks::<Vec<u8>, u32>::alloc_from_infos(&module, &glwe_infos);
59    let mut a_enc_prep: FheUintBlocksPrepared<Vec<u8>, u32, BE> =
60        FheUintBlocksPrepared::<Vec<u8>, u32, BE>::alloc(&module, &ggsw_infos);
61    let mut b_enc_prep: FheUintBlocksPrepared<Vec<u8>, u32, BE> =
62        FheUintBlocksPrepared::<Vec<u8>, u32, BE>::alloc(&module, &ggsw_infos);
63
64    let a: u32 = source.next_u32();
65    let b: u32 = source.next_u32() & 15;
66
67    source.fill_bytes(&mut scratch.borrow().data);
68    a_enc_prep.encrypt_sk(
69        &module,
70        a,
71        &sk_glwe_prep,
72        &mut source_xa,
73        &mut source_xe,
74        scratch.borrow(),
75    );
76    source.fill_bytes(&mut scratch.borrow().data);
77    b_enc_prep.encrypt_sk(
78        &module,
79        b,
80        &sk_glwe_prep,
81        &mut source_xa,
82        &mut source_xe,
83        scratch.borrow(),
84    );
85
86    res.srl(&module, &a_enc_prep, &b_enc_prep, scratch.borrow());
87
88    assert_eq!(
89        res.decrypt(&module, &sk_glwe_prep, scratch.borrow()),
90        a >> b
91    );
92}