From b464c7aa855372bea16f6db40acb689b172f8e49 Mon Sep 17 00:00:00 2001
From: nindanaoto <matsuoka.kotaro@gmail.com>
Date: Sun, 11 Feb 2024 17:19:28 +0000
Subject: [PATCH] May be nussbaumer is working?

---
 include/nussbaumer.hpp | 113 +++++++++++++++++++++++++++++++++++++++++
 include/params.hpp     |   1 +
 test/nussbaumer.cpp    |  69 +++++++++++++++++++++++++
 3 files changed, 183 insertions(+)
 create mode 100644 include/nussbaumer.hpp
 create mode 100644 test/nussbaumer.cpp

diff --git a/include/nussbaumer.hpp b/include/nussbaumer.hpp
new file mode 100644
index 0000000..599ff3e
--- /dev/null
+++ b/include/nussbaumer.hpp
@@ -0,0 +1,113 @@
+#pragma once
+#include <span>
+
+namespace Nussbaumer{
+
+template <typename T, uint rbit>
+inline void PolynomialMulByXai(const std::span<T,1ull<<rbit> res, const size_t a)
+{
+    if (a == 0)
+        return;
+    else{
+        constexpr size_t r = 1ull<<rbit;
+        std::array<T,r> temp;
+        std::copy(res.begin(),res.end(),temp.begin());
+        if (a < r) {
+            for (int i = 0; i < a; i++) res[i] = -temp[i - a + r];
+            for (int i = a; i < r; i++) res[i] = temp[i - a];
+        }
+        else {
+            const size_t aa = a - r;
+            for (int i = 0; i < aa; i++) res[i] = temp[i - aa + r];
+            for (int i = aa; i < r; i++) res[i] = -temp[i - aa];
+        }
+    }
+}
+
+template<typename T, uint mbit, uint rbit>
+void NussbaumerButterfly(const std::span<T,(1u<<(rbit+mbit))> res){
+    constexpr size_t m = 1ull<<mbit;
+    constexpr size_t r = 1ull<<rbit;
+    for(int i = 0; i < m/2; i++)
+        for(int j = 0; j < r; j++){
+            const T temp = res[i*r+j];
+            res[i*r+j] += res[(i+m/2)*r+j];
+            res[(i+m/2)*r+j] = temp - res[(i+m/2)*r+j];
+        }
+    if constexpr(mbit!=1){
+        constexpr size_t stride = 1ull<<(rbit-mbit);
+        for(int i = 1; i < m/2; i++) PolynomialMulByXai<T, rbit>(static_cast<std::span<T,r>>(res.subspan((i+m/2)*r,r)),i*stride);
+        NussbaumerButterfly<T,mbit-1,rbit>(res.template subspan<0,m*r/2>());
+        NussbaumerButterfly<T,mbit-1,rbit>(res.template subspan<m*r/2,m*r/2>());
+    }
+}
+
+template<typename T, uint Nbit>
+void NussbaumerTransform(std::span<T,(1ull<<Nbit)> res){
+    if constexpr(Nbit == 1){
+        const T temp = res[0];
+        res[0] += res[1];
+        res[1] = temp - res[1];
+        return; 
+    }else{
+        //initialize
+        constexpr uint mbit = Nbit/2;
+        constexpr size_t m = 1ull<<mbit;
+        constexpr uint rbit = Nbit-mbit;
+        constexpr size_t r = 1ull<<rbit;
+        std::array<T,(1ull<<Nbit)> temp;
+        std::copy(res.begin(),res.end(),temp.begin());
+        //reorder
+        for(int i = 0; i < m; i++){
+            for(int j = 0; j < r; j++)
+                res[i*r+j] = temp[m*j+i];
+        }
+        NussbaumerButterfly<T,mbit,rbit>(res);
+        for(int i = 0; i < m; i++)
+            NussbaumerTransform<T,rbit>(static_cast<std::span<T,r>>(res.subspan(i*r,r)));
+    }
+}
+
+template<typename T, uint mbit, uint rbit>
+void InverseNussbaumerButterfly(const std::span<T,(1u<<(rbit+mbit))> res){
+    constexpr size_t m = 1ull<<mbit;
+    constexpr size_t r = 1ull<<rbit;
+    if constexpr(mbit!=1){
+        constexpr size_t stride = 1ull<<(rbit-mbit);
+        InverseNussbaumerButterfly<T,mbit-1,rbit>(res.template subspan<0,m*r/2>());
+        InverseNussbaumerButterfly<T,mbit-1,rbit>(res.template subspan<m*r/2,m*r/2>());
+        for(int i = 1; i < m/2; i++) PolynomialMulByXai<T, rbit>(static_cast<std::span<T,r>>(res.subspan((i+m/2)*r,r)),2*r-i*stride);
+    }
+    for(int i = 0; i < m/2; i++)
+        for(int j = 0; j < r; j++){
+            const T temp = res[i*r+j];
+            res[i*r+j] += res[(i+m/2)*r+j];
+            res[(i+m/2)*r+j] = temp - res[(i+m/2)*r+j];
+        }
+}
+
+template<typename T, uint Nbit>
+void InverseNussbaumerTransform(std::span<T,(1ull<<Nbit)> res){
+    if constexpr(Nbit == 1){
+        const T temp = res[0];
+        res[0] += res[1];
+        res[1] = temp - res[1];
+        return; 
+    }else{
+        //initialize
+        constexpr uint mbit = Nbit/2;
+        constexpr size_t m = 1ull<<mbit;
+        constexpr uint rbit = Nbit-mbit;
+        constexpr size_t r = 1ull<<rbit;
+        for(int i = 0; i < m; i++)
+            InverseNussbaumerTransform<T,rbit>(static_cast<std::span<T,r>>(res.subspan(i*r,r)));
+        InverseNussbaumerButterfly<T,mbit,rbit>(res);
+        std::array<T,(1ull<<Nbit)> temp;
+        std::copy(res.begin(),res.end(),temp.begin());
+        //reorder
+        for(int i = 0; i < m; i++)
+            for(int j = 0; j < r; j++)
+                res[m*j+i] = temp[i*r+j];
+    }
+}
+}
\ No newline at end of file
diff --git a/include/params.hpp b/include/params.hpp
index 0e59443..62cf2eb 100644
--- a/include/params.hpp
+++ b/include/params.hpp
@@ -5,6 +5,7 @@
 
 #include "cuhe++.hpp"
 #include "raintt.hpp"
