From 3aa1b8783f34db056bc14b348ebdff75035601bc Mon Sep 17 00:00:00 2001
From: nindanaoto <matsuoka.kotaro@gmail.com>
Date: Tue, 22 Oct 2024 16:18:43 +0000
Subject: [PATCH] Cheng's Packing

---
 include/keyswitch.hpp  | 43 ++++++++++++++++++++++++++++++++++++
 test/chengspacking.cpp | 50 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+)
 create mode 100644 test/chengspacking.cpp

diff --git a/include/keyswitch.hpp b/include/keyswitch.hpp
index 2760f6e..6cf2c35 100644
--- a/include/keyswitch.hpp
+++ b/include/keyswitch.hpp
@@ -342,4 +342,47 @@ void SubsetPrivKeySwitch(TRLWE<typename P::targetP> &res,
     }
 }
 
+template <class P> 
+void PackLWEs(TRLWE<P> &res, const std::vector<TLWE<P>> &tlwe, const AnnihilateKey<P> &ahk, const uint l, const uint offset, const uint interval)
+{
+    if(l==0) InvSampleExtractIndex<P>(res,tlwe[offset],0);
+    else{
+        TRLWE<P> tempeven;
+        PackLWEs<P>(tempeven, tlwe, ahk, l-1, offset, interval*2);
+        TRLWE<P> tempodd;
+        PackLWEs<P>(tempodd, tlwe, ahk, l-1, offset+interval, interval*2);
+        TRLWE<P> tempoddmul;
+        for(int i = 0; i < P::k+1; i++){
+            PolynomialMulByXai<P>(tempoddmul[i], tempodd[i], P::n>>l);
+            for(int j = 0; j < P::n; j++){
+                tempeven[i][j] /= 2;
+                tempoddmul[i][j] /= 2;
+                tempodd[i][j] = tempeven[i][j] - tempoddmul[i][j];
+                // tempodd[i][j] = (tempeven[i][j] - tempoddmul[i][j])/2;
+            }
+        }
+        EvalAuto<P>(res, tempodd, (1<<l)+1, ahk[P::nbit - l]);
+        for(int i = 0; i < P::k+1; i++)
+            for(int j = 0; j < P::n; j++)
+                res[i][j] += tempeven[i][j] + tempoddmul[i][j];
+                // res[i][j] += (tempeven[i][j] + tempoddmul[i][j])/2;
+    }
+}
+
+template <class P> 
+void TLWE2TRLWEChengsPacking(TRLWE<P> &res, std::vector<TLWE<P>> &tlwe, const AnnihilateKey<P> &ahk)
+{
+    uint l = std::bit_width(tlwe.size()) - 1;
+    if(!std::has_single_bit(tlwe.size())){
+        l++;
+        tlwe.resize(1<<l);
+    }
+    PackLWEs<P>(res, tlwe, ahk, l, 0, 1);
+    for (int i = 0; i < P::nbit - l; i++) {
+        TRLWE<P> evaledauto;
+        for (int j = 0; j < (P::k+1) * P::n; j++) res[0][j] /= 2;
+        EvalAuto<P>(evaledauto, res, (1 << (P::nbit - i)) + 1, ahk[i]);
+        for (int j = 0; j < (P::k+1) * P::n; j++) res[0][j] += evaledauto[0][j];
+    }
+}
 }  // namespace TFHEpp
\ No newline at end of file
diff --git a/test/chengspacking.cpp b/test/chengspacking.cpp
new file mode 100644
index 0000000..6ff809c
--- /dev/null
+++ b/test/chengspacking.cpp
@@ -0,0 +1,50 @@
+#include <cassert>
+#include <chrono>
+#include <iostream>
+#include <memory>
+#include <random>
+#include <tfhe++.hpp>
+
+int main()
+{
+    constexpr uint32_t num_test = 1000;
+    constexpr uint l = 4;
+    constexpr uint numtlwe = 1<<l;
+    std::random_device seed_gen;
+    std::default_random_engine engine(seed_gen());
+    std::uniform_int_distribution<uint32_t> binary(0, 1);
+
+    std::unique_ptr<TFHEpp::SecretKey> sk(new TFHEpp::SecretKey());
+
+    std::vector<std::vector<TFHEpp::TLWE<TFHEpp::lvl1param>>> ca(num_test);
+
+    std::vector<std::vector<uint8_t>> pin(num_test);
+    for (std::vector<uint8_t> &i : pin){
+        i.resize(numtlwe);
+        for (uint8_t &p : i) p = binary(engine);
+    }
+    for (int i = 0; i < num_test; i++) ca[i] = TFHEpp::bootsSymEncrypt<TFHEpp::lvl1param>(pin[i], *sk);
+
+    std::vector<TFHEpp::TRLWE<TFHEpp::lvl1param>> cres(num_test);
+
+    std::unique_ptr<TFHEpp::AnnihilateKey<TFHEpp::lvl1param>> ahk(new TFHEpp::AnnihilateKey<TFHEpp::lvl1param>());
+    TFHEpp::annihilatekeygen<TFHEpp::lvl1param>(*ahk, *sk);
+
+    std::chrono::system_clock::time_point start, end;
+    start = std::chrono::system_clock::now();
+
+    for (int test = 0; test < num_test; test++) {
+        TFHEpp::TLWE2TRLWEChengsPacking<TFHEpp::lvl1param>(cres[test], ca[test], *ahk);
+    }
+
+    end = std::chrono::system_clock::now();
+
+    for (int i = 0; i < num_test; i++){
+        std::array<bool, TFHEpp::lvl1param::n> pres = TFHEpp::trlweSymDecrypt<TFHEpp::lvl1param>(cres[i], sk->key.lvl1);
+        for(int j = 0; j < numtlwe; j++) assert(pres[j*(TFHEpp::lvl1param::n>>l)] == (pin[i][j] > 0));
+    }
+
+    std::cout << "Passed" << std::endl;
+    double elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
+    std::cout << elapsed / num_test << "ms" << std::endl;
+}
\ No newline at end of file