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(
 49            module,
 50            a,
 51            sk,
 52            &mut source_xa,
 53            &mut source_xe,
 54            scratch.borrow(),
 55        );
 56
 57        b_enc.encrypt_sk(
 58            module,
 59            b,
 60            sk,
 61            &mut source_xa,
 62            &mut source_xe,
 63            scratch.borrow(),
 64        );
 65
 66        let mut pt: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(module.n(), 1);
 67        pt.raw_mut()[0] = bit;
 68        s.encrypt_sk(
 69            module,
 70            &pt,
 71            sk,
 72            &mut source_xa,
 73            &mut source_xe,
 74            scratch.borrow(),
 75        );
 76        s_prepared.prepare(module, &s, scratch.borrow());
 77
 78        module.cswap(&mut a_enc, &mut b_enc, &s_prepared, scratch.borrow());
 79
 80        let (a_want, b_want) = if bit == 0 { (a, b) } else { (b, a) };
 81
 82        assert_eq!(a_want, a_enc.decrypt(module, sk, scratch.borrow()));
 83        assert_eq!(b_want, b_enc.decrypt(module, sk, scratch.borrow()));
 84    }
 85}
 86
 87pub fn test_glwe_blind_retrieval_statefull<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
 88where
 89    Module<BE>: GLWEEncryptSk<BE> + GLWEDecrypt<BE> + GLWEBlindRetrieval<BE> + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>,
 90    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 91    Scratch<BE>: ScratchTakeBDD<u32, BE>,
 92{
 93    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
 94    let ggsw_infos: poulpy_core::layouts::GGSWLayout = TEST_GGSW_INFOS;
 95
 96    let module: &Module<BE> = &test_context.module;
 97    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
 98
 99    let mut source_xa: Source = Source::new([2u8; 32]);
100    let mut source_xe: Source = Source::new([3u8; 32]);
101
102    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
103
104    let data: Vec<u32> = (0..25).map(|i| i as u32).collect_vec();
105
106    let mut data_enc: Vec<FheUint<Vec<u8>, u32>> = (0..data.len())
107        .map(|i| {
108            let mut ct: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
109            ct.encrypt_sk(
110                module,
111                data[i],
112                sk,
113                &mut source_xa,
114                &mut source_xe,
115                scratch.borrow(),
116            );
117            ct
118        })
119        .collect_vec();
120
121    for idx in 0..data.len() as u32 {
122        let mut idx_enc = FheUintPrepared::alloc_from_infos(module, &ggsw_infos);
123        idx_enc.encrypt_sk(
124            module,
125            idx,
126            sk,
127            &mut source_xa,
128            &mut source_xe,
129            scratch.borrow(),
130        );
131
132        module.glwe_blind_retrieval_statefull(&mut data_enc, &idx_enc, 0, 5, scratch.borrow());
133
134        assert_eq!(
135            data[idx as usize],
136            data_enc[0].decrypt(module, sk, scratch.borrow())
137        );
138
139        module.glwe_blind_retrieval_statefull_rev(&mut data_enc, &idx_enc, 0, 5, scratch.borrow());
140
141        for i in 0..data.len() {
142            assert_eq!(data[i], data_enc[i].decrypt(module, sk, scratch.borrow()))
143        }
144    }
145}
146
147pub fn test_glwe_blind_retriever<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
148where
149    Module<BE>: GLWEEncryptSk<BE> + GLWEDecrypt<BE> + GLWEBlindRetrieval<BE> + GGSWEncryptSk<BE> + GGSWPreparedFactory<BE>,
150    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
151    Scratch<BE>: ScratchTakeBDD<u32, BE>,
152{
153    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
154    let ggsw_infos: poulpy_core::layouts::GGSWLayout = TEST_GGSW_INFOS;
155
156    let module: &Module<BE> = &test_context.module;
157    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
158
159    let mut source_xa: Source = Source::new([2u8; 32]);
160    let mut source_xe: Source = Source::new([3u8; 32]);
161
162    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
163
164    let data: Vec<u32> = (0..25).map(|i| i as u32).collect_vec();
165
166    let data_enc: Vec<FheUint<Vec<u8>, u32>> = (0..data.len())
167        .map(|i| {
168            let mut ct: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
169            ct.encrypt_sk(
170                module,
171                data[i],
172                sk,
173                &mut source_xa,
174                &mut source_xe,
175                scratch.borrow(),
176            );
177            ct
178        })
179        .collect_vec();
180
181    let mut retriever: GLWEBlindRetriever = GLWEBlindRetriever::alloc(&glwe_infos, data.len());
182    for idx in 0..data.len() as u32 {
183        let offset = 2;
184        let mut idx_enc: FheUintPrepared<Vec<u8>, u32, BE> = FheUintPrepared::alloc_from_infos(module, &ggsw_infos);
185        idx_enc.encrypt_sk(
186            module,
187            idx << offset,
188            sk,
189            &mut source_xa,
190            &mut source_xe,
191            scratch.borrow(),
192        );
193
194        let mut res: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_infos);
195        retriever.retrieve(
196            module,
197            &mut res,
198            &data_enc,
199            &idx_enc,
200            offset,
201            scratch.borrow(),
202        );
203
204        assert_eq!(
205            data[idx as usize],
206            res.decrypt(module, sk, scratch.borrow())
207        );
208    }
209}