fixed encryption

2026-02-10 13:16:44 +01:00 · 2025-10-19 18:11:04 +02:00
parent d6e9805a8f
commit a706b759d6
17 changed files with 289 additions and 196 deletions
--- a/poulpy-core/src/encryption/compressed/gglwe_ct.rs
+++ b/poulpy-core/src/encryption/compressed/gglwe_ct.rs
@@ -11,7 +11,7 @@ use crate::{
        glwe_ct::{GLWEEncryptSk, GLWEEncryptSkInternal},
    },
    layouts::{
-        GGLWEInfos, GLWEPlaintextAlloc, LWEInfos,
+        GGLWECompressedSeedMut, GGLWEInfos, GLWEPlaintextAlloc, GLWESecretPrepared, LWEInfos,
        compressed::{GGLWECompressed, GGLWECompressedToMut},
        prepared::GLWESecretPreparedToRef,
    },
@@ -60,7 +60,7 @@ pub trait GGLWECompressedEncryptSk<BE: Backend> {
        source_xe: &mut Source,
        scratch: &mut Scratch<BE>,
    ) where
-        R: GGLWECompressedToMut,
+        R: GGLWECompressedToMut + GGLWECompressedSeedMut,
        P: ScalarZnxToRef,
        S: GLWESecretPreparedToRef<BE>;
 }
@@ -94,80 +94,85 @@ where
        source_xe: &mut Source,
        scratch: &mut Scratch<BE>,
    ) where
-        R: GGLWECompressedToMut,
+        R: GGLWECompressedToMut + GGLWECompressedSeedMut,
        P: ScalarZnxToRef,
        S: GLWESecretPreparedToRef<BE>,
    {
-        let res: &mut GGLWECompressed<&mut [u8]> = &mut res.to_mut();
-        let pt: &ScalarZnx<&[u8]> = &pt.to_ref();
+        let mut seeds: Vec<[u8; 32]> = vec![[0u8; 32]; res.seed_mut().len()];

-        let sk = &sk.to_ref();
+        {
+            let res: &mut GGLWECompressed<&mut [u8]> = &mut res.to_mut();
+            let pt: &ScalarZnx<&[u8]> = &pt.to_ref();
+            let sk: &GLWESecretPrepared<&[u8], BE> = &sk.to_ref();

-        assert_eq!(
-            res.rank_in(),
-            pt.cols() as u32,
-            "res.rank_in(): {} != pt.cols(): {}",
-            res.rank_in(),
-            pt.cols()
-        );
-        assert_eq!(
-            res.rank_out(),
-            sk.rank(),
-            "res.rank_out(): {} != sk.rank(): {}",
-            res.rank_out(),
-            sk.rank()
-        );
-        assert_eq!(res.n(), sk.n());
-        assert_eq!(pt.n() as u32, sk.n());
-        assert!(
-            scratch.available() >= GGLWECompressed::encrypt_sk_tmp_bytes(self, res),
-            "scratch.available: {} < GGLWECiphertext::encrypt_sk_tmp_bytes: {}",
-            scratch.available(),
-            GGLWECompressed::encrypt_sk_tmp_bytes(self, res)
-        );
-        assert!(
-            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()
-        );
+            assert_eq!(
+                res.rank_in(),
+                pt.cols() as u32,
+                "res.rank_in(): {} != pt.cols(): {}",
+                res.rank_in(),
+                pt.cols()
+            );
+            assert_eq!(
+                res.rank_out(),
+                sk.rank(),
+                "res.rank_out(): {} != sk.rank(): {}",
+                res.rank_out(),
+                sk.rank()
+            );
+            assert_eq!(res.n(), sk.n());
+            assert_eq!(pt.n() as u32, sk.n());
+            assert!(
+                scratch.available() >= GGLWECompressed::encrypt_sk_tmp_bytes(self, res),
+                "scratch.available: {} < GGLWECiphertext::encrypt_sk_tmp_bytes: {}",
+                scratch.available(),
+                GGLWECompressed::encrypt_sk_tmp_bytes(self, res)
+            );
+            assert!(
+                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 = res.dnum().into();
-        let dsize: usize = res.dsize().into();
-        let base2k: usize = res.base2k().into();
-        let rank_in: usize = res.rank_in().into();
-        let cols: usize = (res.rank_out() + 1).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 cols: usize = (res.rank_out() + 1).into();

-        let mut source_xa = Source::new(seed);
+            let mut source_xa = Source::new(seed);

-        let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(self, res);
-        (0..rank_in).for_each(|col_i| {
-            (0..dnum).for_each(|d_i| {
-                // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
-                tmp_pt.data.zero(); // zeroes for next iteration
-                self.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (dsize - 1) + d_i * dsize, pt, col_i);
-                self.vec_znx_normalize_inplace(base2k, &mut tmp_pt.data, 0, scrach_1);
+            let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(self, res);
+            for col_i in 0..rank_in {
+                for d_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
+                    self.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (dsize - 1) + d_i * dsize, pt, col_i);
+                    self.vec_znx_normalize_inplace(base2k, &mut tmp_pt.data, 0, scrach_1);

-                let (seed, mut source_xa_tmp) = source_xa.branch();
-                res.seed[col_i * dnum + d_i] = seed;
+                    let (seed, mut source_xa_tmp) = source_xa.branch();
+                    seeds[col_i * dnum + d_i] = seed;

-                self.glwe_encrypt_sk_internal(
-                    res.base2k().into(),
-                    res.k().into(),
-                    &mut res.at_mut(d_i, col_i).data,
-                    cols,
-                    true,
-                    Some((&tmp_pt, 0)),
-                    sk,
-                    &mut source_xa_tmp,
-                    source_xe,
-                    SIGMA,
-                    scrach_1,
-                );
-            });
-        });
+                    self.glwe_encrypt_sk_internal(
+                        res.base2k().into(),
+                        res.k().into(),
+                        &mut res.at_mut(d_i, col_i).data,
+                        cols,
+                        true,
+                        Some((&tmp_pt, 0)),
+                        sk,
+                        &mut source_xa_tmp,
+                        source_xe,
+                        SIGMA,
+                        scrach_1,
+                    );
+                }
+            }
+        }
+
+        res.seed_mut().copy_from_slice(&seeds);
    }
 }
