From 31ffcf6ec0fd3bf7e9aaf493734f6c4020cc0d3d Mon Sep 17 00:00:00 2001
From: nindanaoto <dirty.knife603@gmail.com>
Date: Mon, 4 Oct 2021 19:03:28 +0900
Subject: [PATCH] Added privkskgen

---
 include/cloudkey.hpp  | 25 ++++++++++++++++++++++++-
 include/keyswitch.hpp |  2 +-
 include/params.hpp    |  2 +-
 src/cloudkey.cpp      | 23 ++++-------------------
 src/keyswitch.cpp     |  4 ++--
 5 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/include/cloudkey.hpp b/include/cloudkey.hpp
index bce2b9b..7e7f888 100644
--- a/include/cloudkey.hpp
+++ b/include/cloudkey.hpp
@@ -82,6 +82,29 @@ inline void ikskgen(KeySwitchingKey<P> &ksk, const SecretKey &sk)
                              (j + 1) * P::basebit)),
                     P::α, sk.key.get<typename P::targetP>());
 }
+template <class P>
+inline void privkskgen(PrivateKeySwitchingKey<P> &privksk, Polynomial<typename P::targetP> func,const SecretKey &sk)
+{
+    std::array<typename P::domainP::T, P::domainP::n + 1> key;
+    for (int i = 0; i < P::domainP::n; i++) key[i] = sk.key.lvl2[i];
+    key[P::domainP::n] = -1;
+#pragma omp parallel for collapse(3)
+    for (int i = 0; i <= P::domainP::n; i++)
+        for (int j = 0; j < P::t; j++)
+            for (typename P::targetP::T u = 0;
+                    u < (1 << P::basebit) - 1; u++) {
+                TRLWE<typename P::targetP> c =
+                    trlweSymEncryptZero<typename P::targetP>(
+                        P::α,
+                        sk.key.get<typename P::targetP>());
+                for(int k = 0; k<P::targetP::n;k++)
+                c[1][k] += (u + 1) * func[k]*key[i]
+                            << (numeric_limits<
+                                    typename P::targetP::T>::digits -
+                                (j + 1) * P::basebit);
+                privksk[i][j][u] = c;
+            }
+}
 
 template <class P>
 inline relinKey<P> relinKeygen(const Key<P> &key)
@@ -151,7 +174,7 @@ struct GateKeyNTT {
 template <class bsP, class privksP>
 struct CircuitKey {
     BootstrappingKeyFFT<bsP> bkfft;
-    std::array<PrivKeySwitchKey<privksP>, 2> privksk;
+    std::array<PrivateKeySwitchingKey<privksP>, 2> privksk;
     CircuitKey(const SecretKey &sk);
     CircuitKey() {}
     template <class Archive>
diff --git a/include/keyswitch.hpp b/include/keyswitch.hpp
index 7bf1bf8..b30ea68 100644
--- a/include/keyswitch.hpp
+++ b/include/keyswitch.hpp
@@ -45,5 +45,5 @@ void AnnihilatePrivateKeySwitching(std::array<TRLWE<P>,num_func> &res, const TRL
 template <class P>
 void PrivKeySwitch(TRLWE<typename P::targetP> &res,
                    const TLWE<typename P::domainP> &tlwe,
-                   const PrivKeySwitchKey<P> &privksk);
+                   const PrivateKeySwitchingKey<P> &privksk);
 }  // namespace TFHEpp
\ No newline at end of file
diff --git a/include/params.hpp b/include/params.hpp
index 0593dca..9360f09 100644
--- a/include/params.hpp
+++ b/include/params.hpp
@@ -87,7 +87,7 @@ using TLWE2TRLWEIKSKey = std::array<
 template <class P>
 using AnnihilateKey = std::array<TRGSWFFT<P>, P::nbit>;
 template <class P>
-using PrivKeySwitchKey = std::array<
+using PrivateKeySwitchingKey = std::array<
     std::array<std::array<TRLWE<typename P::targetP>, (1 << P::basebit) - 1>,
                P::t>,
     P::domainP::n + 1>;
diff --git a/src/cloudkey.cpp b/src/cloudkey.cpp
index bfe16b5..7661937 100644
--- a/src/cloudkey.cpp
+++ b/src/cloudkey.cpp
@@ -41,25 +41,10 @@ CircuitKey<bsP, privksP>::CircuitKey(const SecretKey &sk)
     bkfftgen<bsP>(bkfft, sk);
 
     // Generate privksk
-    array<typename privksP::domainP::T, privksP::domainP::n + 1> key;
-    for (int i = 0; i < privksP::domainP::n; i++) key[i] = sk.key.lvl2[i];
-    key[privksP::domainP::n] = -1;
-#pragma omp parallel for collapse(4)
-    for (int z = 0; z < 2; z++)
-        for (int i = 0; i <= privksP::domainP::n; i++)
-            for (int j = 0; j < privksP::t; j++)
-                for (typename privksP::targetP::T u = 0;
-                     u < (1 << privksP::basebit) - 1; u++) {
-                    TRLWE<typename privksP::targetP> c =
-                        trlweSymEncryptZero<typename privksP::targetP>(
-                            privksP::α,
-                            sk.key.get<typename privksP::targetP>());
-                    c[z][0] += (u + 1) * key[i]
-                               << (numeric_limits<
-                                       typename privksP::targetP::T>::digits -
-                                   (j + 1) * privksP::basebit);
-                    privksk[z][i][j][u] = c;
-                }
+    TFHEpp::Polynomial<typename privksP::targetP> poly = {1};
+    privkskgen<privksP>(privksk[1], poly, sk);
+    for(int i = 0; i<privksP::targetP::n; i++) poly[i] = -sk.key.get<typename privksP::targetP>()[i];
+    privkskgen<privksP>(privksk[0], poly, sk);
 }
 #define INST(bsP, privksP) \
     template CircuitKey<bsP, privksP>::CircuitKey(const SecretKey &sk)
diff --git a/src/keyswitch.cpp b/src/keyswitch.cpp
index f06610b..1eedb08 100644
--- a/src/keyswitch.cpp
+++ b/src/keyswitch.cpp
@@ -135,7 +135,7 @@ TFHEPP_EXPLICIT_INSTANTIATION_TRLWE(INST)
 template <class P>
 void PrivKeySwitch(TRLWE<typename P::targetP> &res,
                    const TLWE<typename P::domainP> &tlwe,
-                   const PrivKeySwitchKey<P> &privksk)
+                   const PrivateKeySwitchingKey<P> &privksk)
 {
     constexpr uint32_t mask = (1 << P::basebit) - 1;
     constexpr uint64_t prec_offset =
@@ -164,7 +164,7 @@ void PrivKeySwitch(TRLWE<typename P::targetP> &res,
 #define INST(P)                                                           \
     template void PrivKeySwitch<P>(TRLWE<typename P::targetP> & res,      \
                                    const TLWE<typename P::domainP> &tlwe, \
-                                   const PrivKeySwitchKey<P> &privksk)
+                                   const PrivateKeySwitchingKey<P> &privksk)
 TFHEPP_EXPLICIT_INSTANTIATION_KEY_SWITCH(INST)
 #undef INST