updated repo for publishing (#74)

poulpy-schemes/src/tfhe/circuit_bootstrapping/key.rs

@@ -0,0 +1,165 @@
+use poulpy_core::layouts::{
+    GGLWEAutomorphismKey, GGLWETensorKey, GLWECiphertext, GLWESecret, LWESecret,
+    prepared::{GGLWEAutomorphismKeyPrepared, GGLWETensorKeyPrepared, GLWESecretPrepared, PrepareAlloc},
+};
+use std::collections::HashMap;
+
+use poulpy_backend::hal::{
+    api::{
+        ScratchAvailable, SvpApply, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx, TakeSvpPPol,
+        TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism,
+        VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxDftToVecZnxBigTmpA,
+        VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace,
+        VecZnxSwithcDegree, VmpPMatAlloc, VmpPrepare,
+    },
+    layouts::{Backend, Data, DataRef, Module, Scratch},
+    source::Source,
+};
+
+use crate::tfhe::blind_rotation::{
+    BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk, BlindRotationKeyPrepared,
+};
+
+pub trait CircuitBootstrappingKeyEncryptSk<B: Backend> {
+    #[allow(clippy::too_many_arguments)]
+    fn encrypt_sk<DLwe, DGlwe>(
+        module: &Module<B>,
+        basek: usize,
+        sk_lwe: &LWESecret<DLwe>,
+        sk_glwe: &GLWESecret<DGlwe>,
+        k_brk: usize,
+        rows_brk: usize,
+        k_trace: usize,
+        rows_trace: usize,
+        k_tsk: usize,
+        rows_tsk: usize,
+        source_xa: &mut Source,
+        source_xe: &mut Source,
+        sigma: f64,
+        scratch: &mut Scratch<B>,
+    ) -> Self
+    where
+        DLwe: DataRef,
+        DGlwe: DataRef;
+}
+
+pub struct CircuitBootstrappingKey<D: Data, BRA: BlindRotationAlgo> {
+    pub(crate) brk: BlindRotationKey<D, BRA>,
+    pub(crate) tsk: GGLWETensorKey<Vec<u8>>,
+    pub(crate) atk: HashMap<i64, GGLWEAutomorphismKey<Vec<u8>>>,
+}
+
+impl<BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyEncryptSk<B> for CircuitBootstrappingKey<Vec<u8>, BRA>
+where
+    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
+    Module<B>: SvpApply<B>
+        + VecZnxDftToVecZnxBigTmpA<B>
+        + VecZnxAddScalarInplace
+        + VecZnxDftAllocBytes
+        + VecZnxBigNormalize<B>
+        + VecZnxDftFromVecZnx<B>
+        + SvpApplyInplace<B>
+        + VecZnxDftToVecZnxBigConsume<B>
+        + VecZnxNormalizeTmpBytes
+        + VecZnxFillUniform
+        + VecZnxSubABInplace
+        + VecZnxAddInplace
+        + VecZnxNormalizeInplace<B>
+        + VecZnxAddNormal
+        + VecZnxNormalize<B>
+        + VecZnxSub
+        + SvpPrepare<B>
+        + VecZnxSwithcDegree
+        + SvpPPolAllocBytes
+        + SvpPPolAlloc<B>
+        + VecZnxAutomorphism,
+    Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeSvpPPol<B> + TakeVecZnxBig<B>,
+{
+    fn encrypt_sk<DLwe, DGlwe>(
+        module: &Module<B>,
+        basek: usize,
+        sk_lwe: &LWESecret<DLwe>,
+        sk_glwe: &GLWESecret<DGlwe>,
+        k_brk: usize,
+        rows_brk: usize,
+        k_trace: usize,
+        rows_trace: usize,
+        k_tsk: usize,
+        rows_tsk: usize,
+        source_xa: &mut Source,
+        source_xe: &mut Source,
+        sigma: f64,
+        scratch: &mut Scratch<B>,
+    ) -> Self
+    where
+        DLwe: DataRef,
+        DGlwe: DataRef,
+        Module<B>:,
+    {
+        let mut auto_keys: HashMap<i64, GGLWEAutomorphismKey<Vec<u8>>> = HashMap::new();
+        let gal_els: Vec<i64> = GLWECiphertext::trace_galois_elements(module);
+        gal_els.iter().for_each(|gal_el| {
+            let mut key: GGLWEAutomorphismKey<Vec<u8>> =
+                GGLWEAutomorphismKey::alloc(sk_glwe.n(), basek, k_trace, rows_trace, 1, sk_glwe.rank());
+            key.encrypt_sk(
+                module, *gal_el, sk_glwe, source_xa, source_xe, sigma, scratch,
+            );
+            auto_keys.insert(*gal_el, key);
+        });
+
+        let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch);
+
+        let mut brk: BlindRotationKey<Vec<u8>, BRA> = BlindRotationKey::<Vec<u8>, BRA>::alloc(
+            sk_glwe.n(),
+            sk_lwe.n(),
+            basek,
+            k_brk,
+            rows_brk,
+            sk_glwe.rank(),
+        );
+
+        brk.encrypt_sk(
+            module,
+            &sk_glwe_prepared,
+            sk_lwe,
+            source_xa,
+            source_xe,
+            sigma,
+            scratch,
+        );
+
+        let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(sk_glwe.n(), basek, k_tsk, rows_tsk, 1, sk_glwe.rank());
+        tsk.encrypt_sk(module, sk_glwe, source_xa, source_xe, sigma, scratch);
+
+        Self {
+            brk,
+            atk: auto_keys,
+            tsk,
+        }
+    }
+}
+
+pub struct CircuitBootstrappingKeyPrepared<D: Data, BRA: BlindRotationAlgo, B: Backend> {
+    pub(crate) brk: BlindRotationKeyPrepared<D, BRA, B>,
+    pub(crate) tsk: GGLWETensorKeyPrepared<Vec<u8>, B>,
+    pub(crate) atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
+}
+
+impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> PrepareAlloc<B, CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B>>
+    for CircuitBootstrappingKey<D, BRA>
+where
+    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
+    BlindRotationKey<D, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
+    GGLWETensorKey<D>: PrepareAlloc<B, GGLWETensorKeyPrepared<Vec<u8>, B>>,
+    GGLWEAutomorphismKey<D>: PrepareAlloc<B, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
+{
+    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> {
+        let brk: BlindRotationKeyPrepared<Vec<u8>, BRA, B> = self.brk.prepare_alloc(module, scratch);
+        let tsk: GGLWETensorKeyPrepared<Vec<u8>, B> = self.tsk.prepare_alloc(module, scratch);
+        let mut atk: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>> = HashMap::new();
+        for (key, value) in &self.atk {
+            atk.insert(*key, value.prepare_alloc(module, scratch));
+        }
+        CircuitBootstrappingKeyPrepared { brk, tsk, atk }
+    }
+}