1use poulpy_core::{
 2    GGSWNoise, GLWEDecrypt, GLWEEncryptSk, GLWENoise, ScratchTakeCore,
 3    layouts::{GGSWLayout, GLWELayout, GLWESecretPreparedFactory, 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, BDDKeyPrepared, BDDKeyPreparedFactory, ExecuteBDDCircuit2WTo1W, FheUint, FheUintPrepare,
15        FheUintPrepareDebug, FheUintPrepared, FheUintPreparedEncryptSk, FheUintPreparedFactory, Slt,
16        tests::test_suite::{TEST_GGSW_INFOS, TEST_GLWE_INFOS, TestContext},
17    },
18    blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory},
19};
20
21pub fn test_bdd_slt<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<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 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    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 res: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
54    let mut a_enc_prep: FheUintPrepared<Vec<u8>, u32, BE> =
55        FheUintPrepared::<Vec<u8>, u32, BE>::alloc_from_infos(module, &ggsw_infos);
56    let mut b_enc_prep: FheUintPrepared<Vec<u8>, u32, BE> =
57        FheUintPrepared::<Vec<u8>, u32, BE>::alloc_from_infos(module, &ggsw_infos);
58
59    let a: u32 = source.next_u32();
60    let b: u32 = source.next_u32();
61
62    source.fill_bytes(&mut scratch.borrow().data);
63    a_enc_prep.encrypt_sk(
64        module,
65        a,
66        sk_glwe_prep,
67        &mut source_xa,
68        &mut source_xe,
69        scratch.borrow(),
70    );
71    source.fill_bytes(&mut scratch.borrow().data);
72    b_enc_prep.encrypt_sk(
73        module,
74        b,
75        sk_glwe_prep,
76        &mut source_xa,
77        &mut source_xe,
78        scratch.borrow(),
79    );
80
81    // d + a
82    res.slt(
83        module,
84        &a_enc_prep,
85        &b_enc_prep,
86        bdd_key_prepared,
87        scratch.borrow(),
88    );
89
90    assert_eq!(
91        res.decrypt(module, sk_glwe_prep, scratch.borrow()),
92        ((a as i32) < (b as i32)) as u32
93    );
94}