wip

poulpy-core/src/encryption/gglwe_ct.rs

@@ -1,16 +1,19 @@
 use poulpy_hal::{
     api::{
-        ScratchAvailable, SvpApplyDftToDftInplace, TakeVecZnx, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal,
-        VecZnxAddScalarInplace, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxFillUniform,
-        VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace,
+        ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxAddScalarInplace, VecZnxDftAllocBytes, VecZnxNormalizeInplace,
+        VecZnxNormalizeTmpBytes,
     },
-    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch, ZnxZero},
+    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, ScalarZnxToRef, Scratch, ZnxZero},
     source::Source,
 };
 
 use crate::{
     TakeGLWEPt,
-    layouts::{GGLWECiphertext, GGLWEInfos, GLWECiphertext, GLWEPlaintext, LWEInfos, prepared::GLWESecretPrepared},
+    encryption::glwe_ct::GLWEEncryptSk,
+    layouts::{
+        GGLWECiphertext, GGLWECiphertextToMut, GGLWEInfos, GLWECiphertext, GLWEPlaintext, LWEInfos,
+        prepared::{GLWESecretPrepared, GLWESecretPreparedToRef},
+    },
 };
 
 impl GGLWECiphertext<Vec<u8>> {
@@ -31,78 +34,89 @@ impl GGLWECiphertext<Vec<u8>> {
     }
 }
 
