1#include "ntt120_dft.h"
  2
  3#include "mod_q120.h"
  4
  5// @brief alternative version of the NTT
  6
  7/** for all s=k/2^17, root_of_unity(s) = omega_0^k */
  8static mod_q120 root_of_unity(double s) {
  9  static mod_q120 omega_2pow17{OMEGA1, OMEGA2, OMEGA3, OMEGA4};
 10  static double _2pow17 = 1 << 17;
 11  return pow(omega_2pow17, s * _2pow17);
 12}
 13static mod_q120 root_of_unity_inv(double s) {
 14  static mod_q120 omega_2pow17{OMEGA1, OMEGA2, OMEGA3, OMEGA4};
 15  static double _2pow17 = 1 << 17;
 16  return pow(omega_2pow17, -s * _2pow17);
 17}
 18
 19/** recursive naive ntt */
 20static void q120_ntt_naive_rec(uint64_t n, double entry_pwr, mod_q120* data) {
 21  if (n == 1) return;
 22  const uint64_t h = n / 2;
 23  const double s = entry_pwr / 2.;
 24  mod_q120 om = root_of_unity(s);
 25  for (uint64_t j = 0; j < h; ++j) {
 26    mod_q120 om_right = data[h + j] * om;
 27    data[h + j] = data[j] - om_right;
 28    data[j] = data[j] + om_right;
 29  }
 30  q120_ntt_naive_rec(h, s, data);
 31  q120_ntt_naive_rec(h, s + 0.5, data + h);
 32}
 33static void q120_intt_naive_rec(uint64_t n, double entry_pwr, mod_q120* data) {
 34  if (n == 1) return;
 35  const uint64_t h = n / 2;
 36  const double s = entry_pwr / 2.;
 37  q120_intt_naive_rec(h, s, data);
 38  q120_intt_naive_rec(h, s + 0.5, data + h);
 39  mod_q120 om = root_of_unity_inv(s);
 40  for (uint64_t j = 0; j < h; ++j) {
 41    mod_q120 dat_diff = half(data[j] - data[h + j]);
 42    data[j] = half(data[j] + data[h + j]);
 43    data[h + j] = dat_diff * om;
 44  }
 45}
 46
 47/** user friendly version */
 48q120_nttvec simple_ntt120(const znx_i64& polynomial) {
 49  const uint64_t n = polynomial.nn();
 50  q120_nttvec res(n);
 51  for (uint64_t i = 0; i < n; ++i) {
 52    int64_t xi = polynomial.get_coeff(i);
 53    res.v[i] = mod_q120(xi, xi, xi, xi);
 54  }
 55  q120_ntt_naive_rec(n, 0.5, res.v.data());
 56  return res;
 57}
 58
 59znx_i128 simple_intt120(const q120_nttvec& fftvec) {
 60  const uint64_t n = fftvec.nn();
 61  q120_nttvec copy = fftvec;
 62  znx_i128 res(n);
 63  q120_intt_naive_rec(n, 0.5, copy.v.data());
 64  for (uint64_t i = 0; i < n; ++i) {
 65    res.set_coeff(i, copy.v[i].to_int128());
 66  }
 67  return res;
 68}
 69bool operator==(const q120_nttvec& a, const q120_nttvec& b) { return a.v == b.v; }
 70
 71std::vector<mod_q120> q120_ntt_naive(const std::vector<mod_q120>& x) {
 72  std::vector<mod_q120> res = x;
 73  q120_ntt_naive_rec(res.size(), 0.5, res.data());
 74  return res;
 75}
 76q120_nttvec::q120_nttvec(uint64_t n) : v(n) {}
 77q120_nttvec::q120_nttvec(uint64_t n, const q120b* data) : v(n) {
 78  int64_t* d = (int64_t*)data;
 79  for (uint64_t i = 0; i < n; ++i) {
 80    v[i] = mod_q120::from_q120b(d + 4 * i);
 81  }
 82}
 83q120_nttvec::q120_nttvec(uint64_t n, const q120c* data) : v(n) {
 84  int64_t* d = (int64_t*)data;
 85  for (uint64_t i = 0; i < n; ++i) {
 86    v[i] = mod_q120::from_q120c(d + 4 * i);
 87  }
 88}
 89uint64_t q120_nttvec::nn() const { return v.size(); }
 90q120_nttvec q120_nttvec::zero(uint64_t n) { return q120_nttvec(n); }
 91void q120_nttvec::save_as(q120a* dest) const {
 92  int64_t* const d = (int64_t*)dest;
 93  const uint64_t n = nn();
 94  for (uint64_t i = 0; i < n; ++i) {
 95    v[i].save_as_q120a(d + 4 * i);
 96  }
 97}
 98void q120_nttvec::save_as(q120b* dest) const {
 99  int64_t* const d = (int64_t*)dest;
100  const uint64_t n = nn();
101  for (uint64_t i = 0; i < n; ++i) {
102    v[i].save_as_q120b(d + 4 * i);
103  }
104}
105void q120_nttvec::save_as(q120c* dest) const {
106  int64_t* const d = (int64_t*)dest;
107  const uint64_t n = nn();
108  for (uint64_t i = 0; i < n; ++i) {
109    v[i].save_as_q120c(d + 4 * i);
110  }
111}
112mod_q120 q120_nttvec::get_blk(uint64_t blk) const {
113  REQUIRE_DRAMATICALLY(blk < nn(), "blk overflow");
114  return v[blk];
115}
116q120_nttvec q120_nttvec::random(uint64_t n) {
117  q120_nttvec res(n);
118  for (uint64_t i = 0; i < n; ++i) {
119    res.v[i] = uniform_q120();
120  }
121  return res;
122}