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