/*
 * Crystals Dilithium Polynomial Adapter
 *
 * (C) 2022-2023 Jack Lloyd
 * (C) 2022      Manuel Glaser - Rohde & Schwarz Cybersecurity
 * (C) 2022-2023 Michael Boric, René Meusel - Rohde & Schwarz Cybersecurity
 * (C) 2024 René Meusel, Rohde & Schwarz Cybersecurity
 *
 * Botan is released under the Simplified BSD License (see license.txt)
 */

#ifndef BOTAN_DILITHIUM_POLYNOMIAL_H_
#define BOTAN_DILITHIUM_POLYNOMIAL_H_

#include <botan/mem_ops.h>
#include <botan/internal/dilithium_constants.h>
#include <botan/internal/pqcrystals.h>
#include <botan/internal/pqcrystals_helpers.h>

namespace Botan {

class DilithiumPolyTraits final : public CRYSTALS::Trait_Base<DilithiumConstants, DilithiumPolyTraits> {
   private:
      friend class CRYSTALS::Trait_Base<DilithiumConstants, DilithiumPolyTraits>;

      static constexpr T montgomery_reduce_coefficient(T2 a) {
         const T2 t = static_cast<T>(static_cast<T2>(static_cast<T>(a)) * Q_inverse);
         return static_cast<T>((a - static_cast<T2>(t) * Q) >> (sizeof(T) * 8));
      }

      static constexpr T barrett_reduce_coefficient(T a) {
         // 2**22 is roughly Q/2 and 2**23 is roughly Q
         const T t = (a + (1 << 22)) >> 23;
         a = a - t * Q;
         return a;
      }

   public:
      /**
       * NIST FIPS 204, Algorithm 41 (NTT)
       *
       * Note: ntt(), inverse_ntt() and operator* have side effects on the
       *       montgomery factor of the involved coefficients!
       *       It is assumed that EXACTLY ONE vector or matrix multiplication
       *       is performed between transforming in and out of NTT domain.
       *
       * Produces the result of the NTT transformation without any montgomery
       * factors in the coefficients.
       */
      static constexpr void ntt(std::span<T, N> coeffs) {
         size_t j = 0;
         size_t k = 0;

         for(size_t len = N / 2; len > 0; len >>= 1) {
            for(size_t start = 0; start < N; start = j + len) {
               const T zeta = zetas[++k];
               for(j = start; j < start + len; ++j) {
                  // Zetas contain the montgomery parameter 2^32 mod q
                  const T t = fqmul(zeta, coeffs[j + len]);
                  coeffs[j + len] = coeffs[j] - t;
                  coeffs[j] = coeffs[j] + t;
               }
            }
         }
      }

      /**
       * NIST FIPS 204, Algorithm 42 (NTT^-1).
       *
       * The output is effectively multiplied by the montgomery parameter 2^32
       * mod q so that the input factors 2^(-32) mod q are eliminated. Note
       * that factors 2^(-32) mod q are introduced by multiplication and
       * reduction of values not in montgomery domain.
       *
       * Produces the result of the inverse NTT transformation with a montgomery
       * factor of (2^32 mod q) added (!). See above.
       */
      static constexpr void inverse_ntt(std::span<T, N> coeffs) {
         size_t j = 0;
         size_t k = N;
         for(size_t len = 1; len < N; len <<= 1) {
            for(size_t start = 0; start < N; start = j + len) {
               const T zeta = -zetas[--k];
               for(j = start; j < start + len; ++j) {
                  const T t = coeffs[j];
                  coeffs[j] = t + coeffs[j + len];
                  coeffs[j + len] = t - coeffs[j + len];
                  // Zetas contain the montgomery parameter 2^32 mod q
                  coeffs[j + len] = fqmul(zeta, coeffs[j + len]);
               }
            }
         }

         for(auto& coeff : coeffs) {
            coeff = fqmul(coeff, F_WITH_MONTY_SQUARED);
         }
      }

      /**
       * Multiplication of two polynomials @p lhs and @p rhs in NTT domain.
       *
       * Produces the result of the multiplication in NTT domain, with a factor
       * of (2^-32 mod q) in each element due to montgomery reduction.
       */
      static constexpr void poly_pointwise_montgomery(std::span<T, N> result,
                                                      std::span<const T, N> lhs,
                                                      std::span<const T, N> rhs) {
         for(size_t i = 0; i < N; ++i) {
            result[i] = fqmul(lhs[i], rhs[i]);
         }
      }
};

}  // namespace Botan

#endif
