M poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint_prepared.rs → poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint_prepared.rs

@@ -217,6 +217,18 @@ DM: DataMut,
         DB: DataRef,
         DK: DataRef,
         K: BDDKeyHelper<DK, BRA, BE>;
+    fn fhe_uint_prepare_partial<DM, DB, DK, K>(
+        &self,
+        res: &mut FheUintPrepared<DM, T, BE>,
+        bits: &FheUint<DB, T>,
+        count: usize,
+        key: &K,
+        scratch: &mut Scratch<BE>,
+    ) where
+        DM: DataMut,
+        DB: DataRef,
+        DK: DataRef,
+        K: BDDKeyHelper<DK, BRA, BE>;        
 }
 
 impl<BRA: BlindRotationAlgo, BE: Backend, T: UnsignedInteger> FheUintPrepare<BRA, T, BE> for Module<BE>
@@ -259,6 +271,30 @@ cbt.execute_to_constant(self, &mut tmp_ggsw, &lwe, 1, 1, scratch_1);
             dst.prepare(self, &tmp_ggsw, scratch_1);
         }
     }
+
+    fn fhe_uint_prepare_partial<DM, DB, DK, K>(
+        &self,
+        res: &mut FheUintPrepared<DM, T, BE>,
+        bits: &FheUint<DB, T>,
+        count: usize,
+        key: &K,
+        scratch: &mut Scratch<BE>,
+    ) where
+        DM: DataMut,
+        DB: DataRef,
+        DK: DataRef,
+        K: BDDKeyHelper<DK, BRA, BE>,
+    {
+        let (cbt, ks) = key.get_cbt_key();
+
+        let mut lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(bits); //TODO: add TakeLWE
+        let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(res);
+        for (bit, dst) in res.bits[0..count].iter_mut().enumerate() { // TODO: set the rest of the bits to a prepared zero GGSW
+            bits.get_bit_lwe(self, bit, &mut lwe, ks, scratch_1);
+            cbt.execute_to_constant(self, &mut tmp_ggsw, &lwe, 1, 1, scratch_1);
+            dst.prepare(self, &tmp_ggsw, scratch_1);
+        }
+    }    
 }
 
 impl<D: DataMut, T: UnsignedInteger, BE: Backend> FheUintPrepared<D, T, BE> {
@@ -273,4 +309,15 @@ Scratch<BE>: ScratchTakeCore<BE>,
     {
         module.fhe_uint_prepare(self, other, key, scratch);
     }
+    pub fn prepare_partial<BRA, M, O, K, DK>(&mut self, module: &M, other: &FheUint<O, T>, count: usize, key: &K, scratch: &mut Scratch<BE>)
+    where
+        BRA: BlindRotationAlgo,
+        O: DataRef,
+        DK: DataRef,
+        K: BDDKeyHelper<DK, BRA, BE>,
+        M: FheUintPrepare<BRA, T, BE>,
+        Scratch<BE>: ScratchTakeCore<BE>,
+    {
+        module.fhe_uint_prepare_partial(self, other, count, key, scratch);
+    }    
 }