--- a/poulpy-core/src/encryption/compressed/gglwe_tsk.rs
+++ b/poulpy-core/src/encryption/compressed/gglwe_tsk.rs
@@ -9,11 +9,10 @@ use poulpy_hal::{
 };

 use crate::{
-    ScratchTakeCore,
-    encryption::compressed::gglwe_ksk::GLWESwitchingKeyCompressedEncryptSk,
-    encryption::gglwe_tsk::TensorKeyEncryptSk,
+    GetDistribution, ScratchTakeCore,
+    encryption::{compressed::gglwe_ksk::GLWESwitchingKeyCompressedEncryptSk, gglwe_tsk::TensorKeyEncryptSk},
    layouts::{
-        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, GetDist, LWEInfos, Rank,
+        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos, Rank,
        compressed::{TensorKeyCompressed, TensorKeyCompressedToMut},
    },
 };
@@ -37,7 +36,7 @@ impl<DataSelf: DataMut> TensorKeyCompressed<DataSelf> {
        source_xe: &mut Source,
        scratch: &mut Scratch<BE>,
    ) where
-        S: GLWESecretToRef + GetDist,
+        S: GLWESecretToRef + GetDistribution,
        M: GGLWETensorKeyCompressedEncryptSk<BE>,
    {
        module.gglwe_tensor_key_encrypt_sk(self, sk, seed_xa, source_xe, scratch);
@@ -58,7 +57,7 @@ pub trait GGLWETensorKeyCompressedEncryptSk<BE: Backend> {
        scratch: &mut Scratch<BE>,
    ) where
        R: TensorKeyCompressedToMut,
-        S: GLWESecretToRef + GetDist;
+        S: GLWESecretToRef + GetDistribution;
 }

 impl<BE: Backend> GGLWETensorKeyCompressedEncryptSk<BE> for Module<BE>
@@ -95,7 +94,7 @@ where
        scratch: &mut Scratch<BE>,
    ) where
        R: TensorKeyCompressedToMut,
-        S: GLWESecretToRef + GetDist,
+        S: GLWESecretToRef + GetDistribution,
    {
        let res: &mut TensorKeyCompressed<&mut [u8]> = &mut res.to_mut();

--- a/poulpy-core/src/encryption/compressed/ggsw_ct.rs
+++ b/poulpy-core/src/encryption/compressed/ggsw_ct.rs
@@ -8,7 +8,7 @@ use crate::{
    ScratchTakeCore,
    encryption::{SIGMA, ggsw_ct::GGSWEncryptSk, glwe_ct::GLWEEncryptSkInternal},
    layouts::{
-        GGSWInfos, GLWEInfos, LWEInfos,
+        GGSWCompressedSeedMut, GGSWInfos, GLWEInfos, LWEInfos,
        compressed::{GGSWCompressed, GGSWCompressedToMut},
        prepared::{GLWESecretPrepared, GLWESecretPreparedToRef},
    },
@@ -57,7 +57,7 @@ pub trait GGSWCompressedEncryptSk<BE: Backend> {
        source_xe: &mut Source,
        scratch: &mut Scratch<BE>,
    ) where
-        R: GGSWCompressedToMut,
+        R: GGSWCompressedToMut + GGSWCompressedSeedMut,
        P: ScalarZnxToRef,
        S: GLWESecretPreparedToRef<BE>;
 }
@@ -83,62 +83,66 @@ where
        source_xe: &mut Source,
        scratch: &mut Scratch<BE>,
    ) where
