Add splice u8 and u16

poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs

@@ -1,5 +1,5 @@
 use poulpy_core::{
-    GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEPacking, GLWERotate, LWEFromGLWE, ScratchTakeCore,
+    GLWEAdd, GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEPacking, GLWERotate, GLWESub, GLWETrace, LWEFromGLWE, ScratchTakeCore,
     layouts::{
         Base2K, Degree, GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWEPlaintextLayout, GLWESecretPreparedToRef, GLWEToMut,
         GLWEToRef, LWEInfos, LWEToMut, Rank, TorusPrecision,
@@ -170,54 +170,109 @@ impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
         module.glwe_copy(&mut self.bits, cts.remove(&0).unwrap());
     }
 
-    // pub fn copy_byte<D0, D1, BRA, M, BE: Backend>(
+    #[allow(clippy::too_many_arguments)]
-    // &mut self,
+    pub fn splice_u16<D0, A, B, BRA, M, BE: Backend>(
-    // module: &M,
+        &mut self,
-    // byte_self: usize,
+        module: &M,
-    // byte_a: usize,
+        dst: usize,
-    // a: &FheUint<D1, T>,
+        src: usize,
-    // keys: &BDDKeyPrepared<D0, BRA, BE>,
+        a: &A,
-    // scratch: &mut Scratch<BE>,
+        b: &B,
-    // ) where
+        keys: &BDDKeyPrepared<D0, BRA, BE>,
-    // D0: DataRef,
+        scratch: &mut Scratch<BE>,
-    // D1: DataRef,
+    ) where
-    // BRA: BlindRotationAlgo,
+        D0: DataRef,
-    // M:ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd,
+        A: GLWEToRef + GLWEInfos,
-    // Scratch<BE>: ScratchTakeBDD<T, BE>,
+        B: GLWEToRef + GLWEInfos,
-    // {
+        BRA: BlindRotationAlgo,
-    // let (mut tmp_fhe_uint_byte, scratch_1) = scratch.take_fhe_uint(a);
+        M: ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWECopy,
-    //
+        Scratch<BE>: ScratchTakeBDD<T, BE>,
-    //
+    {
-    // let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
+        assert!(dst < (T::BITS >> 4) as usize);
-    // module.glwe_rotate(-((T::bit_index(byte_a << 3) << log_gap) as i64), tmp_fhe_uint_byte, self);
+        assert!(src < (T::BITS >> 4) as usize);
-    // module.glwe_trace_inplace(&mut tmp_fhe_uint_byte, module.log_n() - 3, module.log_n(),&keys.cbt.atk, scratch);
+
-    //
+        let (mut tmp, scratch_1) = scratch.take_fhe_uint(self);
-    // let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
+        tmp.splice_u8(module, dst << 1, src << 1, a, b, keys, scratch_1);
-    // let rot: i64 = (T::bit_index(byte_self << 3) << log_gap) as i64;
+        self.splice_u8(
-    //
+            module,
-    // Move starting byte index to first coefficient
+            (dst << 1) + 1,
-    // module.glwe_rotate_inplace(-rot, &mut self.bits, scratch);
+            (src << 1) + 1,
-    //
+            &tmp,
+            b,
+            keys,
+            scratch_1,
+        );
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    // Store on the receiver a where the byte_a-th byte of a has been replaced by byte_src2 of src2.
+    pub fn splice_u8<D0, A, B, BRA, M, BE: Backend>(
+        &mut self,
+        module: &M,
+        dst: usize,
+        src: usize,
+        a: &A,
+        b: &B,
+        keys: &BDDKeyPrepared<D0, BRA, BE>,
+        scratch: &mut Scratch<BE>,
+    ) where
+        D0: DataRef,
+        A: GLWEToRef + GLWEInfos,
+        B: GLWEToRef + GLWEInfos,
+        BRA: BlindRotationAlgo,
+        M: ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWECopy,
+        Scratch<BE>: ScratchTakeBDD<T, BE>,
+    {
+        assert!(dst < (T::BITS >> 3) as usize);
+        assert!(src < (T::BITS >> 3) as usize);
+
+        // 1) Zero the byte receiver
+        let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
+        let trace_start = (T::LOG_BITS - T::LOG_BYTES) as usize;
+        let rot: i64 = (T::bit_index(dst << 3) << log_gap) as i64;
+
+        // Move a to self and align byte
+        module.glwe_rotate(-rot, &mut self.bits, a);
+
         // Stores this byte (everything else zeroed) into tmp_trace
-    // let (mut tmp_trace, scratch_1) = scratch.take_glwe(self);
+        let (mut tmp_trace, scratch_1) = scratch.take_glwe(a);
-    // module.glwe_trace(
+        module.glwe_trace(
-    // &mut tmp_trace,
+            &mut tmp_trace,
-    // module.log_n() - 3,
+            trace_start,
-    // module.log_n(),
+            module.log_n(),
-    // self,
+            self,
-    // &keys.cbt.atk,
+            &keys.cbt.atk,
-    // scratch_1,
+            scratch_1,
-    // );
+        );
-    //
+
-    // Subtracts the byte
+        // Subtracts to self to zero it
-    // module.glwe_sub_inplace(&mut self.bits, &tmp_trace);
+        module.glwe_sub_inplace(&mut self.bits, &tmp_trace);
-    //
+
-    // module.glwe_add_inplace(&mut self.bits, &tmp_fhe_uint_byte);
+        // Isolate the byte to transfer from a
-    //
+        let (mut tmp_fhe_uint_byte, scratch_1) = scratch.take_fhe_uint(b);
-    // Moves back into the original position
+
-    // module.glwe_rotate_inplace(-rot, &mut self.bits, scratch);
+        // Move a[byte_a] into a[0]
-    //
+        module.glwe_rotate(
-    // }
+            -((T::bit_index(src << 3) << log_gap) as i64),
+            &mut tmp_fhe_uint_byte,
+            b,
+        );
+
+        // Zeroes all other bytes
+        module.glwe_trace_inplace(
+            &mut tmp_fhe_uint_byte,
+            trace_start,
+            module.log_n(),
+            &keys.cbt.atk,
+            scratch_1,
+        );
+
+        // Add self[0] += a[0]
+        module.glwe_add_inplace(&mut self.bits, &tmp_fhe_uint_byte);
+
+        // Moves back self[0] to self[byte_tg]
+        module.glwe_rotate_inplace(rot, &mut self.bits, scratch);
+    }
 }
 
 impl<D: DataMut, T: UnsignedInteger> GLWEToMut for FheUint<D, T> {

poulpy-schemes/src/tfhe/bdd_arithmetic/mod.rs

@@ -22,7 +22,7 @@ pub trait UnsignedInteger: Copy + 'static {
 
     #[inline(always)]
     fn bit_index(i: usize) -> usize {
-        ((i & Self::LOG_BYTES_MASK) << 3) | (i >> Self::LOG_BYTES)
+        ((i & 7) << Self::LOG_BYTES) | (i >> 3)
     }
 }
 

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/fft64_ref.rs

@@ -5,8 +5,8 @@ use poulpy_backend::FFT64Ref;
 use crate::tfhe::{
     bdd_arithmetic::tests::test_suite::{
         TestContext, test_bdd_add, test_bdd_and, test_bdd_or, test_bdd_prepare, test_bdd_sll, test_bdd_slt, test_bdd_sltu,
-        test_bdd_sra, test_bdd_srl, test_bdd_sub, test_bdd_xor, test_glwe_to_glwe_blind_rotation,
+        test_bdd_sra, test_bdd_srl, test_bdd_sub, test_bdd_xor, test_fhe_uint_splice_u8, test_fhe_uint_splice_u16,
-        test_scalar_to_ggsw_blind_rotation,
+        test_glwe_to_glwe_blind_rotation, test_scalar_to_ggsw_blind_rotation,
     },
     blind_rotation::CGGI,
 };
@@ -14,6 +14,16 @@ use crate::tfhe::{
 static TEST_CONTEXT_CGGI_FFT64_REF: LazyLock<TestContext<CGGI, FFT64Ref>> =
     LazyLock::new(|| TestContext::<CGGI, FFT64Ref>::new());
 
+#[test]
+fn test_fhe_uint_splice_u8_fft64_ref() {
+    test_fhe_uint_splice_u8(&TEST_CONTEXT_CGGI_FFT64_REF)
+}
+
+#[test]
+fn test_fhe_uint_splice_u16_fft64_ref() {
+    test_fhe_uint_splice_u16(&TEST_CONTEXT_CGGI_FFT64_REF)
+}
+
 #[test]
 fn test_glwe_to_glwe_blind_rotation_fft64_ref() {
     test_glwe_to_glwe_blind_rotation(&TEST_CONTEXT_CGGI_FFT64_REF)

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/fheuint.rs

@@ -0,0 +1,128 @@
+use poulpy_core::{
+    GLWEAdd, GLWEDecrypt, GLWEEncryptSk, GLWERotate, GLWESub, GLWETrace,
+    layouts::{GLWELayout, GLWESecretPrepared},
+};
+use poulpy_hal::{
+    api::{ScratchOwnedAlloc, ScratchOwnedBorrow},
+    layouts::{Backend, Module, Scratch, ScratchOwned},
+    source::Source,
+};
+
+use crate::tfhe::{
+    bdd_arithmetic::{
+        BDDKeyPrepared, FheUint, ScratchTakeBDD,
+        tests::test_suite::{TEST_GLWE_INFOS, TestContext},
+    },
+    blind_rotation::BlindRotationAlgo,
+};
+
+pub fn test_fhe_uint_splice_u8<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
+where
+    Module<BE>: GLWEEncryptSk<BE> + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWEDecrypt<BE>,
+    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
+    Scratch<BE>: ScratchTakeBDD<u32, BE>,
+{
+    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
+
+    let module: &Module<BE> = &test_context.module;
+    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
+    let keys: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
+
+    let mut source_xa: Source = Source::new([2u8; 32]);
+    let mut source_xe: Source = Source::new([3u8; 32]);
+
+    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
+
+    let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
+    let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
+    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
+
+    let a: u32 = 0xFFFFFFFF;
+    let b: u32 = 0xAABBCCDD;
+
+    b_enc.encrypt_sk(
+        module,
+        b,
+        sk,
+        &mut source_xa,
+        &mut source_xe,
+        scratch.borrow(),
+    );
+    a_enc.encrypt_sk(
+        module,
+        a,
+        sk,
+        &mut source_xa,
+        &mut source_xe,
+        scratch.borrow(),
+    );
+
+    for dst in 0..4 {
+        for src in 0..4 {
+            c_enc.splice_u8(module, dst, src, &a_enc, &b_enc, keys, scratch.borrow());
+
+            let rj: u32 = (dst << 3) as u32;
+            let ri: u32 = (src << 3) as u32;
+            let a_r: u32 = a.rotate_right(rj);
+            let b_r: u32 = b.rotate_right(ri);
+
+            let c_want: u32 = ((a_r & 0xFFFF_FF00) | (b_r & 0x0000_00FF)).rotate_left(rj);
+
+            assert_eq!(c_want, c_enc.decrypt(module, sk, scratch.borrow()));
+        }
+    }
+}
+
+pub fn test_fhe_uint_splice_u16<BRA: BlindRotationAlgo, BE: Backend>(test_context: &TestContext<BRA, BE>)
+where
+    Module<BE>: GLWEEncryptSk<BE> + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWEDecrypt<BE>,
+    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
+    Scratch<BE>: ScratchTakeBDD<u32, BE>,
+{
+    let glwe_infos: GLWELayout = TEST_GLWE_INFOS;
+
+    let module: &Module<BE> = &test_context.module;
+    let sk: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
+    let keys: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
+
+    let mut source_xa: Source = Source::new([2u8; 32]);
+    let mut source_xe: Source = Source::new([3u8; 32]);
+
+    let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 22);
+
+    let mut a_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
+    let mut b_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
+    let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::<Vec<u8>, u32>::alloc_from_infos(&glwe_infos);
+
+    let a: u32 = 0xFFFFFFFF;
+    let b: u32 = 0xAABBCCDD;
+
+    b_enc.encrypt_sk(
+        module,
+        b,
+        sk,
+        &mut source_xa,
+        &mut source_xe,
+        scratch.borrow(),
+    );
+    a_enc.encrypt_sk(
+        module,
+        a,
+        sk,
+        &mut source_xa,
+        &mut source_xe,
+        scratch.borrow(),
+    );
+
+    for dst in 0..2 {
+        for src in 0..2 {
+            c_enc.splice_u16(module, dst, src, &a_enc, &b_enc, keys, scratch.borrow());
+            let rj: u32 = (dst << 4) as u32;
+            let ri: u32 = (src << 4) as u32;
+            let a_r: u32 = a.rotate_right(rj);
+            let b_r: u32 = b.rotate_right(ri);
+            let c_want: u32 = ((a_r & 0xFFFF_0000) | (b_r & 0x0000_FFFF)).rotate_left(rj);
+            assert_eq!(c_want, c_enc.decrypt(module, sk, scratch.borrow()));
+        }
+    }
+}

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/mod.rs

@@ -1,5 +1,6 @@
 mod add;
 mod and;
+mod fheuint;
 mod ggsw_blind_rotations;
 mod glwe_blind_rotation;
 mod or;
@@ -14,6 +15,7 @@ mod xor;
 
 pub use add::*;
 pub use and::*;
+pub use fheuint::*;
 pub use ggsw_blind_rotations::*;
 pub use glwe_blind_rotation::*;
 pub use or::*;