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};
 10
 11use crate::tfhe::{
 12    bdd_arithmetic::{
 13        BDDKeyPrepared, FheUint, ScratchTakeBDD,
 14        tests::test_suite::{TEST_GLWE_INFOS, TestContext},
 15    },
 16    blind_rotation::BlindRotationAlgo,
 17};
 18
 19pub fn test_fhe_uint_sext<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
 20where
 21    Module<BE>: GLWEEncryptSk<BE> + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWEDecrypt<BE>,
 22    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 23    Scratch<BE>: ScratchTakeBDD<u32, BE>,
 24{
 25    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
 26
 27    let module: &Module<BE> = &test_context.module;
 28    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
 29    let keys: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
 30
 31    let mut source_xa: Source = Source::new([2u8; 32]);
 32    let mut source_xe: Source = Source::new([3u8; 32]);
 33
 34    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
 35
 36    let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 37
 38    for j in 0..3 {
 39        for i in 0..32 {
 40            let a: u32 = 0xFFFFFFFF >> i;
 41
 42            a_enc.encrypt_sk(
 43                module,
 44                a,
 45                sk,
 46                &mut source_xa,
 47                &mut source_xe,
 48                scratch.borrow(),
 49            );
 50
 51            a_enc.sext(module, j, keys, scratch.borrow());
 52
 53            // println!("{:08x} -> {:08x} {:08x}", a, sext(a, j), a_enc.decrypt(module, sk, scratch.borrow()));
 54
 55            assert_eq!(sext(a, j), a_enc.decrypt(module, sk, scratch.borrow()));
 56        }
 57    }
 58}
 59
 60pub fn sext(x: u32, byte: usize) -> u32 {
 61    x | (((x >> (byte << 3)) & 1) * (0xFFFF_FFFF << (byte << 3)))
 62}
 63
 64pub fn test_fhe_uint_splice_u8<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
 65where
 66    Module<BE>: GLWEEncryptSk<BE> + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWEDecrypt<BE>,
 67    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 68    Scratch<BE>: ScratchTakeBDD<u32, BE>,
 69{
 70    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
 71
 72    let module: &Module<BE> = &test_context.module;
 73    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
 74    let keys: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
 75
 76    let mut source_xa: Source = Source::new([2u8; 32]);
 77    let mut source_xe: Source = Source::new([3u8; 32]);
 78
 79    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
 80
 81    let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 82    let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 83    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 84
 85    let a: u32 = 0xFFFFFFFF;
 86    let b: u32 = 0xAABBCCDD;
 87
 88    b_enc.encrypt_sk(
 89        module,
 90        b,
 91        sk,
 92        &mut source_xa,
 93        &mut source_xe,
 94        scratch.borrow(),
 95    );
 96    a_enc.encrypt_sk(
 97        module,
 98        a,
 99        sk,
100        &mut source_xa,
101        &mut source_xe,
102        scratch.borrow(),
103    );
104
105    for dst in 0..4 {
106        for src in 0..4 {
107            c_enc.splice_u8(module, dst, src, &a_enc, &b_enc, keys, scratch.borrow());
108
109            let rj: u32 = (dst << 3) as u32;
110            let ri: u32 = (src << 3) as u32;
111            let a_r: u32 = a.rotate_right(rj);
112            let b_r: u32 = b.rotate_right(ri);
113
114            let c_want: u32 = ((a_r & 0xFFFF_FF00) | (b_r & 0x0000_00FF)).rotate_left(rj);
115
116            assert_eq!(c_want, c_enc.decrypt(module, sk, scratch.borrow()));
117        }
118    }
119}
120
121pub fn test_fhe_uint_splice_u16<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
122where
123    Module<BE>: GLWEEncryptSk<BE> + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWEDecrypt<BE>,
124    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
125    Scratch<BE>: ScratchTakeBDD<u32, BE>,
126{
127    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
128
129    let module: &Module<BE> = &test_context.module;
130    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
131    let keys: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
132
133    let mut source_xa: Source = Source::new([2u8; 32]);
134    let mut source_xe: Source = Source::new([3u8; 32]);
135
136    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
137
138    let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
139    let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
140    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
141
142    let a: u32 = 0xFFFFFFFF;
143    let b: u32 = 0xAABBCCDD;
144
145    b_enc.encrypt_sk(
146        module,
147        b,
148        sk,
149        &mut source_xa,
150        &mut source_xe,
151        scratch.borrow(),
152    );
153    a_enc.encrypt_sk(
154        module,
155        a,
156        sk,
157        &mut source_xa,
158        &mut source_xe,
159        scratch.borrow(),
160    );
161
162    for dst in 0..2 {
163        for src in 0..2 {
164            c_enc.splice_u16(module, dst, src, &a_enc, &b_enc, keys, scratch.borrow());
165            let rj: u32 = (dst << 4) as u32;
166            let ri: u32 = (src << 4) as u32;
167            let a_r: u32 = a.rotate_right(rj);
168            let b_r: u32 = b.rotate_right(ri);
169            let c_want: u32 = ((a_r & 0xFFFF_0000) | (b_r & 0x0000_FFFF)).rotate_left(rj);
170            assert_eq!(c_want, c_enc.decrypt(module, sk, scratch.borrow()));
171        }
172    }
173}