-        R: GGSWCompressedToMut,
+        R: GGSWCompressedToMut + GGSWCompressedSeedMut,
        P: ScalarZnxToRef,
        S: GLWESecretPreparedToRef<BE>,
    {
-        let res: &mut GGSWCompressed<&mut [u8]> = &mut res.to_mut();
-        let sk: &GLWESecretPrepared<&[u8], BE> = &sk.to_ref();
-        let pt: &ScalarZnx<&[u8]> = &pt.to_ref();
+        let mut seeds: Vec<[u8; 32]> = vec![[0u8; 32]; res.seed_mut().len()];

-        #[cfg(debug_assertions)]
        {
-            use poulpy_hal::layouts::ZnxInfos;
+            let res: &mut GGSWCompressed<&mut [u8]> = &mut res.to_mut();
+            let sk: &GLWESecretPrepared<&[u8], BE> = &sk.to_ref();
+            let pt: &ScalarZnx<&[u8]> = &pt.to_ref();

-            assert_eq!(res.rank(), sk.rank());
-            assert_eq!(res.n(), sk.n());
-            assert_eq!(pt.n() as u32, sk.n());
-        }
+            #[cfg(debug_assertions)]
+            {
+                use poulpy_hal::layouts::ZnxInfos;

-        let base2k: usize = res.base2k().into();
-        let rank: usize = res.rank().into();
-        let cols: usize = rank + 1;
-        let dsize: usize = res.dsize().into();
+                assert_eq!(res.rank(), sk.rank());
+                assert_eq!(res.n(), sk.n());
+                assert_eq!(pt.n() as u32, sk.n());
+            }

-        let (mut tmp_pt, scratch_1) = scratch.take_glwe_pt(self, &res.glwe_layout());
+            let base2k: usize = res.base2k().into();
+            let rank: usize = res.rank().into();
+            let cols: usize = rank + 1;
+            let dsize: usize = res.dsize().into();

-        let mut source = Source::new(seed_xa);
+            let (mut tmp_pt, scratch_1) = scratch.take_glwe_pt(self, &res.glwe_layout());

-        res.seed = vec![[0u8; 32]; res.dnum().0 as usize * cols];
+            let mut source = Source::new(seed_xa);

-        for row_i in 0..res.dnum().into() {
-            tmp_pt.data.zero();
+            for row_i in 0..res.dnum().into() {
+                tmp_pt.data.zero();

-            // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
-            self.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (dsize - 1) + row_i * dsize, pt, 0);
-            self.vec_znx_normalize_inplace(base2k, &mut tmp_pt.data, 0, scratch_1);
+                // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
+                self.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (dsize - 1) + row_i * dsize, pt, 0);
+                self.vec_znx_normalize_inplace(base2k, &mut tmp_pt.data, 0, scratch_1);

-            for col_j in 0..rank + 1 {
-                // rlwe encrypt of vec_znx_pt into vec_znx_ct
+                for col_j in 0..rank + 1 {
+                    // rlwe encrypt of vec_znx_pt into vec_znx_ct

-                let (seed, mut source_xa_tmp) = source.branch();
+                    let (seed, mut source_xa_tmp) = source.branch();

-                res.seed[row_i * cols + col_j] = seed;
+                    seeds[row_i * cols + col_j] = seed;

-                self.glwe_encrypt_sk_internal(
-                    res.base2k().into(),
-                    res.k().into(),
-                    &mut res.at_mut(row_i, col_j).data,
-                    cols,
-                    true,
-                    Some((&tmp_pt, col_j)),
-                    sk,
-                    &mut source_xa_tmp,
-                    source_xe,
-                    SIGMA,
-                    scratch_1,
-                );
+                    self.glwe_encrypt_sk_internal(
+                        res.base2k().into(),
+                        res.k().into(),
+                        &mut res.at_mut(row_i, col_j).data,
+                        cols,
+                        true,
+                        Some((&tmp_pt, col_j)),
+                        sk,
+                        &mut source_xa_tmp,
+                        source_xe,
+                        SIGMA,
+                        scratch_1,
+                    );
+                }
            }
        }
