1use itertools::Itertools;
  2use poulpy_core::{
  3    GGSWEncryptSk, GLWEDecrypt, GLWEEncryptSk,
  4    layouts::{GGSW, GGSWPrepared, GGSWPreparedFactory, GLWELayout, GLWESecretPrepared},
  5};
  6use poulpy_hal::{
  7    api::{ScratchOwnedAlloc, ScratchOwnedBorrow},
  8    layouts::{Backend, Module, ScalarZnx, Scratch, ScratchOwned, ZnxViewMut},
  9    source::Source,
 10};
 11use rand::RngCore;
 12
 13use crate::bin_fhe::{
 14    bdd_arithmetic::{
 15        Cswap, FheUint, FheUintPrepared, GLWEBlindRetrieval, GLWEBlindRetriever, ScratchTakeBDD,
 16        tests::test_suite::{TEST_GGSW_INFOS, TEST_GLWE_INFOS, TestContext},
 17    },
 18    blind_rotation::BlindRotationAlgo,
 19};
 20
 21pub fn test_fhe_uint_swap<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
 22where
 23    Module<BE>: GLWEEncryptSk<BE> + GLWEDecrypt<BE> + Cswap<BE> + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>,
 24    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 25    Scratch<BE>: ScratchTakeBDD<u32, BE>,
 26{
 27    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
 28    let ggsw_infos: poulpy_core::layouts::GGSWLayout = TEST_GGSW_INFOS;
 29
 30    let module: &Module<BE> = &test_context.module;
 31    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
 32
 33    let mut source_xa: Source = Source::new([2u8; 32]);
 34    let mut source_xe: Source = Source::new([3u8; 32]);
 35
 36    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
 37
 38    let mut s: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
 39    let mut s_prepared: GGSWPrepared<Vec<u8>, BE> = GGSWPrepared::alloc_from_infos(module, &ggsw_infos);
 40
 41    let a: u32 = source_xa.next_u32();
 42    let b: u32 = source_xa.next_u32();
 43
 44    for bit in [0, 1] {
 45        let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 46        let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 47
 48        a_enc.encrypt_sk(module, a, sk, &mut source_xa, &mut source_xe, scratch.borrow());
 49
 50        b_enc.encrypt_sk(module, b, sk, &mut source_xa, &mut source_xe, scratch.borrow());
 51
 52        let mut pt: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(module.n(), 1);
 53        pt.raw_mut()[0] = bit;
 54        s.encrypt_sk(module, &pt, sk, &mut source_xa, &mut source_xe, scratch.borrow());
 55        s_prepared.prepare(module, &s, scratch.borrow());
 56
 57        module.cswap(&mut a_enc, &mut b_enc, &s_prepared, scratch.borrow());
 58
 59        let (a_want, b_want) = if bit == 0 { (a, b) } else { (b, a) };
 60
 61        assert_eq!(a_want, a_enc.decrypt(module, sk, scratch.borrow()));
 62        assert_eq!(b_want, b_enc.decrypt(module, sk, scratch.borrow()));
 63    }
 64}
 65
 66pub fn test_glwe_blind_retrieval_statefull<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
 67where
 68    Module<BE>: GLWEEncryptSk<BE> + GLWEDecrypt<BE> + GLWEBlindRetrieval<BE> + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>,
 69    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 70    Scratch<BE>: ScratchTakeBDD<u32, BE>,
 71{
 72    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
 73    let ggsw_infos: poulpy_core::layouts::GGSWLayout = TEST_GGSW_INFOS;
 74
 75    let module: &Module<BE> = &test_context.module;
 76    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
 77
 78    let mut source_xa: Source = Source::new([2u8; 32]);
 79    let mut source_xe: Source = Source::new([3u8; 32]);
 80
 81    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
 82
 83    let data: Vec<u32> = (0..25).map(|i| i as u32).collect_vec();
 84
 85    let mut data_enc: Vec<FheUint<Vec<u8>, u32>> = (0..data.len())
 86        .map(|i| {
 87            let mut ct: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 88            ct.encrypt_sk(module, data[i], sk, &mut source_xa, &mut source_xe, scratch.borrow());
 89            ct
 90        })
 91        .collect_vec();
 92
 93    for idx in 0..data.len() as u32 {
 94        let mut idx_enc = FheUintPrepared::alloc_from_infos(module, &ggsw_infos);
 95        idx_enc.encrypt_sk(module, idx, sk, &mut source_xa, &mut source_xe, scratch.borrow());
 96
 97        module.glwe_blind_retrieval_statefull(&mut data_enc, &idx_enc, 0, 5, scratch.borrow());
 98
 99        assert_eq!(data[idx as usize], data_enc[0].decrypt(module, sk, scratch.borrow()));
100
101        module.glwe_blind_retrieval_statefull_rev(&mut data_enc, &idx_enc, 0, 5, scratch.borrow());
102
103        for i in 0..data.len() {
104            assert_eq!(data[i], data_enc[i].decrypt(module, sk, scratch.borrow()))
105        }
106    }
107}
108
109pub fn test_glwe_blind_retriever<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
110where
111    Module<BE>: GLWEEncryptSk<BE> + GLWEDecrypt<BE> + GLWEBlindRetrieval<BE> + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>,
112    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
113    Scratch<BE>: ScratchTakeBDD<u32, BE>,
114{
115    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
116    let ggsw_infos: poulpy_core::layouts::GGSWLayout = TEST_GGSW_INFOS;
117
118    let module: &Module<BE> = &test_context.module;
119    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
120
121    let mut source_xa: Source = Source::new([2u8; 32]);
122    let mut source_xe: Source = Source::new([3u8; 32]);
123
124    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
125
126    let data: Vec<u32> = (0..25).map(|i| i as u32).collect_vec();
127
128    let data_enc: Vec<FheUint<Vec<u8>, u32>> = (0..data.len())
129        .map(|i| {
130            let mut ct: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
131            ct.encrypt_sk(module, data[i], sk, &mut source_xa, &mut source_xe, scratch.borrow());
132            ct
133        })
134        .collect_vec();
135
136    let mut retriever: GLWEBlindRetriever = GLWEBlindRetriever::alloc(&glwe_infos, data.len());
137    for idx in 0..data.len() as u32 {
138        let offset = 2;
139        let mut idx_enc: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(module, &ggsw_infos);
140        idx_enc.encrypt_sk(module, idx << offset, sk, &mut source_xa, &mut source_xe, scratch.borrow());
141
142        let mut res: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_infos);
143        retriever.retrieve(module, &mut res, &data_enc, &idx_enc, offset, scratch.borrow());
144
145        assert_eq!(data[idx as usize], res.decrypt(module, sk, scratch.borrow()));
146    }
147}