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_splice_u8<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    let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 38    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 39
 40    let a: u32 = 0xFFFFFFFF;
 41    let b: u32 = 0xAABBCCDD;
 42
 43    b_enc.encrypt_sk(
 44        module,
 45        b,
 46        sk,
 47        &mut source_xa,
 48        &mut source_xe,
 49        scratch.borrow(),
 50    );
 51    a_enc.encrypt_sk(
 52        module,
 53        a,
 54        sk,
 55        &mut source_xa,
 56        &mut source_xe,
 57        scratch.borrow(),
 58    );
 59
 60    for dst in 0..4 {
 61        for src in 0..4 {
 62            c_enc.splice_u8(module, dst, src, &a_enc, &b_enc, keys, scratch.borrow());
 63
 64            let rj: u32 = (dst << 3) as u32;
 65            let ri: u32 = (src << 3) as u32;
 66            let a_r: u32 = a.rotate_right(rj);
 67            let b_r: u32 = b.rotate_right(ri);
 68
 69            let c_want: u32 = ((a_r & 0xFFFF_FF00) | (b_r & 0x0000_00FF)).rotate_left(rj);
 70
 71            assert_eq!(c_want, c_enc.decrypt(module, sk, scratch.borrow()));
 72        }
 73    }
 74}
 75
 76pub fn test_fhe_uint_splice_u16<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
 77where
 78    Module<BE>: GLWEEncryptSk<BE> + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWEDecrypt<BE>,
 79    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 80    Scratch<BE>: ScratchTakeBDD<u32, BE>,
 81{
 82    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
 83
 84    let module: &Module<BE> = &test_context.module;
 85    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
 86    let keys: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
 87
 88    let mut source_xa: Source = Source::new([2u8; 32]);
 89    let mut source_xe: Source = Source::new([3u8; 32]);
 90
 91    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
 92
 93    let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 94    let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 95    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
 96
 97    let a: u32 = 0xFFFFFFFF;
 98    let b: u32 = 0xAABBCCDD;
 99
100    b_enc.encrypt_sk(
101        module,
102        b,
103        sk,
104        &mut source_xa,
105        &mut source_xe,
106        scratch.borrow(),
107    );
108    a_enc.encrypt_sk(
109        module,
110        a,
111        sk,
112        &mut source_xa,
113        &mut source_xe,
114        scratch.borrow(),
115    );
116
117    for dst in 0..2 {
118        for src in 0..2 {
119            c_enc.splice_u16(module, dst, src, &a_enc, &b_enc, keys, scratch.borrow());
120            let rj: u32 = (dst << 4) as u32;
121            let ri: u32 = (src << 4) as u32;
122            let a_r: u32 = a.rotate_right(rj);
123            let b_r: u32 = b.rotate_right(ri);
124            let c_want: u32 = ((a_r & 0xFFFF_0000) | (b_r & 0x0000_FFFF)).rotate_left(rj);
125            assert_eq!(c_want, c_enc.decrypt(module, sk, scratch.borrow()));
126        }
127    }
128}