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