-impl<DataSelf: DataMut> GGLWECiphertext<DataSelf> {
-    #[allow(clippy::too_many_arguments)]
-    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
-        &mut self,
-        module: &Module<B>,
-        pt: &ScalarZnx<DataPt>,
-        sk: &GLWESecretPrepared<DataSk, B>,
+pub trait GGLWEEncryptSk<B: Backend> {
+    fn gglwe_encrypt_sk<R, P, S>(
+        &self,
+        res: &mut R,
+        pt: &P,
+        sk: &S,
         source_xa: &mut Source,
         source_xe: &mut Source,
         scratch: &mut Scratch<B>,
     ) where
-        Module<B>: VecZnxAddScalarInplace
-            + VecZnxDftAllocBytes
-            + VecZnxBigNormalize<B>
-            + VecZnxDftApply<B>
-            + SvpApplyDftToDftInplace<B>
-            + VecZnxIdftApplyConsume<B>
-            + VecZnxNormalizeTmpBytes
-            + VecZnxFillUniform
-            + VecZnxSubInplace
-            + VecZnxAddInplace
-            + VecZnxNormalizeInplace<B>
-            + VecZnxAddNormal
-            + VecZnxNormalize<B>
-            + VecZnxSub,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
+        R: GGLWECiphertextToMut,
+        P: ScalarZnxToRef,
+        S: GLWESecretPreparedToRef<B>;
+}
+
+impl<B: Backend> GGLWEEncryptSk<B> for Module<B>
+where
+    Module<B>:
+        GLWEEncryptSk<B> + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxAddScalarInplace + VecZnxNormalizeInplace<B>,
+    Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
+{
+    fn gglwe_encrypt_sk<R, P, S>(
+        &self,
+        res: &mut R,
+        pt: &P,
+        sk: &S,
+        source_xa: &mut Source,
+        source_xe: &mut Source,
+        scratch: &mut Scratch<B>,
+    ) where
+        R: GGLWECiphertextToMut,
+        P: ScalarZnxToRef,
+        S: GLWESecretPreparedToRef<B>,
     {
+        let res: &mut GGLWECiphertext<&mut [u8]> = &mut res.to_mut();
+        let pt: &ScalarZnx<&[u8]> = &pt.to_ref();
+
         #[cfg(debug_assertions)]
         {
             use poulpy_hal::layouts::ZnxInfos;
+            let sk: GLWESecretPrepared<&[u8], B> = sk.to_ref();
 
             assert_eq!(
-                self.rank_in(),
+                res.rank_in(),
                 pt.cols() as u32,
-                "self.rank_in(): {} != pt.cols(): {}",
-                self.rank_in(),
+                "res.rank_in(): {} != pt.cols(): {}",
+                res.rank_in(),
                 pt.cols()
             );
             assert_eq!(
-                self.rank_out(),
+                res.rank_out(),
                 sk.rank(),
-                "self.rank_out(): {} != sk.rank(): {}",
-                self.rank_out(),
+                "res.rank_out(): {} != sk.rank(): {}",
+                res.rank_out(),
                 sk.rank()
             );
-            assert_eq!(self.n(), sk.n());
+            assert_eq!(res.n(), sk.n());
             assert_eq!(pt.n() as u32, sk.n());
             assert!(
-                scratch.available() >= GGLWECiphertext::encrypt_sk_scratch_space(module, self),
-                "scratch.available: {} < GGLWECiphertext::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
+                scratch.available() >= GGLWECiphertext::encrypt_sk_scratch_space(self, res),
+                "scratch.available: {} < GGLWECiphertext::encrypt_sk_scratch_space(self, res.rank()={}, res.size()={}): {}",
                 scratch.available(),
-                self.rank_out(),
-                self.size(),
-                GGLWECiphertext::encrypt_sk_scratch_space(module, self)
+                res.rank_out(),
+                res.size(),
+                GGLWECiphertext::encrypt_sk_scratch_space(self, res)
             );
             assert!(
-                self.dnum().0 * self.dsize().0 * self.base2k().0 <= self.k().0,
-                "self.dnum() : {} * self.dsize() : {} * self.base2k() : {} = {} >= self.k() = {}",
-                self.dnum(),
-                self.dsize(),
-                self.base2k(),
-                self.dnum().0 * self.dsize().0 * self.base2k().0,
-                self.k()
+                res.dnum().0 * res.dsize().0 * res.base2k().0 <= res.k().0,
+                "res.dnum() : {} * res.dsize() : {} * res.base2k() : {} = {} >= res.k() = {}",
+                res.dnum(),
+                res.dsize(),
+                res.base2k(),
+                res.dnum().0 * res.dsize().0 * res.base2k().0,
+                res.k()
             );
         }
 
-        let dnum: usize = self.dnum().into();
-        let dsize: usize = self.dsize().into();
-        let base2k: usize = self.base2k().into();
-        let rank_in: usize = self.rank_in().into();
+        let dnum: usize = res.dnum().into();
+        let dsize: usize = res.dsize().into();
+        let base2k: usize = res.base2k().into();
+        let rank_in: usize = res.rank_in().into();
 
-        let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(self);
+        let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(res);
         // For each input column (i.e. rank) produces a GGLWE ciphertext of rank_out+1 columns
         //
         // Example for ksk rank 2 to rank 3:
@@ -114,17 +128,39 @@ impl<DataSelf: DataMut> GGLWECiphertext<DataSelf> {
         //
         // (-(a*s) + s0, a)
         // (-(b*s) + s1, b)
-        (0..rank_in).for_each(|col_i| {
-            (0..dnum).for_each(|row_i| {
+
+        for col_i in 0..rank_in {
+            for row_i in 0..dnum {
                 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
                 tmp_pt.data.zero(); // zeroes for next iteration
-                module.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (dsize - 1) + row_i * dsize, pt, col_i);
-                module.vec_znx_normalize_inplace(base2k, &mut tmp_pt.data, 0, scrach_1);
-
-                // rlwe encrypt of vec_znx_pt into vec_znx_ct
-                self.at_mut(row_i, col_i)
-                    .encrypt_sk(module, &tmp_pt, sk, source_xa, source_xe, scrach_1);
-            });
-        });
+                self.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (dsize - 1) + row_i * dsize, pt, col_i);
+                self.vec_znx_normalize_inplace(base2k, &mut tmp_pt.data, 0, scrach_1);
+                self.glwe_encrypt_sk(
+                    &mut res.at_mut(row_i, col_i),
+                    &tmp_pt,
+                    sk,
+                    source_xa,
+                    source_xe,
+                    scrach_1,
+                );
+            }
+        }
+    }
+}
+
+impl<DataSelf: DataMut> GGLWECiphertext<DataSelf> {
+    #[allow(clippy::too_many_arguments)]
+    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
+        &mut self,
+        module: &Module<B>,
+        pt: &ScalarZnx<DataPt>,
+        sk: &GLWESecretPrepared<DataSk, B>,
+        source_xa: &mut Source,
+        source_xe: &mut Source,
+        scratch: &mut Scratch<B>,
+    ) where
+        Module<B>: GGLWEEncryptSk<B>,
+    {
+        module.gglwe_encrypt_sk(self, pt, sk, source_xa, source_xe, scratch);
     }
 }