+
+        res.seed_mut().copy_from_slice(&seeds);
    }
 }
--- a/poulpy-core/src/encryption/compressed/glwe_ct.rs
+++ b/poulpy-core/src/encryption/compressed/glwe_ct.rs
@@ -9,7 +9,7 @@ use crate::{
        glwe_ct::{GLWEEncryptSk, GLWEEncryptSkInternal},
    },
    layouts::{
-        GLWEInfos, GLWEPlaintextToRef, LWEInfos,
+        GLWECompressedSeedMut, GLWEInfos, GLWEPlaintextToRef, LWEInfos,
        compressed::{GLWECompressed, GLWECompressedToMut},
        prepared::GLWESecretPreparedToRef,
    },
@@ -58,7 +58,7 @@ pub trait GLWECompressedEncryptSk<BE: Backend> {
        source_xe: &mut Source,
        scratch: &mut Scratch<BE>,
    ) where
-        R: GLWECompressedToMut,
+        R: GLWECompressedToMut + GLWECompressedSeedMut,
        P: GLWEPlaintextToRef,
        S: GLWESecretPreparedToRef<BE>;
 }
@@ -83,28 +83,30 @@ where
        source_xe: &mut Source,
        scratch: &mut Scratch<BE>,
    ) where
-        R: GLWECompressedToMut,
+        R: GLWECompressedToMut + GLWECompressedSeedMut,
        P: GLWEPlaintextToRef,
        S: GLWESecretPreparedToRef<BE>,
    {
-        let res: &mut GLWECompressed<&mut [u8]> = &mut res.to_mut();
-        let mut source_xa: Source = Source::new(seed_xa);
-        let cols: usize = (res.rank() + 1).into();
+        {
+            let res: &mut GLWECompressed<&mut [u8]> = &mut res.to_mut();
+            let mut source_xa: Source = Source::new(seed_xa);
+            let cols: usize = (res.rank() + 1).into();

-        self.glwe_encrypt_sk_internal(
-            res.base2k().into(),
-            res.k().into(),
-            &mut res.data,
-            cols,
-            true,
-            Some((pt, 0)),
-            sk,
-            &mut source_xa,
-            source_xe,
-            SIGMA,
-            scratch,
-        );
+            self.glwe_encrypt_sk_internal(
+                res.base2k().into(),
+                res.k().into(),
+                &mut res.data,
+                cols,
+                true,
+                Some((pt, 0)),
+                sk,
+                &mut source_xa,
+                source_xe,
+                SIGMA,
+                scratch,
+            );
+        }

-        res.seed = seed_xa;
+        res.seed_mut().copy_from_slice(&seed_xa);
    }
 }
--- a/poulpy-core/src/encryption/gglwe_tsk.rs
+++ b/poulpy-core/src/encryption/gglwe_tsk.rs
@@ -8,10 +8,10 @@ use poulpy_hal::{
 };

 use crate::{
-    ScratchTakeCore,
+    GetDistribution, ScratchTakeCore,
    encryption::gglwe_ksk::GLWESwitchingKeyEncryptSk,
    layouts::{
-        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, GLWESwitchingKey, GetDist, LWEInfos, Rank, TensorKey, TensorKeyToMut,
+        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, GLWESwitchingKey, LWEInfos, Rank, TensorKey, TensorKeyToMut,
        prepared::{GLWESecretPrepare, GLWESecretPrepared, GLWESecretPreparedAlloc},
    },
 };
@@ -36,7 +36,7 @@ impl<DataSelf: DataMut> TensorKey<DataSelf> {
        scratch: &mut Scratch<BE>,
    ) where
        M: TensorKeyEncryptSk<BE>,
-        S: GLWESecretToRef + GetDist,
+        S: GLWESecretToRef + GetDistribution,
        Scratch<BE>: ScratchTakeCore<BE>,
    {
        module.tensor_key_encrypt_sk(self, sk, source_xa, source_xe, scratch);
@@ -57,7 +57,7 @@ pub trait TensorKeyEncryptSk<BE: Backend> {
        scratch: &mut Scratch<BE>,
    ) where
        R: TensorKeyToMut,
-        S: GLWESecretToRef + GetDist;
+        S: GLWESecretToRef + GetDistribution;
 }

 impl<BE: Backend> TensorKeyEncryptSk<BE> for Module<BE>
@@ -95,7 +95,7 @@ where
        scratch: &mut Scratch<BE>,
    ) where
        R: TensorKeyToMut,
