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