+#include "nussbaumer.hpp"
 
 namespace TFHEpp {
 
diff --git a/test/nussbaumer.cpp b/test/nussbaumer.cpp
new file mode 100644
index 0000000..06a5cdf
--- /dev/null
+++ b/test/nussbaumer.cpp
@@ -0,0 +1,69 @@
+#include <algorithm>
+#include <cassert>
+#include <iostream>
+#include <random>
+#include <tfhe++.hpp>
+
+int main()
+{
+    constexpr uint32_t num_test = 1000;
+    std::random_device seed_gen;
+    std::default_random_engine engine(seed_gen());
+    std::uniform_int_distribution<uint32_t> Bgdist(0, TFHEpp::lvl1param::Bg);
+    std::uniform_int_distribution<uint32_t> Torus32dist(0, UINT32_MAX);
+
+    // std::cout << "Start LVL1 test." << std::endl;
+    for (int test = 0; test < num_test; test++) {
+        using T = uint64_t;
+        std::array<T,TFHEpp::lvl1param::n> a,res;
+        for (T &i : a) i = Torus32dist(engine);
+        res = a;
+        Nussbaumer::NussbaumerTransform<T,TFHEpp::lvl1param::nbit>(std::span{res});
+        Nussbaumer::InverseNussbaumerTransform<T,TFHEpp::lvl1param::nbit>(std::span{res});
+        for (int i = 0; i < TFHEpp::lvl1param::n; i++)
+            assert(abs(static_cast<int32_t>(a[i] - res[i]/TFHEpp::lvl1param::n) <= 1));
+    }
+    std::cout << "Id Passed" << std::endl;
+
+    // for (int test = 0; test < num_test; test++) {
+    //     array<typename TFHEpp::lvl1param::T, lvl1param::n> a;
+    //     for (int i = 0; i < lvl1param::n; i++)
+    //         a[i] = Bgdist(engine) - lvl1param::Bg / 2;
+    //     for (typename TFHEpp::lvl1param::T &i : a)
+    //         i = Bgdist(engine) - lvl1param::Bg / 2;
+    //     array<typename TFHEpp::lvl1param::T, lvl1param::n> b;
+    //     for (typename TFHEpp::lvl1param::T &i : b) i = Torus32dist(engine);
+
+    //     Polynomial<lvl1param> polymul;
+    //     TFHEpp::PolyMul<lvl1param>(polymul, a, b);
+    //     Polynomial<lvl1param> naieve = {};
+    //     for (int i = 0; i < lvl1param::n; i++) {
+    //         for (int j = 0; j <= i; j++)
+    //             naieve[i] += static_cast<int32_t>(a[j]) * b[i - j];
+    //         for (int j = i + 1; j < lvl1param::n; j++)
+    //             naieve[i] -=
+    //                 static_cast<int32_t>(a[j]) * b[lvl1param::n + i - j];
+    //     }
+    //     for (int i = 0; i < lvl1param::n; i++)
+    //         assert(abs(static_cast<int32_t>(naieve[i] - polymul[i])) <= 1);
+    // }
+    // cout << "PolyMul Passed" << endl;
+
+    // uniform_int_distribution<uint64_t> Bgbardist(0, lvl2param::Bg);
+    // uniform_int_distribution<uint64_t> Torus64dist(0, UINT64_MAX);
+
+    // cout << "Start LVL2 test." << endl;
+    // for (int test = 0; test < num_test; test++) {
+    //     Polynomial<lvl2param> a;
+    //     for (uint64_t &i : a) i = Torus64dist(engine);
+    //     PolynomialInFD<lvl2param> resfft;
+    //     TFHEpp::TwistIFFT<lvl2param>(resfft, a);
+    //     Polynomial<lvl2param> res;
+    //     TFHEpp::TwistFFT<lvl2param>(res, resfft);
+    //     for (int i = 0; i < lvl2param::n; i++)
+    //         assert(abs(static_cast<int64_t>(a[i] - res[i])) <= (1 << 14));
+    // }
+    // cout << "FFT Passed" << endl;
+
+    return 0;
+}
\ No newline at end of file