-        S: GLWESecretToRef + GetDist,
+        S: GLWESecretToRef + GetDistribution,
    {
        let res: &mut TensorKey<&mut [u8]> = &mut res.to_mut();

--- a/poulpy-core/src/encryption/glwe_ct.rs
+++ b/poulpy-core/src/encryption/glwe_ct.rs
@@ -510,6 +510,8 @@ where
                // ct[i] = uniform (+ pt)
                self.vec_znx_fill_uniform(base2k, ct, col_ct, source_xa);

+                println!("vec_znx_fill_uniform: {}", ct);
+
                let (mut ci_dft, scratch_3) = scratch_2.take_vec_znx_dft(self, 1, size);

                // ci = ct[i] - pt
--- a/poulpy-core/src/encryption/glwe_pk.rs
+++ b/poulpy-core/src/encryption/glwe_pk.rs
@@ -5,7 +5,7 @@ use poulpy_hal::{
 };

 use crate::{
-    Distribution, ScratchTakeCore,
+    Distribution, GetDistribution, GetDistributionMut, ScratchTakeCore,
    encryption::glwe_ct::GLWEEncryptSk,
    layouts::{
        GLWE, GLWEPublicKey, GLWEPublicKeyToMut, LWEInfos,
@@ -16,7 +16,7 @@ use crate::{
 impl<D: DataMut> GLWEPublicKey<D> {
    pub fn generate<S, M, BE: Backend>(&mut self, module: &M, sk: &S, source_xa: &mut Source, source_xe: &mut Source)
    where
-        S: GLWESecretPreparedToRef<BE>,
+        S: GLWESecretPreparedToRef<BE> + GetDistribution,
        M: GLWEPublicKeyGenerate<BE>,
    {
        module.glwe_public_key_generate(self, sk, source_xa, source_xe);
@@ -26,8 +26,8 @@ impl<D: DataMut> GLWEPublicKey<D> {
 pub trait GLWEPublicKeyGenerate<BE: Backend> {
    fn glwe_public_key_generate<R, S>(&self, res: &mut R, sk: &S, source_xa: &mut Source, source_xe: &mut Source)
    where
-        R: GLWEPublicKeyToMut,
-        S: GLWESecretPreparedToRef<BE>;
+        R: GLWEPublicKeyToMut + GetDistributionMut,
+        S: GLWESecretPreparedToRef<BE> + GetDistribution;
 }

 impl<BE: Backend> GLWEPublicKeyGenerate<BE> for Module<BE>
@@ -38,25 +38,27 @@ where
 {
    fn glwe_public_key_generate<R, S>(&self, res: &mut R, sk: &S, source_xa: &mut Source, source_xe: &mut Source)
    where
-        R: GLWEPublicKeyToMut,
-        S: GLWESecretPreparedToRef<BE>,
+        R: GLWEPublicKeyToMut + GetDistributionMut,
+        S: GLWESecretPreparedToRef<BE> + GetDistribution,
    {
-        let res: &mut GLWEPublicKey<&mut [u8]> = &mut res.to_mut();
-        let sk: &GLWESecretPrepared<&[u8], BE> = &sk.to_ref();
+        {
+            let res: &mut GLWEPublicKey<&mut [u8]> = &mut res.to_mut();
+            let sk: &GLWESecretPrepared<&[u8], BE> = &sk.to_ref();

-        assert_eq!(res.n(), self.n() as u32);
-        assert_eq!(sk.n(), self.n() as u32);
+            assert_eq!(res.n(), self.n() as u32);
+            assert_eq!(sk.n(), self.n() as u32);

-        if sk.dist == Distribution::NONE {
-            panic!("invalid sk: SecretDistribution::NONE")
+            if sk.dist == Distribution::NONE {
+                panic!("invalid sk: SecretDistribution::NONE")
+            }
+
+            // Its ok to allocate scratch space here since pk is usually generated only once.
+            let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(self.glwe_encrypt_sk_tmp_bytes(res));
+
+            let mut tmp: GLWE<Vec<u8>> = GLWE::alloc_from_infos(self, res);
+
+            tmp.encrypt_zero_sk(self, sk, source_xa, source_xe, scratch.borrow());
        }
-
-        // Its ok to allocate scratch space here since pk is usually generated only once.
-        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(self.glwe_encrypt_sk_tmp_bytes(res));
-
-        let mut tmp: GLWE<Vec<u8>> = GLWE::alloc_from_infos(self, res);
-
-        tmp.encrypt_zero_sk(self, sk, source_xa, source_xe, scratch.borrow());
-        res.dist = sk.dist;
+        *res.dist_mut() = *sk.dist();
    }
 }