1use poulpy_core::{
  2    GLWEAdd, GLWEDecrypt, GLWEEncryptSk, GLWERotate, GLWESub, GLWETrace,
  3    layouts::{GLWELayout, GLWESecretPrepared},
  4};
  5use poulpy_hal::{
  6    api::{ScratchOwnedAlloc, ScratchOwnedBorrow},
  7    layouts::{Backend, Module, Scratch, ScratchOwned},
  8    source::Source,
  9};
 10use rand::RngCore;
 11
 12use crate::tfhe::{
 13    bdd_arithmetic::{
 14        BDDKeyPrepared, FheUint, ScratchTakeBDD, ToBits,
 15        tests::test_suite::{TEST_GLWE_INFOS, TestContext},
 16    },
 17    blind_rotation::BlindRotationAlgo,
 18};
 19
 20pub fn test_fhe_uint_sext<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
 21where
 22    Module<BE>: GLWEEncryptSk<BE> + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWEDecrypt<BE>,
 23    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 24    Scratch<BE>: ScratchTakeBDD<u32, BE>,
 25{
 26    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
 27
 28    let module: &Module<BE> = &test_context.module;
 29    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
 30    let keys: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
 31
 32    let mut source_xa: Source = Source::new([2u8; 32]);
 33    let mut source_xe: Source = Source::new([3u8; 32]);
 34
 35    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
 36
 37    let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 38
 39    for j in 0..3 {
 40        let a: u32 = 0x8483_8281;
 41        a_enc.encrypt_sk(
 42            module,
 43            a,
 44            sk,
 45            &mut source_xa,
 46            &mut source_xe,
 47            scratch.borrow(),
 48        );
 49
 50        a_enc.sext(module, j, keys, scratch.borrow());
 51
 52        assert_eq!(
 53            sext(a, ((1 + j as u32) << 3) - 1),
 54            a_enc.decrypt(module, sk, scratch.borrow())
 55        );
 56    }
 57
 58    for j in 0..3 {
 59        let a: u32 = 0x4443_4241;
 60        a_enc.encrypt_sk(
 61            module,
 62            a,
 63            sk,
 64            &mut source_xa,
 65            &mut source_xe,
 66            scratch.borrow(),
 67        );
 68
 69        a_enc.sext(module, j, keys, scratch.borrow());
 70
 71        assert_eq!(
 72            sext(a, ((1 + j as u32) << 3) - 1),
 73            a_enc.decrypt(module, sk, scratch.borrow())
 74        );
 75    }
 76}
 77
 78pub(crate) fn sext(x: u32, bits: u32) -> u32 {
 79    let lo: u32 = x << (u32::BITS - bits) >> (u32::BITS - bits);
 80    let hi: u32 = ((x >> bits) & 1) * (0xFFFF_FFFF << bits);
 81    hi | lo
 82}
 83
 84pub fn test_fhe_uint_splice_u8<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
 85where
 86    Module<BE>: GLWEEncryptSk<BE> + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWEDecrypt<BE>,
 87    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 88    Scratch<BE>: ScratchTakeBDD<u32, BE>,
 89{
 90    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
 91
 92    let module: &Module<BE> = &test_context.module;
 93    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
 94    let keys: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
 95
 96    let mut source_xa: Source = Source::new([2u8; 32]);
 97    let mut source_xe: Source = Source::new([3u8; 32]);
 98
 99    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
100
101    let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
102    let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
103    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
104
105    let a: u32 = 0xFFFFFFFF;
106    let b: u32 = 0xAABBCCDD;
107
108    b_enc.encrypt_sk(
109        module,
110        b,
111        sk,
112        &mut source_xa,
113        &mut source_xe,
114        scratch.borrow(),
115    );
116    a_enc.encrypt_sk(
117        module,
118        a,
119        sk,
120        &mut source_xa,
121        &mut source_xe,
122        scratch.borrow(),
123    );
124
125    for dst in 0..4 {
126        for src in 0..4 {
127            c_enc.splice_u8(module, dst, src, &a_enc, &b_enc, keys, scratch.borrow());
128
129            let rj: u32 = (dst << 3) as u32;
130            let ri: u32 = (src << 3) as u32;
131            let a_r: u32 = a.rotate_right(rj);
132            let b_r: u32 = b.rotate_right(ri);
133
134            let c_want: u32 = ((a_r & 0xFFFF_FF00) | (b_r & 0x0000_00FF)).rotate_left(rj);
135
136            assert_eq!(c_want, c_enc.decrypt(module, sk, scratch.borrow()));
137        }
138    }
139}
140
141pub fn test_fhe_uint_splice_u16<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
142where
143    Module<BE>: GLWEEncryptSk<BE> + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWEDecrypt<BE>,
144    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
145    Scratch<BE>: ScratchTakeBDD<u32, BE>,
146{
147    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
148
149    let module: &Module<BE> = &test_context.module;
150    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
151    let keys: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
152
153    let mut source_xa: Source = Source::new([2u8; 32]);
154    let mut source_xe: Source = Source::new([3u8; 32]);
155
156    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
157
158    let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
159    let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
160    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
161
162    let a: u32 = 0xFFFFFFFF;
163    let b: u32 = 0xAABBCCDD;
164
165    b_enc.encrypt_sk(
166        module,
167        b,
168        sk,
169        &mut source_xa,
170        &mut source_xe,
171        scratch.borrow(),
172    );
173    a_enc.encrypt_sk(
174        module,
175        a,
176        sk,
177        &mut source_xa,
178        &mut source_xe,
179        scratch.borrow(),
180    );
181
182    for dst in 0..2 {
183        for src in 0..2 {
184            c_enc.splice_u16(module, dst, src, &a_enc, &b_enc, keys, scratch.borrow());
185            let rj: u32 = (dst << 4) as u32;
186            let ri: u32 = (src << 4) as u32;
187            let a_r: u32 = a.rotate_right(rj);
188            let b_r: u32 = b.rotate_right(ri);
189            let c_want: u32 = ((a_r & 0xFFFF_0000) | (b_r & 0x0000_FFFF)).rotate_left(rj);
190            assert_eq!(c_want, c_enc.decrypt(module, sk, scratch.borrow()));
191        }
192    }
193}
194
195pub fn test_fhe_uint_get_bit_glwe<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
196where
197    Module<BE>: GLWEEncryptSk<BE> + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWEDecrypt<BE>,
198    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
199    Scratch<BE>: ScratchTakeBDD<u32, BE>,
200{
201    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
202
203    let module: &Module<BE> = &test_context.module;
204    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
205    let keys: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
206
207    let mut source_xa: Source = Source::new([2u8; 32]);
208    let mut source_xe: Source = Source::new([3u8; 32]);
209
210    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
211
212    let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
213    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
214
215    let a: u32 = source_xa.next_u32();
216
217    a_enc.encrypt_sk(
218        module,
219        a,
220        sk,
221        &mut source_xa,
222        &mut source_xe,
223        scratch.borrow(),
224    );
225
226    for i in 0..32 {
227        a_enc.get_bit_glwe(module, i, &mut c_enc, keys, scratch.borrow());
228        assert_eq!(a.bit(i) as u32, c_enc.decrypt(module, sk, scratch.borrow()));
229    }
230}