Add FheUint from Vec<GLWE>

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

@@ -54,7 +54,7 @@ where
         self.execute_bdd_circuit(&mut out_bits, &helper, circuit, scratch_1);
 
         // Repacks the bits
-        out.pack(self, out_bits, &key.cbt.atk, scratch_1);
+        out.pack(self, out_bits, key, scratch_1);
     }
 }
 

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

@@ -1,8 +1,8 @@
 use poulpy_core::{
     GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEPacking, GLWERotate, LWEFromGLWE, ScratchTakeCore,
     layouts::{
-        Base2K, Degree, GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWEPlaintextLayout, GLWESecretPreparedToRef, GLWEToRef,
-        LWEInfos, LWEToMut, Rank, TorusPrecision, prepared::GLWEAutomorphismKeyPrepared,
+        Base2K, Degree, GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWEPlaintextLayout, GLWESecretPreparedToRef, GLWEToMut,
+        GLWEToRef, LWEInfos, LWEToMut, Rank, TorusPrecision,
     },
 };
 use poulpy_hal::{
@@ -12,7 +12,10 @@ use poulpy_hal::{
 };
 use std::{collections::HashMap, marker::PhantomData};
 
-use crate::tfhe::bdd_arithmetic::{FromBits, ToBits, UnsignedInteger};
+use crate::tfhe::{
+    bdd_arithmetic::{BDDKeyPrepared, FromBits, ToBits, UnsignedInteger},
+    blind_rotation::BlindRotationAlgo,
+};
 
 /// An FHE ciphertext encrypting the bits of an [UnsignedInteger].
 pub struct FheUint<D: Data, T: UnsignedInteger> {
@@ -141,37 +144,113 @@ impl<D: DataRef, T: UnsignedInteger + FromBits> FheUint<D, T> {
 }
 
 impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
-    #[allow(dead_code)]
-    pub(crate) fn pack<D1, ATK, M, BE: Backend>(
+    /// Packs Vec<GLWE(bit[i])> into [FheUint].
+    pub fn pack<G, D1, M, BRA: BlindRotationAlgo, BE: Backend>(
         &mut self,
         module: &M,
-        mut tmp_res: Vec<GLWE<D1>>,
-        auto_keys: &HashMap<i64, GLWEAutomorphismKeyPrepared<ATK, BE>>,
+        mut bits: Vec<G>,
+        key: &BDDKeyPrepared<D1, BRA, BE>,
         scratch: &mut Scratch<BE>,
     ) where
-        D1: DataMut,
-        ATK: DataRef,
+        G: GLWEToMut + GLWEToRef + GLWEInfos,
+        D1: DataRef,
         M: ModuleLogN + GLWEPacking<BE> + GLWECopy,
         Scratch<BE>: ScratchTakeCore<BE>,
     {
         // Repacks the GLWE ciphertexts bits
         let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
-        let mut cts: HashMap<usize, &mut GLWE<D1>> = HashMap::new();
-        for (i, ct) in tmp_res.iter_mut().enumerate().take(T::BITS as usize) {
+        let mut cts: HashMap<usize, &mut G> = HashMap::new();
+        for (i, ct) in bits.iter_mut().enumerate().take(T::BITS as usize) {
             cts.insert(T::bit_index(i) << log_gap, ct);
         }
 
-        module.glwe_pack(&mut cts, log_gap, auto_keys, scratch);
+        module.glwe_pack(&mut cts, log_gap, &key.cbt.atk, scratch);
 
         // And copies the repacked ciphertext on the receiver.
         module.glwe_copy(&mut self.bits, cts.remove(&0).unwrap());
     }
+
+    // pub fn copy_byte<D0, D1, BRA, M, BE: Backend>(
+    // &mut self,
+    // module: &M,
+    // byte_self: usize,
+    // byte_a: usize,
+    // a: &FheUint<D1, T>,
+    // keys: &BDDKeyPrepared<D0, BRA, BE>,
+    // scratch: &mut Scratch<BE>,
+    // ) where
+    // D0: DataRef,
+    // D1: DataRef,
+    // BRA: BlindRotationAlgo,
+    // M:ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd,
+    // Scratch<BE>: ScratchTakeBDD<T, BE>,
+    // {
+    // let (mut tmp_fhe_uint_byte, scratch_1) = scratch.take_fhe_uint(a);
+    //
+    //
+    // let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
+    // module.glwe_rotate(-((T::bit_index(byte_a << 3) << log_gap) as i64), tmp_fhe_uint_byte, self);
+    // module.glwe_trace_inplace(&mut tmp_fhe_uint_byte, module.log_n() - 3, module.log_n(),&keys.cbt.atk, scratch);
+    //
+    // let log_gap: usize = module.log_n() - T::LOG_BITS as usize;
+    // let rot: i64 = (T::bit_index(byte_self << 3) << log_gap) as i64;
+    //
+    // Move starting byte index to first coefficient
+    // module.glwe_rotate_inplace(-rot, &mut self.bits, scratch);
+    //
+    // Stores this byte (everything else zeroed) into tmp_trace
+    // let (mut tmp_trace, scratch_1) = scratch.take_glwe(self);
+    // module.glwe_trace(
+    // &mut tmp_trace,
+    // module.log_n() - 3,
+    // module.log_n(),
+    // self,
+    // &keys.cbt.atk,
+    // scratch_1,
+    // );
+    //
+    // Subtracts the byte
+    // module.glwe_sub_inplace(&mut self.bits, &tmp_trace);
+    //
+    // module.glwe_add_inplace(&mut self.bits, &tmp_fhe_uint_byte);
+    //
+    // Moves back into the original position
+    // module.glwe_rotate_inplace(-rot, &mut self.bits, scratch);
+    //
+    // }
 }
 
-impl<D: DataRef, T: UnsignedInteger> FheUint<D, T> {
-    pub fn get_bit<L, K, M, BE: Backend>(&self, module: &M, bit: usize, res: &mut L, ks: &K, scratch: &mut Scratch<BE>)
+impl<D: DataMut, T: UnsignedInteger> GLWEToMut for FheUint<D, T> {
+    fn to_mut(&mut self) -> GLWE<&mut [u8]> {
+        self.bits.to_mut()
+    }
+}
+
+pub trait ScratchTakeBDD<T: UnsignedInteger, BE: Backend>
+where
+    Self: ScratchTakeCore<BE>,
+{
+    fn take_fhe_uint<A>(&mut self, infos: &A) -> (FheUint<&mut [u8], T>, &mut Self)
     where
-        L: LWEToMut,
+        A: GLWEInfos,
+    {
+        let (glwe, scratch) = self.take_glwe(infos);
+        (
+            FheUint {
+                bits: glwe,
+                _phantom: PhantomData,
+            },
+            scratch,
+        )
+    }
+}
+
+impl<T: UnsignedInteger, BE: Backend> ScratchTakeBDD<T, BE> for Scratch<BE> where Self: ScratchTakeCore<BE> {}
+
+impl<D: DataRef, T: UnsignedInteger> FheUint<D, T> {
+    pub fn get_bit<R, K, M, BE: Backend>(&self, module: &M, bit: usize, res: &mut R, ks: &K, scratch: &mut Scratch<BE>)
+    where
+        R: LWEToMut,
         K: GGLWEPreparedToRef<BE> + GGLWEInfos,
         M: ModuleLogN + LWEFromGLWE<BE> + GLWERotate<BE>,
         Scratch<BE>: ScratchTakeCore<BE>,