fix blind rotation

poulpy-schemes/src/tfhe/blind_rotation/algorithms/cggi/key.rs

@@ -0,0 +1,79 @@
+use poulpy_hal::{
+    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, ScalarZnxToRef, Scratch, ZnxView, ZnxViewMut},
+    source::Source,
+};
+
+use std::marker::PhantomData;
+
+use poulpy_core::{
+    Distribution, GGSWEncryptSk, GetDistribution, ScratchTakeCore,
+    layouts::{GGSW, GGSWInfos, GLWEInfos, GLWESecretPreparedToRef, LWEInfos, LWESecret, LWESecretToRef},
+};
+
+use crate::tfhe::blind_rotation::{
+    BlindRotationKey, BlindRotationKeyEncryptSk, BlindRotationKeyFactory, BlindRotationKeyInfos, CGGI,
+};
+
+impl<D: DataRef> BlindRotationKeyFactory<CGGI> for BlindRotationKey<D, CGGI> {
+    fn blind_rotation_key_alloc<A>(infos: &A) -> BlindRotationKey<Vec<u8>, CGGI>
+    where
+        A: BlindRotationKeyInfos,
+    {
+        BlindRotationKey {
+            keys: (0..infos.n_lwe().as_usize())
+                .map(|_| GGSW::alloc_from_infos(infos))
+                .collect(),
+            dist: Distribution::NONE,
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<BE: Backend> BlindRotationKeyEncryptSk<BE, CGGI> for Module<BE>
+where
+    Self: GGSWEncryptSk<BE>,
+    Scratch<BE>: ScratchTakeCore<BE>,
+{
+    fn blind_rotation_key_encrypt_sk_tmp_bytes<A: GGSWInfos>(&self, infos: &A) -> usize {
+        self.ggsw_encrypt_sk_tmp_bytes(infos)
+    }
+
+    fn blind_rotation_key_encrypt_sk<D, S0, S1>(
+        &self,
+        res: &mut BlindRotationKey<D, CGGI>,
+        sk_glwe: &S0,
+        sk_lwe: &S1,
+        source_xa: &mut Source,
+        source_xe: &mut Source,
+        scratch: &mut Scratch<BE>,
+    ) where
+        D: DataMut,
+        S0: GLWESecretPreparedToRef<BE> + GLWEInfos,
+        S1: LWESecretToRef + LWEInfos + GetDistribution,
+    {
+        assert_eq!(res.keys.len() as u32, sk_lwe.n());
+        assert!(sk_glwe.n() <= self.n() as u32);
+        assert_eq!(sk_glwe.rank(), res.rank());
+
+        match sk_lwe.dist() {
+            Distribution::BinaryBlock(_) | Distribution::BinaryFixed(_) | Distribution::BinaryProb(_) | Distribution::ZERO => {}
+            _ => {
+                panic!("invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)")
+            }
+        }
+
+        {
+            let sk_lwe: &LWESecret<&[u8]> = &sk_lwe.to_ref();
+
+            res.dist = sk_lwe.dist();
+
+            let mut pt: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(sk_glwe.n().into(), 1);
+            let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data().to_ref();
+
+            for (i, ggsw) in res.keys.iter_mut().enumerate() {
+                pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
+                ggsw.encrypt_sk(self, &pt, sk_glwe, source_xa, source_xe, scratch);
+            }
+        }
+    }
+}